feat(transport): Addlocal_addr to Request o (#1327)

hyperium · Mar 31, 2023 · b54ce23 · b54ce23
1 parent 44ccf49
commit b54ce23
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 0 deletions.
diff --git a/tests/integration_tests/tests/connect_info.rs b/tests/integration_tests/tests/connect_info.rs
@@ -14,6 +14,7 @@ async fn getting_connect_info() {
     #[tonic::async_trait]
     impl test_server::Test for Svc {
         async fn unary_call(&self, req: Request<Input>) -> Result<Response<Output>, Status> {
+            assert!(req.local_addr().is_some());
             assert!(req.remote_addr().is_some());
             assert!(req.extensions().get::<TcpConnectInfo>().is_some());
 
@@ -73,6 +74,7 @@ pub mod unix {
             let conn_info = req.extensions().get::<UdsConnectInfo>().unwrap();
 
             // Client-side unix sockets are unnamed.
+            assert!(req.local_addr().is_none());
             assert!(req.remote_addr().is_none());
             assert!(conn_info.peer_addr.as_ref().unwrap().is_unnamed());
             // This should contain process credentials for the client socket.

diff --git a/tonic/src/request.rs b/tonic/src/request.rs
@@ -203,6 +203,40 @@ impl<T> Request<T> {
         }
     }
 
+    /// Get the local address of this connection.
+    ///
+    /// This will return `None` if the `IO` type used
+    /// does not implement `Connected` or when using a unix domain socket.
+    /// This currently only works on the server side.
+    pub fn local_addr(&self) -> Option<SocketAddr> {
+        #[cfg(feature = "transport")]
+        {
+            #[cfg(feature = "tls")]
+            {
+                self.extensions()
+                    .get::<TcpConnectInfo>()
+                    .and_then(|i| i.local_addr())
+                    .or_else(|| {
+                        self.extensions()
+                            .get::<TlsConnectInfo<TcpConnectInfo>>()
+                            .and_then(|i| i.get_ref().local_addr())
+                    })
+            }
+
+            #[cfg(not(feature = "tls"))]
+            {
+                self.extensions()
+                    .get::<TcpConnectInfo>()
+                    .and_then(|i| i.local_addr())
+            }
+        }
+
+        #[cfg(not(feature = "transport"))]
+        {
+            None
+        }
+    }
+
     /// Get the remote address of this connection.
     ///
     /// This will return `None` if the `IO` type used

diff --git a/tonic/src/transport/server/conn.rs b/tonic/src/transport/server/conn.rs
@@ -68,10 +68,16 @@ pub trait Connected {
 /// [ext]: crate::Request::extensions
 #[derive(Debug, Clone)]
 pub struct TcpConnectInfo {
+    local_addr: Option<SocketAddr>,
     remote_addr: Option<SocketAddr>,
 }
 
 impl TcpConnectInfo {
+    /// Return the local address the IO resource is connected.
+    pub fn local_addr(&self) -> Option<SocketAddr> {
+        self.local_addr
+    }
+
     /// Return the remote address the IO resource is connected too.
     pub fn remote_addr(&self) -> Option<SocketAddr> {
         self.remote_addr
@@ -83,6 +89,7 @@ impl Connected for AddrStream {
 
     fn connect_info(&self) -> Self::ConnectInfo {
         TcpConnectInfo {
+            local_addr: Some(self.local_addr()),
             remote_addr: Some(self.remote_addr()),
         }
     }
@@ -93,6 +100,7 @@ impl Connected for TcpStream {
 
     fn connect_info(&self) -> Self::ConnectInfo {
         TcpConnectInfo {
+            local_addr: self.local_addr().ok(),
             remote_addr: self.peer_addr().ok(),
         }
     }