From b54ce2321a5cba1c32261f4eda2b27d1110b893d Mon Sep 17 00:00:00 2001
From: YiiSh <mokeyish@hotmail.com>
Date: Sat, 1 Apr 2023 01:09:50 +0800
Subject: [PATCH] feat(transport): Add`local_addr` to `Request o` (#1327)

---
 tests/integration_tests/tests/connect_info.rs |  2 ++
 tonic/src/request.rs                          | 34 +++++++++++++++++++
 tonic/src/transport/server/conn.rs            |  8 +++++
 3 files changed, 44 insertions(+)

diff --git a/tests/integration_tests/tests/connect_info.rs b/tests/integration_tests/tests/connect_info.rs
index 4b09a6063..cf612e849 100644
--- 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
index 86cacd6f7..6a57ec93f 100644
--- 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
index 40a303d93..188bff0f2 100644
--- 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(),
         }
     }