diff --git a/CHANGELOG.md b/CHANGELOG.md index 6ba45462..2a4e63a7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Implement `Sink` for `WebSocket` ([#52](https://github.com/tokio-rs/axum/pull/52)) - Implement `Deref` most extractors ([#56](https://github.com/tokio-rs/axum/pull/56)) - Return `405 Method Not Allowed` for unsupported method for route ([#63](https://github.com/tokio-rs/axum/pull/63)) +- Add extractor for remote connection info ([#55](https://github.com/tokio-rs/axum/pull/55)) ## Breaking changes diff --git a/Cargo.toml b/Cargo.toml index 89f38452..d75bab3f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -23,7 +23,7 @@ bytes = "1.0" futures-util = "0.3" http = "0.2" http-body = "0.4" -hyper = { version = "0.14", features = ["server", "tcp"] } +hyper = { version = "0.14", features = ["server", "tcp", "http1"] } pin-project = "1.0" regex = "1.5" serde = "1.0" diff --git a/examples/unix_domain_socket.rs b/examples/unix_domain_socket.rs index dc5259a8..9077d011 100644 --- a/examples/unix_domain_socket.rs +++ b/examples/unix_domain_socket.rs @@ -1,4 +1,7 @@ -use axum::prelude::*; +use axum::{ + extract::connect_info::{self, ConnectInfo}, + prelude::*, +}; use futures::ready; use http::{Method, StatusCode, Uri}; use hyper::{ @@ -9,9 +12,10 @@ use std::{ io, path::PathBuf, pin::Pin, + sync::Arc, task::{Context, Poll}, }; -use tokio::net::UnixListener; +use tokio::net::{unix::UCred, UnixListener}; use tokio::{ io::{AsyncRead, AsyncWrite}, net::UnixStream, @@ -35,10 +39,10 @@ async fn main() { let uds = UnixListener::bind(path.clone()).unwrap(); tokio::spawn(async { - let app = route("/", get(|| async { "Hello, World!" })); + let app = route("/", get(handler)); hyper::Server::builder(ServerAccept { uds }) - .serve(app.into_make_service()) + .serve(app.into_make_service_with_connect_info::()) .await .unwrap(); }); @@ -67,6 +71,12 @@ async fn main() { assert_eq!(body, "Hello, World!"); } +async fn handler(ConnectInfo(info): ConnectInfo) -> &'static str { + println!("new connection from `{:?}`", info); + + "Hello, World!" +} + struct ServerAccept { uds: UnixListener, } @@ -124,3 +134,23 @@ impl Connection for ClientConnection { Connected::new() } } + +#[derive(Clone, Debug)] +struct UdsConnectInfo { + peer_addr: Arc, + peer_cred: UCred, +} + +impl connect_info::Connected<&UnixStream> for UdsConnectInfo { + type ConnectInfo = Self; + + fn connect_info(target: &UnixStream) -> Self::ConnectInfo { + let peer_addr = target.peer_addr().unwrap(); + let peer_cred = target.peer_cred().unwrap(); + + Self { + peer_addr: Arc::new(peer_addr), + peer_cred, + } + } +} diff --git a/src/extract/connect_info.rs b/src/extract/connect_info.rs new file mode 100644 index 00000000..cb92e095 --- /dev/null +++ b/src/extract/connect_info.rs @@ -0,0 +1,203 @@ +//! Extractor for getting connection information from a client. +//! +//! See [`RoutingDsl::into_make_service_with_connect_info`] for more details. +//! +//! [`RoutingDsl::into_make_service_with_connect_info`]: crate::routing::RoutingDsl::into_make_service_with_connect_info + +use super::{Extension, FromRequest, RequestParts}; +use async_trait::async_trait; +use hyper::server::conn::AddrStream; +use std::{ + convert::Infallible, + fmt, + marker::PhantomData, + net::SocketAddr, + task::{Context, Poll}, +}; +use tower::Service; +use tower_http::add_extension::AddExtension; + +/// A [`MakeService`] created from a router. +/// +/// See [`RoutingDsl::into_make_service_with_connect_info`] for more details. +/// +/// [`MakeService`]: tower::make::MakeService +/// [`RoutingDsl::into_make_service_with_connect_info`]: crate::routing::RoutingDsl::into_make_service_with_connect_info +pub struct IntoMakeServiceWithConnectInfo { + svc: S, + _connect_info: PhantomData C>, +} + +impl IntoMakeServiceWithConnectInfo { + pub(crate) fn new(svc: S) -> Self { + Self { + svc, + _connect_info: PhantomData, + } + } +} + +impl fmt::Debug for IntoMakeServiceWithConnectInfo +where + S: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("IntoMakeServiceWithConnectInfo") + .field("svc", &self.svc) + .finish() + } +} + +/// Trait that connected IO resources implement and use to produce information +/// about the connection. +/// +/// The goal for this trait is to allow users to implement custom IO types that +/// can still provide the same connection metadata. +/// +/// See [`RoutingDsl::into_make_service_with_connect_info`] for more details. +/// +/// [`RoutingDsl::into_make_service_with_connect_info`]: crate::routing::RoutingDsl::into_make_service_with_connect_info +pub trait Connected { + /// The connection information type the IO resources generates. + type ConnectInfo: Clone + Send + Sync + 'static; + + /// Create type holding information about the connection. + fn connect_info(target: T) -> Self::ConnectInfo; +} + +impl Connected<&AddrStream> for SocketAddr { + type ConnectInfo = SocketAddr; + + fn connect_info(target: &AddrStream) -> Self::ConnectInfo { + target.remote_addr() + } +} + +impl Service for IntoMakeServiceWithConnectInfo +where + S: Clone, + C: Connected, +{ + type Response = AddExtension>; + type Error = Infallible; + type Future = ResponseFuture; + + fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn call(&mut self, target: T) -> Self::Future { + let connect_info = ConnectInfo(C::connect_info(target)); + let svc = AddExtension::new(self.svc.clone(), connect_info); + ResponseFuture(futures_util::future::ok(svc)) + } +} + +opaque_future! { + /// Response future for [`IntoMakeServiceWithConnectInfo`]. + pub type ResponseFuture = + futures_util::future::Ready>; +} + +/// Extractor for getting connection information produced by a [`Connected`]. +/// +/// Note this extractor requires you to use +/// [`RoutingDsl::into_make_service_with_connect_info`] to run your app +/// otherwise it will fail at runtime. +/// +/// See [`RoutingDsl::into_make_service_with_connect_info`] for more details. +/// +/// [`RoutingDsl::into_make_service_with_connect_info`]: crate::routing::RoutingDsl::into_make_service_with_connect_info +#[derive(Clone, Copy, Debug)] +pub struct ConnectInfo(pub T); + +#[async_trait] +impl FromRequest for ConnectInfo +where + B: Send, + T: Clone + Send + Sync + 'static, +{ + type Rejection = as FromRequest>::Rejection; + + async fn from_request(req: &mut RequestParts) -> Result { + let Extension(connect_info) = Extension::::from_request(req).await?; + Ok(connect_info) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::prelude::*; + use hyper::Server; + use std::net::{SocketAddr, TcpListener}; + + #[tokio::test] + async fn socket_addr() { + async fn handler(ConnectInfo(addr): ConnectInfo) -> String { + format!("{}", addr) + } + + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = listener.local_addr().unwrap(); + + let (tx, rx) = tokio::sync::oneshot::channel(); + tokio::spawn(async move { + let app = route("/", get(handler)); + let server = Server::from_tcp(listener) + .unwrap() + .serve(app.into_make_service_with_connect_info::()); + tx.send(()).unwrap(); + server.await.expect("server error"); + }); + rx.await.unwrap(); + + let client = reqwest::Client::new(); + + let res = client.get(format!("http://{}", addr)).send().await.unwrap(); + let body = res.text().await.unwrap(); + assert!(body.starts_with("127.0.0.1:")); + } + + #[tokio::test] + async fn custom() { + #[derive(Clone, Debug)] + struct MyConnectInfo { + value: &'static str, + } + + impl Connected<&AddrStream> for MyConnectInfo { + type ConnectInfo = Self; + + fn connect_info(_target: &AddrStream) -> Self::ConnectInfo { + Self { + value: "it worked!", + } + } + } + + async fn handler(ConnectInfo(addr): ConnectInfo) -> &'static str { + addr.value + } + + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = listener.local_addr().unwrap(); + + let (tx, rx) = tokio::sync::oneshot::channel(); + tokio::spawn(async move { + let app = route("/", get(handler)); + let server = Server::from_tcp(listener) + .unwrap() + .serve(app.into_make_service_with_connect_info::()); + tx.send(()).unwrap(); + server.await.expect("server error"); + }); + rx.await.unwrap(); + + let client = reqwest::Client::new(); + + let res = client.get(format!("http://{}", addr)).send().await.unwrap(); + let body = res.text().await.unwrap(); + assert_eq!(body, "it worked!"); + } +} diff --git a/src/extract/mod.rs b/src/extract/mod.rs index d6a7d804..6d59636b 100644 --- a/src/extract/mod.rs +++ b/src/extract/mod.rs @@ -260,12 +260,16 @@ use std::{ task::{Context, Poll}, }; +pub mod connect_info; pub mod extractor_middleware; pub mod rejection; #[doc(inline)] pub use self::extractor_middleware::extractor_middleware; +#[doc(inline)] +pub use self::connect_info::ConnectInfo; + #[cfg(feature = "multipart")] #[cfg_attr(docsrs, doc(cfg(feature = "multipart")))] pub mod multipart; @@ -904,6 +908,8 @@ where /// # hyper::Server::bind(&"".parse().unwrap()).serve(app.into_make_service()).await.unwrap(); /// # }; /// ``` +/// +/// [`Stream`]: https://docs.rs/futures/latest/futures/stream/trait.Stream.html #[derive(Debug)] pub struct BodyStream(B); diff --git a/src/routing.rs b/src/routing.rs index 75c20f3b..a7ed9ab2 100644 --- a/src/routing.rs +++ b/src/routing.rs @@ -3,6 +3,7 @@ use crate::{ body::{box_body, BoxBody}, buffer::MpscBuffer, + extract::connect_info::{Connected, IntoMakeServiceWithConnectInfo}, response::IntoResponse, util::ByteStr, }; @@ -266,6 +267,93 @@ pub trait RoutingDsl: crate::sealed::Sealed + Sized { { tower::make::Shared::new(self) } + + /// Convert this router into a [`MakeService`], that will store `C`'s + /// associated `ConnectInfo` in a request extension such that [`ConnectInfo`] + /// can extract it. + /// + /// This enables extracting things like the client's remote address. + /// + /// Extracting [`std::net::SocketAddr`] is supported out of the box: + /// + /// ``` + /// use axum::{prelude::*, extract::ConnectInfo}; + /// use std::net::SocketAddr; + /// + /// let app = route("/", get(handler)); + /// + /// async fn handler(ConnectInfo(addr): ConnectInfo) -> String { + /// format!("Hello {}", addr) + /// } + /// + /// # async { + /// hyper::Server::bind(&"0.0.0.0:3000".parse().unwrap()) + /// .serve( + /// app.into_make_service_with_connect_info::() + /// ) + /// .await + /// .expect("server failed"); + /// # }; + /// ``` + /// + /// You can implement custom a [`Connected`] like so: + /// + /// ``` + /// use axum::{ + /// prelude::*, + /// extract::connect_info::{ConnectInfo, Connected}, + /// }; + /// use hyper::server::conn::AddrStream; + /// + /// let app = route("/", get(handler)); + /// + /// async fn handler( + /// ConnectInfo(my_connect_info): ConnectInfo, + /// ) -> String { + /// format!("Hello {:?}", my_connect_info) + /// } + /// + /// #[derive(Clone, Debug)] + /// struct MyConnectInfo { + /// // ... + /// } + /// + /// impl Connected<&AddrStream> for MyConnectInfo { + /// type ConnectInfo = MyConnectInfo; + /// + /// fn connect_info(target: &AddrStream) -> Self::ConnectInfo { + /// MyConnectInfo { + /// // ... + /// } + /// } + /// } + /// + /// # async { + /// hyper::Server::bind(&"0.0.0.0:3000".parse().unwrap()) + /// .serve( + /// app.into_make_service_with_connect_info::() + /// ) + /// .await + /// .expect("server failed"); + /// # }; + /// ``` + /// + /// See the [unix domain socket example][uds] for an example of how to use + /// this to collect UDS connection info. + /// + /// [`MakeService`]: tower::make::MakeService + /// [`Connected`]: crate::extract::connect_info::Connected + /// [`ConnectInfo`]: crate::extract::connect_info::ConnectInfo + /// [uds]: https://github.com/tokio-rs/axum/blob/main/examples/unix_domain_socket.rs + fn into_make_service_with_connect_info( + self, + ) -> IntoMakeServiceWithConnectInfo + where + Self: Clone, + C: Connected, + { + IntoMakeServiceWithConnectInfo::new(self) + } } impl RoutingDsl for Route {}