Skip to content

Commit

Permalink
Add trait IntoResponseHeaders (#649)
Browse files Browse the repository at this point in the history
* 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
  • Loading branch information
jplatte authored Jan 23, 2022
1 parent 9f5dc45 commit 5bb8df1
Show file tree
Hide file tree
Showing 4 changed files with 186 additions and 93 deletions.
106 changes: 27 additions & 79 deletions axum-core/src/response/headers.rs
Original file line number Diff line number Diff line change
@@ -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.
Expand Down Expand Up @@ -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::*;
Expand Down
138 changes: 127 additions & 11 deletions axum-core/src/response/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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()))
Expand Down Expand Up @@ -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,
Expand All @@ -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(());
}
}
}
}
}
30 changes: 29 additions & 1 deletion axum-debug/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -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`)
Expand Down
5 changes: 3 additions & 2 deletions axum-debug/tests/fail/wrong_return_type.stderr
Original file line number Diff line number Diff line change
@@ -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
|
Expand Down

0 comments on commit 5bb8df1

Please sign in to comment.