From b7f815f10dd1f86507817bb48bbf6857629be8e7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20Ml=C3=A1dek?= Date: Fri, 25 Apr 2025 19:13:10 +0200 Subject: [PATCH] axum-extra: implement `FromRequest` for `Either*` --- axum-extra/src/either.rs | 164 ++++++++++++++++++++++++++++++++++++++- 1 file changed, 162 insertions(+), 2 deletions(-) diff --git a/axum-extra/src/either.rs b/axum-extra/src/either.rs index 5d0e33eb..37d48c31 100755 --- a/axum-extra/src/either.rs +++ b/axum-extra/src/either.rs @@ -1,6 +1,6 @@ //! `Either*` types for combining extractors or responses into a single type. //! -//! # As an extractor +//! # As an `FromRequestParts` extractor //! //! ``` //! use axum_extra::either::Either3; @@ -54,6 +54,42 @@ //! Note that if all the inner extractors reject the request, the rejection from the last //! extractor will be returned. For the example above that would be [`BytesRejection`]. //! +//! # As an `FromRequest` extractor +//! +//! In the following example, we can first try to deserialize the payload as JSON, if that fails try +//! to interpret it as a UTF-8 string, and lastly just take the raw bytes. +//! +//! It might be preferable to instead extract `Bytes` directly and then fallibly convert them to +//! `String` and then deserialize the data inside the handler. +//! +//! ``` +//! use axum_extra::either::Either3; +//! use axum::{ +//! body::Bytes, +//! Json, +//! Router, +//! routing::get, +//! extract::FromRequestParts, +//! }; +//! +//! #[derive(serde::Deserialize)] +//! struct Payload { +//! user: String, +//! request_id: u32, +//! } +//! +//! async fn handler( +//! body: Either3, String, Bytes>, +//! ) { +//! match body { +//! Either3::E1(json) => { /* ... */ } +//! Either3::E2(string) => { /* ... */ } +//! Either3::E3(bytes) => { /* ... */ } +//! } +//! } +//! # +//! # let _: axum::routing::MethodRouter = axum::routing::get(handler); +//! ``` //! # As a response //! //! ``` @@ -93,9 +129,10 @@ use std::task::{Context, Poll}; use axum::{ - extract::FromRequestParts, + extract::{rejection::BytesRejection, FromRequest, FromRequestParts, Request}, response::{IntoResponse, Response}, }; +use bytes::Bytes; use http::request::Parts; use tower_layer::Layer; use tower_service::Service; @@ -226,6 +263,28 @@ pub enum Either8 { E8(E8), } +/// Rejection used for [`Either`], [`Either3`], etc. +/// +/// Contains one variant for a case when the whole request could not be loaded and one variant +/// containing the rejection of the last variant if all extractors failed.. +#[derive(Debug)] +pub enum EitherRejection { + /// Buffering of the request body failed. + Bytes(BytesRejection), + + /// All extractors failed. This contains the error returned by the last extractor. + LastRejection(E), +} + +impl IntoResponse for EitherRejection { + fn into_response(self) -> Response { + match self { + EitherRejection::Bytes(rejection) => rejection.into_response(), + EitherRejection::LastRejection(rejection) => rejection.into_response(), + } + } +} + macro_rules! impl_traits_for_either { ( $either:ident => @@ -251,6 +310,43 @@ macro_rules! impl_traits_for_either { } } + impl FromRequest for $either<$($ident),*, $last> + where + S: Send + Sync, + $($ident: FromRequest),*, + $last: FromRequest, + $($ident::Rejection: Send),*, + $last::Rejection: IntoResponse + Send, + { + type Rejection = EitherRejection<$last::Rejection>; + + async fn from_request(req: Request, state: &S) -> Result { + let (parts, body) = req.into_parts(); + let bytes = Bytes::from_request(Request::from_parts(parts.clone(), body), state) + .await + .map_err(EitherRejection::Bytes)?; + + $( + let req = Request::from_parts( + parts.clone(), + axum::body::Body::new(http_body_util::Full::new(bytes.clone())), + ); + if let Ok(extracted) = $ident::from_request(req, state).await { + return Ok(Self::$ident(extracted)); + } + )* + + let req = Request::from_parts( + parts.clone(), + axum::body::Body::new(http_body_util::Full::new(bytes.clone())), + ); + match $last::from_request(req, state).await { + Ok(extracted) => Ok(Self::$last(extracted)), + Err(error) => Err(EitherRejection::LastRejection(error)), + } + } + } + impl<$($ident),*, $last> IntoResponse for $either<$($ident),*, $last> where $($ident: IntoResponse),*, @@ -312,3 +408,67 @@ where } } } + +#[cfg(test)] +mod tests { + use std::future::Future; + + use axum::body::Body; + use axum::extract::rejection::StringRejection; + use axum::extract::{FromRequest, Request, State}; + use bytes::Bytes; + use http_body_util::Full; + + use super::*; + + struct False; + + impl FromRequestParts for False { + type Rejection = (); + + fn from_request_parts( + _parts: &mut Parts, + _state: &S, + ) -> impl Future> + Send { + std::future::ready(Err(())) + } + } + + #[tokio::test] + async fn either_from_request() { + // The body is by design not valid UTF-8. + let request = Request::new(Body::new(Full::new(Bytes::from_static(&[255])))); + + let either = Either4::::from_request(request, &()) + .await + .unwrap(); + + assert!(matches!(either, Either4::E3(_))); + } + + #[tokio::test] + async fn either_from_request_rejection() { + // The body is by design not valid UTF-8. + let request = Request::new(Body::new(Full::new(Bytes::from_static(&[255])))); + + let either = Either::::from_request(request, &()) + .await + .unwrap_err(); + + assert!(matches!( + either, + EitherRejection::LastRejection(StringRejection::InvalidUtf8(_)) + )); + } + + #[tokio::test] + async fn either_from_request_parts() { + let (mut parts, _) = Request::new(Body::empty()).into_parts(); + + let either = Either3::>::from_request_parts(&mut parts, &()) + .await + .unwrap(); + + assert!(matches!(either, Either3::E3(State(())))); + } +}