From fd96bce34d9e5428c0325b364feb60ccc03ec3e4 Mon Sep 17 00:00:00 2001
From: Ingvar Stepanyan <me@rreverser.com>
Date: Fri, 3 Mar 2023 14:35:49 +0000
Subject: [PATCH] Gracefully handle missing headers (#1810)

Co-authored-by: David Pedersen <david.pdrsn@gmail.com>
---
 axum/CHANGELOG.md        |  2 ++
 axum/src/typed_header.rs | 44 ++++++++++++++++++++++++++++------------
 2 files changed, 33 insertions(+), 13 deletions(-)

diff --git a/axum/CHANGELOG.md b/axum/CHANGELOG.md
index 5e6734c88f..d6f6f5f17d 100644
--- a/axum/CHANGELOG.md
+++ b/axum/CHANGELOG.md
@@ -9,11 +9,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - **fixed:** Add `#[must_use]` attributes to types that do nothing unless used ([#1809])
 - **fixed:** Add `#[must_use]` to `WebSocketUpgrade::on_upgrade` ([#1801])
+- **fixed:** Gracefully handle missing headers in the `TypedHeader` extractor ([#1810])
 - **fixed:** Fix routing issues when loading a `Router` via a dynamic library ([#1806])
 
 [#1801]: https://github.com/tokio-rs/axum/pull/1801
 [#1806]: https://github.com/tokio-rs/axum/pull/1806
 [#1809]: https://github.com/tokio-rs/axum/pull/1809
+[#1810]: https://github.com/tokio-rs/axum/pull/1810
 
 # 0.6.9 (24. February, 2023)
 
diff --git a/axum/src/typed_header.rs b/axum/src/typed_header.rs
index c047af74c5..717bc2409a 100644
--- a/axum/src/typed_header.rs
+++ b/axum/src/typed_header.rs
@@ -62,17 +62,19 @@ where
     type Rejection = TypedHeaderRejection;
 
     async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
-        match parts.headers.typed_try_get::<T>() {
-            Ok(Some(value)) => Ok(Self(value)),
-            Ok(None) => Err(TypedHeaderRejection {
+        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: TypedHeaderRejectionReason::Missing,
-            }),
-            Err(err) => Err(TypedHeaderRejection {
-                name: T::name(),
-                reason: TypedHeaderRejectionReason::Error(err),
-            }),
-        }
+                reason: if is_missing {
+                    // Report a more precise rejection for the missing header case.
+                    TypedHeaderRejectionReason::Missing
+                } else {
+                    TypedHeaderRejectionReason::Error(err)
+                },
+            })
     }
 }
 
@@ -175,19 +177,35 @@ mod tests {
     async fn typed_header() {
         async fn handle(
             TypedHeader(user_agent): TypedHeader<headers::UserAgent>,
+            TypedHeader(cookies): TypedHeader<headers::Cookie>,
         ) -> impl IntoResponse {
-            user_agent.to_string()
+            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, "foobar");
+        assert_eq!(body, r#"User-Agent="foobar", Cookie=[]"#);
 
-        let res = client.get("/").send().await;
+        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");
     }