From 5bd23d64f5a47cfbd8209d6a4f6b2fb55d0e53a8 Mon Sep 17 00:00:00 2001
From: Anthony Green <agreen@starry.com>
Date: Tue, 30 Nov 2021 09:10:24 -0500
Subject: [PATCH] tonic: implement the server Connected trait for tokio
 UnixStream

This impl is needed in order to use a tokio UnixStream as the
`incoming` argument in methods like
`tonic::transport::server::Router::serve_with_incoming_shutdown`

Fixes: #856

Signed-off-by: Anthony Green <agreen@starry.com>
---
 tonic/src/request.rs               |  4 ++--
 tonic/src/transport/server/conn.rs | 32 +++++++++++++++++++++++++++++-
 tonic/src/transport/server/mod.rs  |  3 +++
 3 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/tonic/src/request.rs b/tonic/src/request.rs
index 46a2d486d..64dd042cf 100644
--- a/tonic/src/request.rs
+++ b/tonic/src/request.rs
@@ -202,8 +202,8 @@ impl<T> Request<T> {
     /// Get the remote address of this connection.
     ///
     /// This will return `None` if the `IO` type used
-    /// does not implement `Connected`. This currently,
-    /// only works on the server side.
+    /// does not implement `Connected` or when using a unix domain socket.
+    /// This currently only works on the server side.
     pub fn remote_addr(&self) -> Option<SocketAddr> {
         #[cfg(feature = "transport")]
         {
diff --git a/tonic/src/transport/server/conn.rs b/tonic/src/transport/server/conn.rs
index 40d60b232..e03d7f847 100644
--- a/tonic/src/transport/server/conn.rs
+++ b/tonic/src/transport/server/conn.rs
@@ -4,7 +4,7 @@ use tokio::net::TcpStream;
 
 #[cfg(feature = "tls")]
 use crate::transport::Certificate;
-#[cfg(feature = "tls")]
+#[cfg(any(unix, feature = "tls"))]
 use std::sync::Arc;
 #[cfg(feature = "tls")]
 use tokio_rustls::{rustls::Session, server::TlsStream};
@@ -98,6 +98,36 @@ impl Connected for TcpStream {
     }
 }
 
+/// Connection info for Unix domain socket streams.
+///
+/// This type will be accessible through [request extensions][ext] if you're using
+/// a unix stream.
+///
+/// See [`Connected`] for more details.
+///
+/// [ext]: crate::Request::extensions
+#[cfg(unix)]
+#[cfg_attr(docsrs, doc(cfg(unix)))]
+#[derive(Clone, Debug)]
+pub struct UdsConnectInfo {
+    /// Peer address. This will be "unnamed" for client unix sockets.
+    pub peer_addr: Option<Arc<tokio::net::unix::SocketAddr>>,
+    /// Process credentials for the unix socket.
+    pub peer_cred: Option<tokio::net::unix::UCred>,
+}
+
+#[cfg(unix)]
+impl Connected for tokio::net::UnixStream {
+    type ConnectInfo = UdsConnectInfo;
+
+    fn connect_info(&self) -> Self::ConnectInfo {
+        UdsConnectInfo {
+            peer_addr: self.peer_addr().ok().map(Arc::new),
+            peer_cred: self.peer_cred().ok(),
+        }
+    }
+}
+
 impl Connected for tokio::io::DuplexStream {
     type ConnectInfo = ();
 
diff --git a/tonic/src/transport/server/mod.rs b/tonic/src/transport/server/mod.rs
index adba6f896..bfefff276 100644
--- a/tonic/src/transport/server/mod.rs
+++ b/tonic/src/transport/server/mod.rs
@@ -17,6 +17,9 @@ pub use conn::TlsConnectInfo;
 #[cfg(feature = "tls")]
 use super::service::TlsAcceptor;
 
+#[cfg(unix)]
+pub use conn::UdsConnectInfo;
+
 use incoming::TcpIncoming;
 
 #[cfg(feature = "tls")]