diff --git a/axum/src/middleware/from_fn.rs b/axum/src/middleware/from_fn.rs index daa76ec6..c4866566 100644 --- a/axum/src/middleware/from_fn.rs +++ b/axum/src/middleware/from_fn.rs @@ -1,4 +1,3 @@ -use crate::response::{IntoResponse, Response}; use axum_core::extract::{FromRequest, FromRequestParts, Request}; use futures_util::future::BoxFuture; use std::{ @@ -11,10 +10,14 @@ use std::{ task::{Context, Poll}, }; use tower::util::BoxCloneSyncService; -use tower::ServiceBuilder; use tower_layer::Layer; use tower_service::Service; +use crate::{ + response::{IntoResponse, Response}, + util::MapIntoResponse, +}; + /// Create a middleware from an async function. /// /// `from_fn` requires the function given to @@ -300,10 +303,7 @@ macro_rules! impl_service { Err(rejection) => return rejection.into_response(), }; - let inner = ServiceBuilder::new() - .layer_fn(BoxCloneSyncService::new) - .map_response(IntoResponse::into_response) - .service(ready_inner); + let inner = BoxCloneSyncService::new(MapIntoResponse::new(ready_inner)); let next = Next { inner }; f($($ty,)* $last, next).await.into_response() diff --git a/axum/src/routing/route.rs b/axum/src/routing/route.rs index fe27d2e2..6cdc58a6 100644 --- a/axum/src/routing/route.rs +++ b/axum/src/routing/route.rs @@ -1,6 +1,7 @@ use crate::{ body::{Body, HttpBody}, response::Response, + util::MapIntoResponse, }; use axum_core::{extract::Request, response::IntoResponse}; use bytes::Bytes; @@ -17,7 +18,7 @@ use std::{ task::{ready, Context, Poll}, }; use tower::{ - util::{BoxCloneSyncService, MapErrLayer, MapResponseLayer, Oneshot}, + util::{BoxCloneSyncService, MapErrLayer, Oneshot}, ServiceExt, }; use tower_layer::Layer; @@ -36,9 +37,7 @@ impl Route { T::Response: IntoResponse + 'static, T::Future: Send + 'static, { - Self(BoxCloneSyncService::new( - svc.map_response(IntoResponse::into_response), - )) + Self(BoxCloneSyncService::new(MapIntoResponse::new(svc))) } /// Variant of [`Route::call`] that takes ownership of the route to avoid cloning. diff --git a/axum/src/util.rs b/axum/src/util.rs index 7c9b7864..e4014c59 100644 --- a/axum/src/util.rs +++ b/axum/src/util.rs @@ -1,5 +1,13 @@ +use axum_core::response::{IntoResponse, Response}; use pin_project_lite::pin_project; -use std::{ops::Deref, sync::Arc}; +use std::{ + future::Future, + ops::Deref, + pin::Pin, + sync::Arc, + task::{ready, Context, Poll}, +}; +use tower::Service; #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub(crate) struct PercentDecodedStr(Arc); @@ -37,6 +45,57 @@ pin_project! { } } +#[derive(Clone)] +pub(crate) struct MapIntoResponse { + inner: S, +} + +impl MapIntoResponse { + pub(crate) fn new(inner: S) -> Self { + Self { inner } + } +} + +impl Service> for MapIntoResponse +where + S: Service>, + S::Response: IntoResponse, +{ + type Response = Response; + type Error = S::Error; + type Future = MapIntoResponseFuture; + + fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { + self.inner.poll_ready(cx) + } + + fn call(&mut self, req: http::Request) -> Self::Future { + MapIntoResponseFuture { + inner: self.inner.call(req), + } + } +} + +pin_project! { + pub(crate) struct MapIntoResponseFuture { + #[pin] + inner: F, + } +} + +impl Future for MapIntoResponseFuture +where + F: Future>, + T: IntoResponse, +{ + type Output = Result; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let res = ready!(self.project().inner.poll(cx)?); + Poll::Ready(Ok(res.into_response())) + } +} + pub(crate) fn try_downcast(k: K) -> Result where T: 'static,