diff --git a/axum-extra/src/routing/mod.rs b/axum-extra/src/routing/mod.rs index 61e8c137..112b67f4 100644 --- a/axum-extra/src/routing/mod.rs +++ b/axum-extra/src/routing/mod.rs @@ -36,7 +36,7 @@ pub trait RouterExt: sealed::Sealed { impl RouterExt for Router where - B: Send + 'static, + B: axum::body::HttpBody + Send + 'static, { fn with(self, routes: T) -> Self where diff --git a/axum-extra/src/routing/resource.rs b/axum-extra/src/routing/resource.rs index 6cfa41ad..830dc404 100644 --- a/axum-extra/src/routing/resource.rs +++ b/axum-extra/src/routing/resource.rs @@ -54,7 +54,10 @@ pub struct Resource { pub(crate) router: Router, } -impl Resource { +impl Resource +where + B: axum::body::HttpBody + Send + 'static, +{ /// Create a `Resource` with the given name. /// /// All routes will be nested at `/{resource_name}`. diff --git a/axum/CHANGELOG.md b/axum/CHANGELOG.md index 9c6f307a..5aee01d8 100644 --- a/axum/CHANGELOG.md +++ b/axum/CHANGELOG.md @@ -51,13 +51,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `PathRejection` - `MatchedPathRejection` - `WebSocketUpgradeRejection` -- **fixed:** Set `Allow` header when responding with `405 Method Not Allowed` +- **fixed:** Set `Allow` header when responding with `405 Method Not Allowed` ([#733]) +- **fixed:** Correctly set the `Content-Length` header for response to `HEAD` + requests ([#734]) [#644]: https://github.com/tokio-rs/axum/pull/644 [#665]: https://github.com/tokio-rs/axum/pull/665 [#698]: https://github.com/tokio-rs/axum/pull/698 [#699]: https://github.com/tokio-rs/axum/pull/699 [#719]: https://github.com/tokio-rs/axum/pull/719 +[#733]: https://github.com/tokio-rs/axum/pull/733 +[#734]: https://github.com/tokio-rs/axum/pull/734 # 0.4.4 (13. January, 2022) diff --git a/axum/src/routing/future.rs b/axum/src/routing/future.rs index c195e12e..013ed894 100644 --- a/axum/src/routing/future.rs +++ b/axum/src/routing/future.rs @@ -1,26 +1,21 @@ //! Future types. use crate::response::Response; -use futures_util::future::Either; -use std::{convert::Infallible, future::ready}; +use std::convert::Infallible; pub use super::{into_make_service::IntoMakeServiceFuture, route::RouteFuture}; opaque_future! { /// Response future for [`Router`](super::Router). - pub type RouterFuture = - futures_util::future::Either< - RouteFuture, - std::future::Ready>, - >; + pub type RouterFuture = RouteFuture; } impl RouterFuture { pub(super) fn from_future(future: RouteFuture) -> Self { - Self::new(Either::Left(future)) + Self::new(future) } pub(super) fn from_response(response: Response) -> Self { - Self::new(Either::Right(ready(Ok(response)))) + Self::new(RouteFuture::from_response(response)) } } diff --git a/axum/src/routing/method_routing.rs b/axum/src/routing/method_routing.rs index 8daa1d68..79f6b04f 100644 --- a/axum/src/routing/method_routing.rs +++ b/axum/src/routing/method_routing.rs @@ -960,7 +960,10 @@ where use crate::routing::future::RouteFuture; -impl Service> for MethodRouter { +impl Service> for MethodRouter +where + B: HttpBody, +{ type Response = Response; type Error = E; type Future = RouteFuture; @@ -980,7 +983,7 @@ impl Service> for MethodRouter { ) => { if $method == Method::$method_variant { if let Some(svc) = $svc { - return RouteFuture::new(svc.0.clone().oneshot($req)) + return RouteFuture::from_future(svc.0.clone().oneshot($req)) .strip_body($method == Method::HEAD); } } @@ -1014,13 +1017,14 @@ impl Service> for MethodRouter { call!(req, method, DELETE, delete); call!(req, method, TRACE, trace); - let future = - match fallback { - Fallback::Default(fallback) => RouteFuture::new(fallback.0.clone().oneshot(req)) - .strip_body(method == Method::HEAD), - Fallback::Custom(fallback) => RouteFuture::new(fallback.0.clone().oneshot(req)) - .strip_body(method == Method::HEAD), - }; + let future = match fallback { + Fallback::Default(fallback) => { + RouteFuture::from_future(fallback.0.clone().oneshot(req)) + .strip_body(method == Method::HEAD) + } + Fallback::Custom(fallback) => RouteFuture::from_future(fallback.0.clone().oneshot(req)) + .strip_body(method == Method::HEAD), + }; match allow_header { AllowHeader::None => future.allow_header(Bytes::new()), diff --git a/axum/src/routing/mod.rs b/axum/src/routing/mod.rs index bf55d4f2..931bcdbb 100644 --- a/axum/src/routing/mod.rs +++ b/axum/src/routing/mod.rs @@ -81,7 +81,7 @@ impl Clone for Router { impl Default for Router where - B: Send + 'static, + B: HttpBody + Send + 'static, { fn default() -> Self { Self::new() @@ -93,7 +93,7 @@ const NEST_TAIL_PARAM_CAPTURE: &str = "/*axum_nest"; impl Router where - B: Send + 'static, + B: HttpBody + Send + 'static, { /// Create a new `Router`. /// @@ -461,7 +461,7 @@ where impl Service> for Router where - B: Send + 'static, + B: HttpBody + Send + 'static, { type Response = Response; type Error = Infallible; diff --git a/axum/src/routing/route.rs b/axum/src/routing/route.rs index f3b4a168..7fea7f44 100644 --- a/axum/src/routing/route.rs +++ b/axum/src/routing/route.rs @@ -1,9 +1,12 @@ use crate::{ - body::{boxed, Body, Empty}, + body::{boxed, Body, Empty, HttpBody}, response::Response, }; use bytes::Bytes; -use http::{header, HeaderValue, Request}; +use http::{ + header::{self, CONTENT_LENGTH}, + HeaderValue, Request, +}; use pin_project_lite::pin_project; use std::{ convert::Infallible, @@ -46,7 +49,10 @@ impl fmt::Debug for Route { } } -impl Service> for Route { +impl Service> for Route +where + B: HttpBody, +{ type Response = Response; type Error = E; type Future = RouteFuture; @@ -58,7 +64,7 @@ impl Service> for Route { #[inline] fn call(&mut self, req: Request) -> Self::Future { - RouteFuture::new(self.0.clone().oneshot(req)) + RouteFuture::from_future(self.0.clone().oneshot(req)) } } @@ -66,21 +72,44 @@ pin_project! { /// Response future for [`Route`]. pub struct RouteFuture { #[pin] - future: Oneshot< - BoxCloneService, Response, E>, - Request, - >, + kind: RouteFutureKind, strip_body: bool, allow_header: Option, } } +pin_project! { + #[project = RouteFutureKindProj] + enum RouteFutureKind { + Future { + #[pin] + future: Oneshot< + BoxCloneService, Response, E>, + Request, + >, + }, + Response { + response: Option, + } + } +} + impl RouteFuture { - pub(crate) fn new( + pub(crate) fn from_future( future: Oneshot, Response, E>, Request>, ) -> Self { - RouteFuture { - future, + Self { + kind: RouteFutureKind::Future { future }, + strip_body: false, + allow_header: None, + } + } + + pub(super) fn from_response(response: Response) -> Self { + Self { + kind: RouteFutureKind::Response { + response: Some(response), + }, strip_body: false, allow_header: None, } @@ -97,38 +126,74 @@ impl RouteFuture { } } -impl Future for RouteFuture { +impl Future for RouteFuture +where + B: HttpBody, +{ type Output = Result; #[inline] - fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - let strip_body = self.strip_body; - let allow_header = self.allow_header.take(); + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let this = self.project(); - match self.project().future.poll(cx) { - Poll::Ready(Ok(res)) => { - let mut res = if strip_body { - res.map(|_| boxed(Empty::new())) - } else { - res - }; - - match allow_header { - Some(allow_header) if !res.headers().contains_key(header::ALLOW) => { - res.headers_mut().insert( - header::ALLOW, - HeaderValue::from_maybe_shared(allow_header) - .expect("invalid `Allow` header"), - ); - } - _ => {} - } - - Poll::Ready(Ok(res)) + let mut res = match this.kind.project() { + RouteFutureKindProj::Future { future } => match future.poll(cx) { + Poll::Ready(Ok(res)) => res, + Poll::Ready(Err(err)) => return Poll::Ready(Err(err)), + Poll::Pending => return Poll::Pending, + }, + RouteFutureKindProj::Response { response } => { + response.take().expect("future polled after completion") } - Poll::Ready(Err(err)) => Poll::Ready(Err(err)), - Poll::Pending => Poll::Pending, + }; + + set_allow_header(&mut res, this.allow_header); + + // make sure to set content-length before removing the body + set_content_length(&mut res); + + let res = if *this.strip_body { + res.map(|_| boxed(Empty::new())) + } else { + res + }; + + Poll::Ready(Ok(res)) + } +} + +fn set_allow_header(res: &mut Response, allow_header: &mut Option) { + match allow_header.take() { + Some(allow_header) if !res.headers().contains_key(header::ALLOW) => { + res.headers_mut().insert( + header::ALLOW, + HeaderValue::from_maybe_shared(allow_header).expect("invalid `Allow` header"), + ); } + _ => {} + } +} + +fn set_content_length(res: &mut Response) +where + B: HttpBody, +{ + if res.headers().contains_key(CONTENT_LENGTH) { + return; + } + + if let Some(size) = res.size_hint().exact() { + let header_value = if size == 0 { + #[allow(clippy::declare_interior_mutable_const)] + const ZERO: HeaderValue = HeaderValue::from_static("0"); + + ZERO + } else { + let mut buffer = itoa::Buffer::new(); + HeaderValue::from_str(buffer.format(size)).unwrap() + }; + + res.headers_mut().insert(CONTENT_LENGTH, header_value); } } diff --git a/axum/src/routing/tests/mod.rs b/axum/src/routing/tests/mod.rs index 418a0943..a30c52fc 100644 --- a/axum/src/routing/tests/mod.rs +++ b/axum/src/routing/tests/mod.rs @@ -618,3 +618,30 @@ async fn merging_routers_with_same_paths_but_different_methods() { let body = res.text().await; assert_eq!(body, "POST"); } + +#[tokio::test] +async fn head_content_length_through_hyper_server() { + let app = Router::new() + .route("/", get(|| async { "foo" })) + .route("/json", get(|| async { Json(json!({ "foo": 1 })) })); + + let client = TestClient::new(app); + + let res = client.head("/").send().await; + assert_eq!(res.headers()["content-length"], "3"); + assert!(res.text().await.is_empty()); + + let res = client.head("/json").send().await; + assert_eq!(res.headers()["content-length"], "9"); + assert!(res.text().await.is_empty()); +} + +#[tokio::test] +async fn head_content_length_through_hyper_server_that_hits_fallback() { + let app = Router::new().fallback((|| async { "foo" }).into_service()); + + let client = TestClient::new(app); + + let res = client.head("/").send().await; + assert_eq!(res.headers()["content-length"], "3"); +} diff --git a/axum/src/test_helpers.rs b/axum/src/test_helpers.rs index 7ffdf1a8..43ae5c1e 100644 --- a/axum/src/test_helpers.rs +++ b/axum/src/test_helpers.rs @@ -59,6 +59,12 @@ impl TestClient { } } + pub(crate) fn head(&self, url: &str) -> RequestBuilder { + RequestBuilder { + builder: self.client.head(format!("http://{}{}", self.addr, url)), + } + } + pub(crate) fn post(&self, url: &str) -> RequestBuilder { RequestBuilder { builder: self.client.post(format!("http://{}{}", self.addr, url)),