Add trait IntoResponseHeaders (#649)

* Introduce IntoResponseHeaders trait

* Implement IntoResponseHeaders for HeaderMap

* Add impl IntoResponse for impl IntoResponseHeaders

… and update IntoResponse impls that use HeaderMap to be generic instead.

* Add impl IntoResponseHeaders for Headers

… and remove IntoResponse impls that use it.

* axum-debug: Fix grammar in docs

* Explain confusing error message in docs

Author: jplatte

axum-core/src/response/headers.rs

@@ -1,11 +1,8 @@
-use super::{IntoResponse, Response};
-use crate::body::boxed;
-use bytes::Bytes;
+use super::{IntoResponse, IntoResponseHeaders, Response};
 use http::{
-    header::{HeaderMap, HeaderName, HeaderValue},
+    header::{HeaderName, HeaderValue},
     StatusCode,
 };
-use http_body::{Empty, Full};
 use std::{convert::TryInto, fmt};
 
 /// A response with headers.
@@ -54,102 +51,53 @@ use std::{convert::TryInto, fmt};
 #[derive(Clone, Copy, Debug)]
 pub struct Headers<H>(pub H);
 
-impl<H> Headers<H> {
-    fn try_into_header_map<K, V>(self) -> Result<HeaderMap, Response>
-    where
-        H: IntoIterator<Item = (K, V)>,
-        K: TryInto<HeaderName>,
-        K::Error: fmt::Display,
-        V: TryInto<HeaderValue>,
-        V::Error: fmt::Display,
-    {
-        self.0
-            .into_iter()
-            .map(|(key, value)| {
-                let key = key.try_into().map_err(Either::A)?;
-                let value = value.try_into().map_err(Either::B)?;
-                Ok((key, value))
-            })
-            .collect::<Result<_, _>>()
-            .map_err(|err| {
-                let err = match err {
-                    Either::A(err) => err.to_string(),
-                    Either::B(err) => err.to_string(),
-                };
-
-                let body = boxed(Full::new(Bytes::copy_from_slice(err.as_bytes())));
-                let mut res = Response::new(body);
-                *res.status_mut() = StatusCode::INTERNAL_SERVER_ERROR;
-                res
-            })
-    }
-}
-
-impl<H, K, V> IntoResponse for Headers<H>
+impl<H, K, V> IntoResponseHeaders for Headers<H>
 where
     H: IntoIterator<Item = (K, V)>,
     K: TryInto<HeaderName>,
     K::Error: fmt::Display,
     V: TryInto<HeaderValue>,
     V::Error: fmt::Display,
 {
-    fn into_response(self) -> Response {
-        let headers = self.try_into_header_map();
-
-        match headers {
-            Ok(headers) => {
-                let mut res = Response::new(boxed(Empty::new()));
-                *res.headers_mut() = headers;
-                res
-            }
-            Err(err) => err,
+    type IntoIter = IntoIter<H::IntoIter>;
+
+    fn into_headers(self) -> Self::IntoIter {
+        IntoIter {
+            inner: self.0.into_iter(),
         }
     }
 }
 
-impl<H, T, K, V> IntoResponse for (Headers<H>, T)
-where
-    T: IntoResponse,
-    H: IntoIterator<Item = (K, V)>,
-    K: TryInto<HeaderName>,
-    K::Error: fmt::Display,
-    V: TryInto<HeaderValue>,
-    V::Error: fmt::Display,
-{
-    fn into_response(self) -> Response {
-        let headers = match self.0.try_into_header_map() {
-            Ok(headers) => headers,
-            Err(res) => return res,
-        };
-
-        (headers, self.1).into_response()
-    }
+#[doc(hidden)]
+#[derive(Debug)]
+pub struct IntoIter<H> {
+    inner: H,
 }
 
-impl<H, T, K, V> IntoResponse for (StatusCode, Headers<H>, T)
+impl<H, K, V> Iterator for IntoIter<H>
 where
-    T: IntoResponse,
-    H: IntoIterator<Item = (K, V)>,
+    H: Iterator<Item = (K, V)>,
     K: TryInto<HeaderName>,
     K::Error: fmt::Display,
     V: TryInto<HeaderValue>,
     V::Error: fmt::Display,
 {
-    fn into_response(self) -> Response {
-        let headers = match self.1.try_into_header_map() {
-            Ok(headers) => headers,
-            Err(res) => return res,
-        };
-
-        (self.0, headers, self.2).into_response()
+    type Item = Result<(Option<HeaderName>, HeaderValue), Response>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.inner.next().map(|(key, value)| {
+            let key = key
+                .try_into()
+                .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()).into_response())?;
+            let value = value
+                .try_into()
+                .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()).into_response())?;
+
+            Ok((Some(key), value))
+        })
     }
 }
 
-enum Either<A, B> {
-    A(A),
-    B(B),
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;

axum-core/src/response/mod.rs

@@ -10,14 +10,14 @@ use crate::{
 };
 use bytes::Bytes;
 use http::{
-    header::{self, HeaderMap, HeaderValue},
+    header::{self, HeaderMap, HeaderName, HeaderValue},
     StatusCode,
 };
 use http_body::{
     combinators::{MapData, MapErr},
     Empty, Full,
 };
-use std::{borrow::Cow, convert::Infallible};
+use std::{borrow::Cow, convert::Infallible, iter};
 
 mod headers;
 
@@ -148,6 +148,42 @@ pub trait IntoResponse {
     fn into_response(self) -> Response;
 }
 
+/// Trait for generating response headers.
+///
+
+/// **Note: If you see this trait not being implemented in an error message, you are almost
+/// certainly being mislead by the compiler¹. Look for the following snippet in the output and
+/// check [`IntoResponse`]'s documentation if you find it:**
+///
+/// ```text
+/// note: required because of the requirements on the impl of `IntoResponse` for `<type>`
+/// ```
+///
+/// Any type that implements this trait automatically implements `IntoResponse` as well, but can
+/// also be used in a tuple like `(StatusCode, Self)`, `(Self, impl IntoResponseHeaders)`,
+/// `(StatusCode, Self, impl IntoResponseHeaders, impl IntoResponse)` and so on.
+///
+/// This trait can't currently be implemented outside of axum.
+///
+/// ¹ See also [this rustc issue](https://github.com/rust-lang/rust/issues/22590)
+pub trait IntoResponseHeaders {
+    /// The return type of [`.into_headers()`].
+    ///
+    /// The iterator item is a `Result` to allow the implementation to return a server error
+    /// instead.
+    ///
+    /// The header name is optional because `HeaderMap`s iterator doesn't yield it multiple times
+    /// for headers that have multiple values, to avoid unnecessary copies.
+    #[doc(hidden)]
+    type IntoIter: IntoIterator<Item = Result<(Option<HeaderName>, HeaderValue), Response>>;
+
+    /// Attempt to turn `self` into a list of headers.
+    ///
+    /// In practice, only the implementation for `axum::response::Headers` ever returns `Err(_)`.
+    #[doc(hidden)]
+    fn into_headers(self) -> Self::IntoIter;
+}
+
 impl IntoResponse for () {
     fn into_response(self) -> Response {
         Response::new(boxed(Empty::new()))
@@ -320,6 +356,21 @@ impl IntoResponse for StatusCode {
     }
 }
 
+impl<H> IntoResponse for H
+where
+    H: IntoResponseHeaders,
+{
+    fn into_response(self) -> Response {
+        let mut res = Response::new(boxed(Empty::new()));
+
+        if let Err(e) = try_extend_headers(res.headers_mut(), self.into_headers()) {
+            return e;
+        }
+
+        res
+    }
+}
+
 impl<T> IntoResponse for (StatusCode, T)
 where
     T: IntoResponse,
@@ -331,33 +382,98 @@ where
     }
 }
 
-impl<T> IntoResponse for (HeaderMap, T)
+impl<H, T> IntoResponse for (H, T)
 where
+    H: IntoResponseHeaders,
     T: IntoResponse,
 {
     fn into_response(self) -> Response {
         let mut res = self.1.into_response();
-        res.headers_mut().extend(self.0);
+
+        if let Err(e) = try_extend_headers(res.headers_mut(), self.0.into_headers()) {
+            return e;
+        }
+
         res
     }
 }
 
-impl<T> IntoResponse for (StatusCode, HeaderMap, T)
+impl<H, T> IntoResponse for (StatusCode, H, T)
 where
+    H: IntoResponseHeaders,
     T: IntoResponse,
 {
     fn into_response(self) -> Response {
         let mut res = self.2.into_response();
         *res.status_mut() = self.0;
-        res.headers_mut().extend(self.1);
+
+        if let Err(e) = try_extend_headers(res.headers_mut(), self.1.into_headers()) {
+            return e;
+        }
+
         res
     }
 }
 
-impl IntoResponse for HeaderMap {
-    fn into_response(self) -> Response {
-        let mut res = Response::new(boxed(Empty::new()));
-        *res.headers_mut() = self;
-        res
+impl IntoResponseHeaders for HeaderMap {
+    // FIXME: Use type_alias_impl_trait when available
+    type IntoIter = iter::Map<
+        http::header::IntoIter<HeaderValue>,
+        fn(
+            (Option<HeaderName>, HeaderValue),
+        ) -> Result<(Option<HeaderName>, HeaderValue), Response>,
+    >;
+
+    fn into_headers(self) -> Self::IntoIter {
+        self.into_iter().map(Ok)
+    }
+}
+
+// Slightly adjusted version of `impl<T> Extend<(Option<HeaderName>, T)> for HeaderMap<T>`.
+// Accepts an iterator that returns Results and short-circuits on an `Err`.
+fn try_extend_headers(
+    headers: &mut HeaderMap,
+    iter: impl IntoIterator<Item = Result<(Option<HeaderName>, HeaderValue), Response>>,
+) -> Result<(), Response> {
+    use http::header::Entry;
+
+    let mut iter = iter.into_iter();
+
+    // The structure of this is a bit weird, but it is mostly to make the
+    // borrow checker happy.
+    let (mut key, mut val) = match iter.next().transpose()? {
+        Some((Some(key), val)) => (key, val),
+        Some((None, _)) => panic!("expected a header name, but got None"),
+        None => return Ok(()),
+    };
+
+    'outer: loop {
+        let mut entry = match headers.entry(key) {
+            Entry::Occupied(mut e) => {
+                // Replace all previous values while maintaining a handle to
+                // the entry.
+                e.insert(val);
+                e
+            }
+            Entry::Vacant(e) => e.insert_entry(val),
+        };
+
+        // As long as `HeaderName` is none, keep inserting the value into
+        // the current entry
+        loop {
+            match iter.next().transpose()? {
+                Some((Some(k), v)) => {
+                    key = k;
+                    val = v;
+                    continue 'outer;
+                }
+                Some((None, v)) => {
+                    entry.append(v);
+                }
+                None => {
+                    return Ok(());
+                }
+            }
+        }
     }
 }

axum-debug/src/lib.rs

@@ -21,7 +21,7 @@
 //! ```
 //!
 //! You will get a long error message about function not implementing [`Handler`] trait. But why
-//! this function does not implement it? To figure it out [`debug_handler`] macro can be used.
+//! does this function not implement it? To figure it out, the [`debug_handler`] macro can be used.
 //!
 //! ```rust,compile_fail
 //! # use axum::{routing::get, Router};
@@ -89,6 +89,34 @@
 //! async fn handler(request: Request<BoxBody>) {}
 //! ```
 //!
+//! # Known limitations
+//!
+//! If your response type doesn't implement `IntoResponse`, you will get a slightly confusing error
+//! message:
+//!
+//! ```text
+//! error[E0277]: the trait bound `bool: axum_core::response::IntoResponseHeaders` is not satisfied
+//!  --> tests/fail/wrong_return_type.rs:4:23
+//!   |
+//! 4 | async fn handler() -> bool {
+//!   |                       ^^^^ the trait `axum_core::response::IntoResponseHeaders` is not implemented for `bool`
+//!   |
+//!   = note: required because of the requirements on the impl of `IntoResponse` for `bool`
+//! note: required by a bound in `__axum_debug_check_handler_into_response::{closure#0}::check`
+//!  --> tests/fail/wrong_return_type.rs:4:23
+//!   |
+//! 4 | async fn handler() -> bool {
+//!   |                       ^^^^ required by this bound in `__axum_debug_check_handler_into_response::{closure#0}::check`
+//! ```
+//!
+//! The main error message when `IntoResponse` isn't implemented will also ways be for a different
+//! trait, `IntoResponseHeaders`, not being implemented. This trait is not meant to be implemented
+//! for types outside of axum and what you really need to do is change your return type or implement
+//! `IntoResponse` for it (if it is your own type that you want to return directly from handlers).
+//!
+//! This issue is not specific to axum and cannot be fixed by us. For more details, see the
+//! [rustc issue about it](https://github.com/rust-lang/rust/issues/22590).
+//!
 //! # Performance
 //!
 //! Macros in this crate have no effect when using release profile. (eg. `cargo build --release`)

axum-debug/tests/fail/wrong_return_type.stderr

@@ -1,9 +1,10 @@
-error[E0277]: the trait bound `bool: IntoResponse` is not satisfied
+error[E0277]: the trait bound `bool: axum_core::response::IntoResponseHeaders` is not satisfied
  --> tests/fail/wrong_return_type.rs:4:23
   |
 4 | async fn handler() -> bool {
-  |                       ^^^^ the trait `IntoResponse` is not implemented for `bool`
+  |                       ^^^^ the trait `axum_core::response::IntoResponseHeaders` is not implemented for `bool`
   |
+  = note: required because of the requirements on the impl of `IntoResponse` for `bool`
 note: required by a bound in `__axum_debug_check_handler_into_response::{closure#0}::check`
  --> tests/fail/wrong_return_type.rs:4:23
   |