Add local_addr for Serve and WithGracefulShutdown (#2881)

tokio-rs · Aug 24, 2024 · 1ac617a · 1ac617a
1 parent ef1448a
commit 1ac617a
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 0 deletions.
diff --git a/axum/CHANGELOG.md b/axum/CHANGELOG.md
@@ -11,10 +11,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - **added:** `axum::serve::Serve::tcp_nodelay` and `axum::serve::WithGracefulShutdown::tcp_nodelay` ([#2653])
 - **added:** `Router::has_routes` function ([#2790])
 - **change:** Update tokio-tungstenite to 0.23 ([#2841])
+- **added:** `Serve::local_addr` and `WithGracefulShutdown::local_addr` functions ([#2881])
 
 [#2653]: https://github.com/tokio-rs/axum/pull/2653
 [#2790]: https://github.com/tokio-rs/axum/pull/2790
 [#2841]: https://github.com/tokio-rs/axum/pull/2841
+[#2881]: https://github.com/tokio-rs/axum/pull/2881
 
 # 0.7.5 (24. March, 2024)
 

diff --git a/axum/src/serve.rs b/axum/src/serve.rs
@@ -173,6 +173,11 @@ impl<M, S> Serve<M, S> {
             ..self
         }
     }
+
+    /// Returns the local address this server is bound to.
+    pub fn local_addr(&self) -> io::Result<SocketAddr> {
+        self.tcp_listener.local_addr()
+    }
 }
 
 #[cfg(all(feature = "tokio", any(feature = "http1", feature = "http2")))]
@@ -307,6 +312,11 @@ impl<M, S, F> WithGracefulShutdown<M, S, F> {
             ..self
         }
     }
+
+    /// Returns the local address this server is bound to.
+    pub fn local_addr(&self) -> io::Result<SocketAddr> {
+        self.tcp_listener.local_addr()
+    }
 }
 
 #[cfg(all(feature = "tokio", any(feature = "http1", feature = "http2")))]
@@ -544,6 +554,10 @@ mod tests {
         routing::get,
         Router,
     };
+    use std::{
+        future::pending,
+        net::{IpAddr, Ipv4Addr},
+    };
 
     #[allow(dead_code, unused_must_use)]
     async fn if_it_compiles_it_works() {
@@ -607,4 +621,29 @@ mod tests {
     }
 
     async fn handler() {}
+
+    #[crate::test]
+    async fn test_serve_local_addr() {
+        let router: Router = Router::new();
+        let addr = "0.0.0.0:0";
+
+        let server = serve(TcpListener::bind(addr).await.unwrap(), router.clone());
+        let address = server.local_addr().unwrap();
+
+        assert_eq!(address.ip(), IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)));
+        assert_ne!(address.port(), 0);
+    }
+
+    #[crate::test]
+    async fn test_with_graceful_shutdown_local_addr() {
+        let router: Router = Router::new();
+        let addr = "0.0.0.0:0";
+
+        let server = serve(TcpListener::bind(addr).await.unwrap(), router.clone())
+            .with_graceful_shutdown(pending());
+        let address = server.local_addr().unwrap();
+
+        assert_eq!(address.ip(), IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)));
+        assert_ne!(address.port(), 0);
+    }
 }