Move TypedHeader to axum-extra (#1850)

Co-authored-by: Michael Scofield <[email protected]>
Co-authored-by: Jonas Platte <[email protected]>

axum-extra/src/extract/cookie/mod.rs

@@ -41,12 +41,14 @@ pub use cookie::Key;
 /// use axum::{
 ///     Router,
 ///     routing::{post, get},
-///     extract::TypedHeader,
 ///     response::{IntoResponse, Redirect},
-///     headers::authorization::{Authorization, Bearer},
 ///     http::StatusCode,
 /// };
-/// use axum_extra::extract::cookie::{CookieJar, Cookie};
+/// use axum_extra::{
+///     TypedHeader,
+///     headers::authorization::{Authorization, Bearer},
+///     extract::cookie::{CookieJar, Cookie},
+/// };
 ///
 /// async fn create_session(
 ///     TypedHeader(auth): TypedHeader<Authorization<Bearer>>,

axum-extra/src/extract/cookie/private.rs

@@ -23,12 +23,15 @@ use std::{convert::Infallible, fmt, marker::PhantomData};
 /// use axum::{
 ///     Router,
 ///     routing::{post, get},
-///     extract::{TypedHeader, FromRef},
+///     extract::FromRef,
 ///     response::{IntoResponse, Redirect},
-///     headers::authorization::{Authorization, Bearer},
 ///     http::StatusCode,
 /// };
-/// use axum_extra::extract::cookie::{PrivateCookieJar, Cookie, Key};
+/// use axum_extra::{
+///     TypedHeader,
+///     headers::authorization::{Authorization, Bearer},
+///     extract::cookie::{PrivateCookieJar, Cookie, Key},
+/// };
 ///
 /// async fn set_secret(
 ///     jar: PrivateCookieJar,

axum-extra/src/extract/cookie/signed.rs

@@ -24,12 +24,15 @@ use std::{convert::Infallible, fmt, marker::PhantomData};
 /// use axum::{
 ///     Router,
 ///     routing::{post, get},
-///     extract::{TypedHeader, FromRef},
+///     extract::FromRef,
 ///     response::{IntoResponse, Redirect},
-///     headers::authorization::{Authorization, Bearer},
 ///     http::StatusCode,
 /// };
-/// use axum_extra::extract::cookie::{SignedCookieJar, Cookie, Key};
+/// use axum_extra::{
+///     TypedHeader,
+///     headers::authorization::{Authorization, Bearer},
+///     extract::cookie::{SignedCookieJar, Cookie, Key},
+/// };
 ///
 /// async fn create_session(
 ///     TypedHeader(auth): TypedHeader<Authorization<Bearer>>,

axum-extra/src/extract/mod.rs

@@ -39,3 +39,7 @@ pub use self::multipart::Multipart;
 #[cfg(feature = "json-lines")]
 #[doc(no_inline)]
 pub use crate::json_lines::JsonLines;
+
+#[cfg(feature = "typed-header")]
+#[doc(no_inline)]
+pub use crate::typed_header::TypedHeader;

axum-extra/src/lib.rs

@@ -21,6 +21,7 @@
 //! `protobuf` | Enables the `Protobuf` extractor and response | No
 //! `query` | Enables the `Query` extractor | No
 //! `typed-routing` | Enables the `TypedPath` routing utilities | No
+//! `typed-header` | Enables the `TypedHeader` extractor and response  | No
 //!
 //! [`axum`]: https://crates.io/crates/axum
 
@@ -80,6 +81,17 @@ pub mod routing;
 #[cfg(feature = "json-lines")]
 pub mod json_lines;
 
+#[cfg(feature = "typed-header")]
+pub mod typed_header;
+
+#[cfg(feature = "typed-header")]
+#[doc(no_inline)]
+pub use headers;
+
+#[cfg(feature = "typed-header")]
+#[doc(inline)]
+pub use typed_header::TypedHeader;
+
 #[cfg(feature = "protobuf")]
 pub mod protobuf;
 

axum-extra/src/response/mod.rs

@@ -88,3 +88,7 @@ mime_response! {
     Wasm,
     "application/wasm",
 }
+
+#[cfg(feature = "typed-header")]
+#[doc(no_inline)]
+pub use crate::typed_header::TypedHeader;

axum-extra/src/typed_header.rs

@@ -0,0 +1,209 @@
+//! Extractor and response for typed headers.
+
+use axum::{
+    async_trait,
+    extract::FromRequestParts,
+    response::{IntoResponse, IntoResponseParts, Response, ResponseParts},
+};
+use headers::{Header, HeaderMapExt};
+use http::request::Parts;
+use std::convert::Infallible;
+
+/// Extractor and response that works with typed header values from [`headers`].
+///
+/// # As extractor
+///
+/// In general, it's recommended to extract only the needed headers via `TypedHeader` rather than
+/// removing all headers with the `HeaderMap` extractor.
+///
+/// ```rust,no_run
+/// use axum::{
+///     routing::get,
+///     Router,
+/// };
+/// use headers::UserAgent;
+/// use axum_extra::TypedHeader;
+///
+/// async fn users_teams_show(
+///     TypedHeader(user_agent): TypedHeader<UserAgent>,
+/// ) {
+///     // ...
+/// }
+///
+/// let app = Router::new().route("/users/:user_id/team/:team_id", get(users_teams_show));
+/// # let _: Router = app;
+/// ```
+///
+/// # As response
+///
+/// ```rust
+/// use axum::{
+///     response::IntoResponse,
+/// };
+/// use headers::ContentType;
+/// use axum_extra::TypedHeader;
+///
+/// async fn handler() -> (TypedHeader<ContentType>, &'static str) {
+///     (
+///         TypedHeader(ContentType::text_utf8()),
+///         "Hello, World!",
+///     )
+/// }
+/// ```
+#[cfg(feature = "typed-header")]
+#[derive(Debug, Clone, Copy)]
+#[must_use]
+pub struct TypedHeader<T>(pub T);
+
+#[async_trait]
+impl<T, S> FromRequestParts<S> for TypedHeader<T>
+where
+    T: Header,
+    S: Send + Sync,
+{
+    type Rejection = TypedHeaderRejection;
+
+    async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
+        let mut values = parts.headers.get_all(T::name()).iter();
+        let is_missing = values.size_hint() == (0, Some(0));
+        T::decode(&mut values)
+            .map(Self)
+            .map_err(|err| TypedHeaderRejection {
+                name: T::name(),
+                reason: if is_missing {
+                    // Report a more precise rejection for the missing header case.
+                    TypedHeaderRejectionReason::Missing
+                } else {
+                    TypedHeaderRejectionReason::Error(err)
+                },
+            })
+    }
+}
+
+axum_core::__impl_deref!(TypedHeader);
+
+impl<T> IntoResponseParts for TypedHeader<T>
+where
+    T: Header,
+{
+    type Error = Infallible;
+
+    fn into_response_parts(self, mut res: ResponseParts) -> Result<ResponseParts, Self::Error> {
+        res.headers_mut().typed_insert(self.0);
+        Ok(res)
+    }
+}
+
+impl<T> IntoResponse for TypedHeader<T>
+where
+    T: Header,
+{
+    fn into_response(self) -> Response {
+        let mut res = ().into_response();
+        res.headers_mut().typed_insert(self.0);
+        res
+    }
+}
+
+/// Rejection used for [`TypedHeader`](TypedHeader).
+#[cfg(feature = "typed-header")]
+#[derive(Debug)]
+pub struct TypedHeaderRejection {
+    name: &'static http::header::HeaderName,
+    reason: TypedHeaderRejectionReason,
+}
+
+impl TypedHeaderRejection {
+    /// Name of the header that caused the rejection
+    pub fn name(&self) -> &http::header::HeaderName {
+        self.name
+    }
+
+    /// Reason why the header extraction has failed
+    pub fn reason(&self) -> &TypedHeaderRejectionReason {
+        &self.reason
+    }
+}
+
+/// Additional information regarding a [`TypedHeaderRejection`]
+#[cfg(feature = "typed-header")]
+#[derive(Debug)]
+#[non_exhaustive]
+pub enum TypedHeaderRejectionReason {
+    /// The header was missing from the HTTP request
+    Missing,
+    /// An error occured when parsing the header from the HTTP request
+    Error(headers::Error),
+}
+
+impl IntoResponse for TypedHeaderRejection {
+    fn into_response(self) -> Response {
+        (http::StatusCode::BAD_REQUEST, self.to_string()).into_response()
+    }
+}
+
+impl std::fmt::Display for TypedHeaderRejection {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match &self.reason {
+            TypedHeaderRejectionReason::Missing => {
+                write!(f, "Header of type `{}` was missing", self.name)
+            }
+            TypedHeaderRejectionReason::Error(err) => {
+                write!(f, "{} ({})", err, self.name)
+            }
+        }
+    }
+}
+
+impl std::error::Error for TypedHeaderRejection {
+    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
+        match &self.reason {
+            TypedHeaderRejectionReason::Error(err) => Some(err),
+            TypedHeaderRejectionReason::Missing => None,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::test_helpers::*;
+    use axum::{response::IntoResponse, routing::get, Router};
+
+    #[tokio::test]
+    async fn typed_header() {
+        async fn handle(
+            TypedHeader(user_agent): TypedHeader<headers::UserAgent>,
+            TypedHeader(cookies): TypedHeader<headers::Cookie>,
+        ) -> impl IntoResponse {
+            let user_agent = user_agent.as_str();
+            let cookies = cookies.iter().collect::<Vec<_>>();
+            format!("User-Agent={user_agent:?}, Cookie={cookies:?}")
+        }
+
+        let app = Router::new().route("/", get(handle));
+
+        let client = TestClient::new(app);
+
+        let res = client
+            .get("/")
+            .header("user-agent", "foobar")
+            .header("cookie", "a=1; b=2")
+            .header("cookie", "c=3")
+            .send()
+            .await;
+        let body = res.text().await;
+        assert_eq!(
+            body,
+            r#"User-Agent="foobar", Cookie=[("a", "1"), ("b", "2"), ("c", "3")]"#
+        );
+
+        let res = client.get("/").header("user-agent", "foobar").send().await;
+        let body = res.text().await;
+        assert_eq!(body, r#"User-Agent="foobar", Cookie=[]"#);
+
+        let res = client.get("/").header("cookie", "a=1").send().await;
+        let body = res.text().await;
+        assert_eq!(body, "Header of type `user-agent` was missing");
+    }
+}