Replace map_response(IntoResponse::into_response) with custom service wrapper

Author: svix-jplatte

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()

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<E> Route<E> {
         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.

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<str>);
@@ -37,6 +45,57 @@ pin_project! {
     }
 }
 
+#[derive(Clone)]
+pub(crate) struct MapIntoResponse<S> {
+    inner: S,
+}
+
+impl<S> MapIntoResponse<S> {
+    pub(crate) fn new(inner: S) -> Self {
+        Self { inner }
+    }
+}
+
+impl<B, S> Service<http::Request<B>> for MapIntoResponse<S>
+where
+    S: Service<http::Request<B>>,
+    S::Response: IntoResponse,
+{
+    type Response = Response;
+    type Error = S::Error;
+    type Future = MapIntoResponseFuture<S::Future>;
+
+    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.inner.poll_ready(cx)
+    }
+
+    fn call(&mut self, req: http::Request<B>) -> Self::Future {
+        MapIntoResponseFuture {
+            inner: self.inner.call(req),
+        }
+    }
+}
+
+pin_project! {
+    pub(crate) struct MapIntoResponseFuture<F> {
+        #[pin]
+        inner: F,
+    }
+}
+
+impl<F, T, E> Future for MapIntoResponseFuture<F>
+where
+    F: Future<Output = Result<T, E>>,
+    T: IntoResponse,
+{
+    type Output = Result<Response, E>;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let res = ready!(self.project().inner.poll(cx)?);
+        Poll::Ready(Ok(res.into_response()))
+    }
+}
+
 pub(crate) fn try_downcast<T, K>(k: K) -> Result<T, K>
 where
     T: 'static,