checkpoint

axum/src/extract/connect_info.rs

@@ -8,10 +8,11 @@ use super::{Extension, FromRequest, RequestParts};
 use crate::{AddExtension, AddExtensionLayer};
 use async_trait::async_trait;
 use hyper::server::conn::AddrStream;
+use pin_project_lite::pin_project;
 use std::{
     convert::Infallible,
     fmt,
-    future::ready,
+    future::{ready, Future},
     marker::PhantomData,
     net::SocketAddr,
     task::{Context, Poll},
@@ -25,44 +26,44 @@ use tower_service::Service;
 ///
 /// [`MakeService`]: tower::make::MakeService
 /// [`Router::into_make_service_with_connect_info`]: crate::routing::Router::into_make_service_with_connect_info
-pub struct IntoMakeServiceWithConnectInfo<S, C> {
+pub struct WithConnectInfo<M, C> {
-    svc: S,
+    make_svc: M,
     _connect_info: PhantomData<fn() -> C>,
 }
 
 #[test]
 fn traits() {
     use crate::test_helpers::*;
-    assert_send::<IntoMakeServiceWithConnectInfo<(), NotSendSync>>();
+    assert_send::<WithConnectInfo<(), NotSendSync>>();
 }
 
-impl<S, C> IntoMakeServiceWithConnectInfo<S, C> {
+impl<M, C> WithConnectInfo<M, C> {
-    pub(crate) fn new(svc: S) -> Self {
+    pub(crate) fn new(make_svc: M) -> Self {
         Self {
-            svc,
+            make_svc,
             _connect_info: PhantomData,
         }
     }
 }
 
-impl<S, C> fmt::Debug for IntoMakeServiceWithConnectInfo<S, C>
+impl<M, C> fmt::Debug for WithConnectInfo<M, C>
 where
-    S: fmt::Debug,
+    M: fmt::Debug,
 {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("IntoMakeServiceWithConnectInfo")
-            .field("svc", &self.svc)
+            .field("svc", &self.make_svc)
             .finish()
     }
 }
 
-impl<S, C> Clone for IntoMakeServiceWithConnectInfo<S, C>
+impl<M, C> Clone for WithConnectInfo<M, C>
 where
-    S: Clone,
+    M: Clone,
 {
     fn clone(&self) -> Self {
         Self {
-            svc: self.svc.clone(),
+            make_svc: self.make_svc.clone(),
             _connect_info: PhantomData,
         }
     }
@@ -79,40 +80,64 @@ where
 /// [`Router::into_make_service_with_connect_info`]: crate::routing::Router::into_make_service_with_connect_info
 pub trait Connected<T>: Clone + Send + Sync + 'static {
     /// Create type holding information about the connection.
-    fn connect_info(target: T) -> Self;
+    fn connect_info(target: &T) -> Self;
 }
 
 impl Connected<&AddrStream> for SocketAddr {
-    fn connect_info(target: &AddrStream) -> Self {
+    fn connect_info(target: &&AddrStream) -> Self {
         target.remote_addr()
     }
 }
 
-impl<S, C, T> Service<T> for IntoMakeServiceWithConnectInfo<S, C>
+impl<M, C, T> Service<T> for WithConnectInfo<M, C>
 where
-    S: Clone,
+    M: Service<T>,
     C: Connected<T>,
 {
-    type Response = AddExtension<S, ConnectInfo<C>>;
+    type Response = AddExtension<M::Response, ConnectInfo<C>>;
-    type Error = Infallible;
+    type Error = M::Error;
-    type Future = ResponseFuture<S, C>;
+    type Future = ResponseFuture<M::Future, C>;
 
     #[inline]
-    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
-        Poll::Ready(Ok(()))
+        self.make_svc.poll_ready(cx)
     }
 
     fn call(&mut self, target: T) -> Self::Future {
-        let connect_info = ConnectInfo(C::connect_info(target));
+        let connect_info = C::connect_info(&target);
-        let svc = AddExtensionLayer::new(connect_info).layer(self.svc.clone());
+        ResponseFuture {
-        ResponseFuture::new(ready(Ok(svc)))
+            future: self.make_svc.call(target),
+            connect_info: Some(connect_info),
+        }
     }
 }
 
-opaque_future! {
+pin_project! {
     /// Response future for [`IntoMakeServiceWithConnectInfo`].
-    pub type ResponseFuture<S, C> =
+    pub struct ResponseFuture<F, C> {
-        std::future::Ready<Result<AddExtension<S, ConnectInfo<C>>, Infallible>>;
+        #[pin]
+        future: F,
+        connect_info: Option<C>,
+    }
+}
+
+impl<F, C, S, E> Future for ResponseFuture<F, C>
+where
+    F: Future<Output = Result<S, E>>,
+    C: Clone,
+{
+    type Output = Result<AddExtension<S, ConnectInfo<C>>, E>;
+
+    fn poll(self: std::pin::Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let this = self.project();
+        let svc = futures_util::ready!(this.future.poll(cx))?;
+        let connect_info = this
+            .connect_info
+            .take()
+            .expect("future polled after completion");
+        let svc = AddExtensionLayer::new(ConnectInfo(connect_info)).layer(svc);
+        Poll::Ready(Ok(svc))
+    }
 }
 
 /// Extractor for getting connection information produced by a [`Connected`].
@@ -161,7 +186,7 @@ mod tests {
             let app = Router::new().route("/", get(handler));
             let server = Server::from_tcp(listener)
                 .unwrap()
-                .serve(app.into_make_service_with_connect_info::<SocketAddr, _>());
+                .serve(app.into_make_service().with_connect_info::<SocketAddr, _>());
             tx.send(()).unwrap();
             server.await.expect("server error");
         });
@@ -182,7 +207,7 @@ mod tests {
         }
 
         impl Connected<&AddrStream> for MyConnectInfo {
-            fn connect_info(_target: &AddrStream) -> Self {
+            fn connect_info(_target: &&AddrStream) -> Self {
                 Self {
                     value: "it worked!",
                 }
@@ -199,9 +224,10 @@ mod tests {
         let (tx, rx) = tokio::sync::oneshot::channel();
         tokio::spawn(async move {
             let app = Router::new().route("/", get(handler));
-            let server = Server::from_tcp(listener)
+            let server = Server::from_tcp(listener).unwrap().serve(
-                .unwrap()
+                app.into_make_service()
-                .serve(app.into_make_service_with_connect_info::<MyConnectInfo, _>());
+                    .with_connect_info::<MyConnectInfo, _>(),
+            );
             tx.send(()).unwrap();
             server.await.expect("server error");
         });

axum/src/handler/mod.rs

@@ -73,7 +73,7 @@
 use crate::{
     body::{boxed, Body, Bytes, HttpBody},
     extract::{
-        connect_info::{Connected, IntoMakeServiceWithConnectInfo},
+        connect_info::{Connected, WithConnectInfo},
         FromRequest, RequestParts,
     },
     response::{IntoResponse, Response},
@@ -218,7 +218,7 @@ pub trait Handler<T, B = Body>: Clone + Send + Sized + 'static {
     /// ```
     ///
     /// [`MakeService`]: tower::make::MakeService
-    fn into_make_service(self) -> IntoMakeService<IntoService<Self, T, B>> {
+    fn into_make_service(self) -> IntoMakeService<crate::routing::Shared<IntoService<Self, T, B>>> {
         IntoMakeService::new(self.into_service())
     }
 
@@ -252,11 +252,11 @@ pub trait Handler<T, B = Body>: Clone + Send + Sized + 'static {
     /// [`Router::into_make_service_with_connect_info`]: crate::routing::Router::into_make_service_with_connect_info
     fn into_make_service_with_connect_info<C, Target>(
         self,
-    ) -> IntoMakeServiceWithConnectInfo<IntoService<Self, T, B>, C>
+    ) -> WithConnectInfo<IntoService<Self, T, B>, C>
     where
         C: Connected<Target>,
     {
-        IntoMakeServiceWithConnectInfo::new(self.into_service())
+        WithConnectInfo::new(self.into_service())
     }
 }
 

axum/src/routing/future.rs

@@ -4,7 +4,7 @@ use crate::response::Response;
 use futures_util::future::Either;
 use std::{convert::Infallible, future::ready};
 
-pub use super::{into_make_service::IntoMakeServiceFuture, route::RouteFuture};
+pub use super::{into_make_service::SharedFuture, route::RouteFuture};
 
 opaque_future! {
     /// Response future for [`Router`](super::Router).

axum/src/routing/into_make_service.rs

@@ -1,31 +1,74 @@
+use crate::extract::connect_info::{Connected, WithConnectInfo};
 use std::{
     convert::Infallible,
     future::ready,
     task::{Context, Poll},
 };
+use tower_layer::Layer;
 use tower_service::Service;
 
 /// A [`MakeService`] that produces axum router services.
 ///
 /// [`MakeService`]: tower::make::MakeService
 #[derive(Debug, Clone)]
-pub struct IntoMakeService<S> {
+pub struct IntoMakeService<M> {
-    svc: S,
+    make_svc: M,
 }
 
-impl<S> IntoMakeService<S> {
+impl<S> IntoMakeService<Shared<S>> {
     pub(crate) fn new(svc: S) -> Self {
-        Self { svc }
+        Self {
+            make_svc: Shared(svc),
+        }
     }
 }
 
-impl<S, T> Service<T> for IntoMakeService<S>
+impl<M> IntoMakeService<M> {
+    pub fn with_connect_info<C, Target>(self) -> IntoMakeService<WithConnectInfo<M, C>>
+    where
+        C: Connected<Target>,
+    {
+        self.layer(tower_layer::layer_fn(WithConnectInfo::new))
+    }
+
+    pub fn layer<L>(self, layer: L) -> IntoMakeService<L::Service>
+    where
+        L: Layer<M>,
+    {
+        IntoMakeService {
+            make_svc: layer.layer(self.make_svc),
+        }
+    }
+}
+
+impl<M, T> Service<T> for IntoMakeService<M>
+where
+    M: Service<T>,
+{
+    type Response = M::Response;
+    type Error = M::Error;
+    type Future = M::Future;
+
+    #[inline]
+    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.make_svc.poll_ready(cx)
+    }
+
+    fn call(&mut self, target: T) -> Self::Future {
+        self.make_svc.call(target)
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct Shared<S>(S);
+
+impl<S, T> Service<T> for Shared<S>
 where
     S: Clone,
 {
     type Response = S;
     type Error = Infallible;
-    type Future = IntoMakeServiceFuture<S>;
+    type Future = SharedFuture<S>;
 
     #[inline]
     fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
@@ -33,13 +76,13 @@ where
     }
 
     fn call(&mut self, _target: T) -> Self::Future {
-        IntoMakeServiceFuture::new(ready(Ok(self.svc.clone())))
+        SharedFuture::new(ready(Ok(self.0.clone())))
     }
 }
 
 opaque_future! {
-    /// Response future for [`IntoMakeService`].
+    /// Response future for [`Shared`].
-    pub type IntoMakeServiceFuture<S> =
+    pub type SharedFuture<S> =
         std::future::Ready<Result<S, Infallible>>;
 }
 

axum/src/routing/mod.rs

@@ -4,7 +4,7 @@ use self::{future::RouterFuture, not_found::NotFound};
 use crate::{
     body::{boxed, Body, Bytes, HttpBody},
     extract::{
-        connect_info::{Connected, IntoMakeServiceWithConnectInfo},
+        connect_info::{Connected, WithConnectInfo},
         MatchedPath, OriginalUri,
     },
     response::Response,
@@ -38,7 +38,11 @@ mod strip_prefix;
 #[cfg(test)]
 mod tests;
 
-pub use self::{into_make_service::IntoMakeService, method_filter::MethodFilter, route::Route};
+pub use self::{
+    into_make_service::{IntoMakeService, Shared},
+    method_filter::MethodFilter,
+    route::Route,
+};
 
 pub use self::method_routing::{
     any, any_service, delete, delete_service, get, get_service, head, head_service, on, on_service,
@@ -381,20 +385,10 @@ where
     /// ```
     ///
     /// [`MakeService`]: tower::make::MakeService
-    pub fn into_make_service(self) -> IntoMakeService<Self> {
+    pub fn into_make_service(self) -> IntoMakeService<Shared<Self>> {
         IntoMakeService::new(self)
     }
 
-    #[doc = include_str!("../docs/routing/into_make_service_with_connect_info.md")]
-    pub fn into_make_service_with_connect_info<C, Target>(
-        self,
-    ) -> IntoMakeServiceWithConnectInfo<Self, C>
-    where
-        C: Connected<Target>,
-    {
-        IntoMakeServiceWithConnectInfo::new(self)
-    }
-
     #[inline]
     fn call_route(&self, match_: matchit::Match<&RouteId>, mut req: Request<B>) -> RouterFuture<B> {
         let id = *match_.value;

examples/low-level-rustls/src/main.rs

@@ -39,7 +39,8 @@ async fn main() {
 
     let mut app = Router::new()
         .route("/", get(handler))
-        .into_make_service_with_connect_info::<SocketAddr, _>();
+        .into_make_service()
+        .with_connect_info::<SocketAddr, _>();
 
     loop {
         let stream = poll_fn(|cx| Pin::new(&mut listener).poll_accept(cx))

examples/unix-domain-socket/src/main.rs

@@ -55,7 +55,10 @@ async fn main() {
         let app = Router::new().route("/", get(handler));
 
         axum::Server::builder(ServerAccept { uds })
-            .serve(app.into_make_service_with_connect_info::<UdsConnectInfo, _>())
+            .serve(
+                app.into_make_service()
+                    .with_connect_info::<UdsConnectInfo, _>(),
+            )
             .await
             .unwrap();
     });
@@ -156,7 +159,7 @@ struct UdsConnectInfo {
 }
 
 impl connect_info::Connected<&UnixStream> for UdsConnectInfo {
-    fn connect_info(target: &UnixStream) -> Self {
+    fn connect_info(target: &&UnixStream) -> Self {
         let peer_addr = target.peer_addr().unwrap();
         let peer_cred = target.peer_cred().unwrap();