diff --git a/axum/src/extract/connect_info.rs b/axum/src/extract/connect_info.rs index 071b0690..4c093bff 100644 --- a/axum/src/extract/connect_info.rs +++ b/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 { - svc: S, +pub struct WithConnectInfo { + make_svc: M, _connect_info: PhantomData C>, } #[test] fn traits() { use crate::test_helpers::*; - assert_send::>(); + assert_send::>(); } -impl IntoMakeServiceWithConnectInfo { - pub(crate) fn new(svc: S) -> Self { +impl WithConnectInfo { + pub(crate) fn new(make_svc: M) -> Self { Self { - svc, + make_svc, _connect_info: PhantomData, } } } -impl fmt::Debug for IntoMakeServiceWithConnectInfo +impl fmt::Debug for WithConnectInfo 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 Clone for IntoMakeServiceWithConnectInfo +impl Clone for WithConnectInfo 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: 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 Service for IntoMakeServiceWithConnectInfo +impl Service for WithConnectInfo where - S: Clone, + M: Service, C: Connected, { - type Response = AddExtension>; - type Error = Infallible; - type Future = ResponseFuture; + type Response = AddExtension>; + type Error = M::Error; + type Future = ResponseFuture; #[inline] - fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll> { - Poll::Ready(Ok(())) + fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { + self.make_svc.poll_ready(cx) } fn call(&mut self, target: T) -> Self::Future { - let connect_info = ConnectInfo(C::connect_info(target)); - let svc = AddExtensionLayer::new(connect_info).layer(self.svc.clone()); - ResponseFuture::new(ready(Ok(svc))) + let connect_info = C::connect_info(&target); + ResponseFuture { + future: self.make_svc.call(target), + connect_info: Some(connect_info), + } } } -opaque_future! { +pin_project! { /// Response future for [`IntoMakeServiceWithConnectInfo`]. - pub type ResponseFuture = - std::future::Ready>, Infallible>>; + pub struct ResponseFuture { + #[pin] + future: F, + connect_info: Option, + } +} + +impl Future for ResponseFuture +where + F: Future>, + C: Clone, +{ + type Output = Result>, E>; + + fn poll(self: std::pin::Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + 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::()); + .serve(app.into_make_service().with_connect_info::()); 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) - .unwrap() - .serve(app.into_make_service_with_connect_info::()); + let server = Server::from_tcp(listener).unwrap().serve( + app.into_make_service() + .with_connect_info::(), + ); tx.send(()).unwrap(); server.await.expect("server error"); }); diff --git a/axum/src/handler/mod.rs b/axum/src/handler/mod.rs index 57a09444..3731f0f0 100644 --- a/axum/src/handler/mod.rs +++ b/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: Clone + Send + Sized + 'static { /// ``` /// /// [`MakeService`]: tower::make::MakeService - fn into_make_service(self) -> IntoMakeService> { + fn into_make_service(self) -> IntoMakeService>> { IntoMakeService::new(self.into_service()) } @@ -252,11 +252,11 @@ pub trait Handler: 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( self, - ) -> IntoMakeServiceWithConnectInfo, C> + ) -> WithConnectInfo, C> where C: Connected, { - IntoMakeServiceWithConnectInfo::new(self.into_service()) + WithConnectInfo::new(self.into_service()) } } diff --git a/axum/src/routing/future.rs b/axum/src/routing/future.rs index c195e12e..62846c9c 100644 --- a/axum/src/routing/future.rs +++ b/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). diff --git a/axum/src/routing/into_make_service.rs b/axum/src/routing/into_make_service.rs index fbc57c4a..72cadd61 100644 --- a/axum/src/routing/into_make_service.rs +++ b/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 { - svc: S, +pub struct IntoMakeService { + make_svc: M, } -impl IntoMakeService { +impl IntoMakeService> { pub(crate) fn new(svc: S) -> Self { - Self { svc } + Self { + make_svc: Shared(svc), + } } } -impl Service for IntoMakeService +impl IntoMakeService { + pub fn with_connect_info(self) -> IntoMakeService> + where + C: Connected, + { + self.layer(tower_layer::layer_fn(WithConnectInfo::new)) + } + + pub fn layer(self, layer: L) -> IntoMakeService + where + L: Layer, + { + IntoMakeService { + make_svc: layer.layer(self.make_svc), + } + } +} + +impl Service for IntoMakeService +where + M: Service, +{ + type Response = M::Response; + type Error = M::Error; + type Future = M::Future; + + #[inline] + fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { + 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); + +impl Service for Shared where S: Clone, { type Response = S; type Error = Infallible; - type Future = IntoMakeServiceFuture; + type Future = SharedFuture; #[inline] fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll> { @@ -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`]. - pub type IntoMakeServiceFuture = + /// Response future for [`Shared`]. + pub type SharedFuture = std::future::Ready>; } diff --git a/axum/src/routing/mod.rs b/axum/src/routing/mod.rs index 7eac8c66..4d285917 100644 --- a/axum/src/routing/mod.rs +++ b/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 { + pub fn into_make_service(self) -> IntoMakeService> { IntoMakeService::new(self) } - #[doc = include_str!("../docs/routing/into_make_service_with_connect_info.md")] - pub fn into_make_service_with_connect_info( - self, - ) -> IntoMakeServiceWithConnectInfo - where - C: Connected, - { - IntoMakeServiceWithConnectInfo::new(self) - } - #[inline] fn call_route(&self, match_: matchit::Match<&RouteId>, mut req: Request) -> RouterFuture { let id = *match_.value; diff --git a/examples/low-level-rustls/src/main.rs b/examples/low-level-rustls/src/main.rs index 611949f0..42ef4c79 100644 --- a/examples/low-level-rustls/src/main.rs +++ b/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::(); + .into_make_service() + .with_connect_info::(); loop { let stream = poll_fn(|cx| Pin::new(&mut listener).poll_accept(cx)) diff --git a/examples/unix-domain-socket/src/main.rs b/examples/unix-domain-socket/src/main.rs index d5c1c3c2..c55c9872 100644 --- a/examples/unix-domain-socket/src/main.rs +++ b/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::()) + .serve( + app.into_make_service() + .with_connect_info::(), + ) .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();