From 97c140cdf727f4dc0200a3ac54eb988fcad412b3 Mon Sep 17 00:00:00 2001 From: David Pedersen Date: Tue, 17 Aug 2021 17:28:02 +0200 Subject: [PATCH] Add `Headers` response (#193) * Add `Headers` Example usage: ```rust use axum::{ route, routing::RoutingDsl, response::{IntoResponse, Headers}, handler::get, }; use http::header::{HeaderName, HeaderValue}; // It works with any `IntoIterator` where `Key` can be // turned into a `HeaderName` and `Value` can be turned into a `HeaderValue` // // Such as `Vec<(HeaderName, HeaderValue)>` async fn just_headers() -> impl IntoResponse { Headers(vec![ (HeaderName::from_static("X-Foo"), HeaderValue::from_static("foo")), ]) } // Or `[(&str, &str)]` async fn from_strings() -> impl IntoResponse { Headers([("X-Foo", "foo")]) } ``` Fixes https://github.com/tokio-rs/axum/issues/187 * Make work on Rust versions without `IntoIterator` for arrays * format * changelog --- .clippy.toml | 2 +- CHANGELOG.md | 3 +- src/response/headers.rs | 229 ++++++++++++++++++++++++++++++++++++++++ src/response/mod.rs | 17 +-- 4 files changed, 243 insertions(+), 8 deletions(-) create mode 100644 src/response/headers.rs diff --git a/.clippy.toml b/.clippy.toml index 0f404fa5..829dd1c5 100644 --- a/.clippy.toml +++ b/.clippy.toml @@ -1 +1 @@ -msrv = "1.40" +msrv = "1.51" diff --git a/CHANGELOG.md b/CHANGELOG.md index c829056d..a17c5f36 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,7 +14,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Add `NestedUri` for extracting request URI in nested services ([#161](https://github.com/tokio-rs/axum/pull/161)) - Implement `FromRequest` for `http::Extensions` - Implement SSE as an `IntoResponse` instead of a service ([#98](https://github.com/tokio-rs/axum/pull/98)) -- Add `Redirect` response. ([#192](https://github.com/tokio-rs/axum/pull/192)) +- Add `Headers` for easily customizing headers on a response ([#193](https://github.com/tokio-rs/axum/pull/193)) +- Add `Redirect` response ([#192](https://github.com/tokio-rs/axum/pull/192)) - Make `RequestParts::{new, try_into_request}` public ([#194](https://github.com/tokio-rs/axum/pull/194)) ## Breaking changes diff --git a/src/response/headers.rs b/src/response/headers.rs new file mode 100644 index 00000000..6df6c081 --- /dev/null +++ b/src/response/headers.rs @@ -0,0 +1,229 @@ +use super::IntoResponse; +use crate::body::{box_body, BoxBody}; +use bytes::Bytes; +use http::header::{HeaderMap, HeaderName, HeaderValue}; +use http::{Response, StatusCode}; +use http_body::{Body, Full}; +use std::{convert::TryInto, fmt}; +use tower::{util::Either, BoxError}; + +/// A response with headers. +/// +/// # Example +/// +/// ```rust +/// use axum::{ +/// route, +/// routing::RoutingDsl, +/// response::{IntoResponse, Headers}, +/// handler::get, +/// }; +/// use http::header::{HeaderName, HeaderValue}; +/// +/// // It works with any `IntoIterator` where `Key` can be +/// // turned into a `HeaderName` and `Value` can be turned into a `HeaderValue` +/// // +/// // Such as `Vec<(HeaderName, HeaderValue)>` +/// async fn just_headers() -> impl IntoResponse { +/// Headers(vec![ +/// (HeaderName::from_static("X-Foo"), HeaderValue::from_static("foo")), +/// ]) +/// } +/// +/// // Or `Vec<(&str, &str)>` +/// async fn from_strings() -> impl IntoResponse { +/// Headers(vec![("X-Foo", "foo")]) +/// } +/// +/// // Or `[(&str, &str)]` if you're on Rust 1.53+ +/// +/// let app = route("/just-headers", get(just_headers)) +/// .route("/from-strings", get(from_strings)); +/// # async { +/// # axum::Server::bind(&"".parse().unwrap()).serve(app.into_make_service()).await.unwrap(); +/// # }; +/// ``` +/// +/// If a conversion to `HeaderName` or `HeaderValue` fails a `500 Internal +/// Server Error` response will be returned. +/// +/// You can also return `(Headers, impl IntoResponse)` to customize the headers +/// of a response, or `(StatusCode, Headeres, impl IntoResponse)` to customize +/// the status code and headers. +#[derive(Clone, Copy, Debug)] +pub struct Headers(pub H); + +impl Headers { + fn try_into_header_map(self) -> Result>> + where + H: IntoIterator, + K: TryInto, + K::Error: fmt::Display, + V: TryInto, + 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::>() + .map_err(|err| { + let err = match err { + Either::A(err) => err.to_string(), + Either::B(err) => err.to_string(), + }; + + let body = Full::new(Bytes::copy_from_slice(err.as_bytes())); + let mut res = Response::new(body); + *res.status_mut() = StatusCode::INTERNAL_SERVER_ERROR; + res + }) + } +} + +impl IntoResponse for Headers +where + H: IntoIterator, + K: TryInto, + K::Error: fmt::Display, + V: TryInto, + V::Error: fmt::Display, +{ + type Body = Full; + type BodyError = ::Error; + + fn into_response(self) -> http::Response { + let headers = self.try_into_header_map(); + + match headers { + Ok(headers) => { + let mut res = Response::new(Full::new(Bytes::new())); + *res.headers_mut() = headers; + res + } + Err(err) => err, + } + } +} + +impl IntoResponse for (Headers, T) +where + T: IntoResponse, + T::Body: Body + Send + Sync + 'static, + ::Error: Into, + H: IntoIterator, + K: TryInto, + K::Error: fmt::Display, + V: TryInto, + V::Error: fmt::Display, +{ + type Body = BoxBody; + type BodyError = ::Error; + + // this boxing could be improved with a EitherBody but thats + // an issue for another time + fn into_response(self) -> Response { + let headers = match self.0.try_into_header_map() { + Ok(headers) => headers, + Err(res) => return res.map(box_body), + }; + + (headers, self.1).into_response().map(box_body) + } +} + +impl IntoResponse for (StatusCode, Headers, T) +where + T: IntoResponse, + T::Body: Body + Send + Sync + 'static, + ::Error: Into, + H: IntoIterator, + K: TryInto, + K::Error: fmt::Display, + V: TryInto, + V::Error: fmt::Display, +{ + type Body = BoxBody; + type BodyError = ::Error; + + fn into_response(self) -> Response { + let headers = match self.1.try_into_header_map() { + Ok(headers) => headers, + Err(res) => return res.map(box_body), + }; + + (self.0, headers, self.2).into_response().map(box_body) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use futures_util::FutureExt; + use http::header::USER_AGENT; + + #[test] + fn vec_of_header_name_and_value() { + let res = Headers(vec![(USER_AGENT, HeaderValue::from_static("axum"))]).into_response(); + + assert_eq!(res.headers()["user-agent"], "axum"); + assert_eq!(res.status(), StatusCode::OK); + } + + #[test] + fn vec_of_strings() { + let res = Headers(vec![("user-agent", "axum")]).into_response(); + + assert_eq!(res.headers()["user-agent"], "axum"); + } + + #[test] + fn with_body() { + let res = (Headers(vec![("user-agent", "axum")]), "foo").into_response(); + + assert_eq!(res.headers()["user-agent"], "axum"); + let body = hyper::body::to_bytes(res.into_body()) + .now_or_never() + .unwrap() + .unwrap(); + assert_eq!(&body[..], b"foo"); + } + + #[test] + fn with_status_and_body() { + let res = ( + StatusCode::NOT_FOUND, + Headers(vec![("user-agent", "axum")]), + "foo", + ) + .into_response(); + + assert_eq!(res.headers()["user-agent"], "axum"); + assert_eq!(res.status(), StatusCode::NOT_FOUND); + let body = hyper::body::to_bytes(res.into_body()) + .now_or_never() + .unwrap() + .unwrap(); + assert_eq!(&body[..], b"foo"); + } + + #[test] + fn invalid_header_name() { + let bytes: &[u8] = &[0, 159, 146, 150]; // invalid utf-8 + let res = Headers(vec![(bytes, "axum")]).into_response(); + + assert_eq!(res.status(), StatusCode::INTERNAL_SERVER_ERROR); + } + + #[test] + fn invalid_header_value() { + let bytes: &[u8] = &[0, 159, 146, 150]; // invalid utf-8 + let res = Headers(vec![("user-agent", bytes)]).into_response(); + + assert!(res.headers().get("user-agent").is_none()); + assert_eq!(res.status(), StatusCode::INTERNAL_SERVER_ERROR); + } +} diff --git a/src/response/mod.rs b/src/response/mod.rs index 39919e30..7aa767ba 100644 --- a/src/response/mod.rs +++ b/src/response/mod.rs @@ -13,15 +13,20 @@ use http_body::{ use std::{borrow::Cow, convert::Infallible}; use tower::{util::Either, BoxError}; +mod headers; +mod redirect; + +pub mod sse; + #[doc(no_inline)] pub use crate::Json; -mod redirect; - -pub use self::redirect::Redirect; - -pub mod sse; -pub use sse::{sse, Sse}; +#[doc(inline)] +pub use self::{ + headers::Headers, + redirect::Redirect, + sse::{sse, Sse}, +}; /// Trait for generating responses. ///