Skip to content

Commit

Permalink
feat(transport): Addlocal_addr to Request o (#1327)
Browse files Browse the repository at this point in the history
  • Loading branch information
mokeyish authored Mar 31, 2023
1 parent 44ccf49 commit b54ce23
Show file tree
Hide file tree
Showing 3 changed files with 44 additions and 0 deletions.
2 changes: 2 additions & 0 deletions tests/integration_tests/tests/connect_info.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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());

Expand Down Expand Up @@ -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.
Expand Down
34 changes: 34 additions & 0 deletions tonic/src/request.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
8 changes: 8 additions & 0 deletions tonic/src/transport/server/conn.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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()),
}
}
Expand All @@ -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(),
}
}
Expand Down

0 comments on commit b54ce23

Please sign in to comment.