feat(transport): provide generic access to connect info

examples/src/uds/server.rs

@@ -26,6 +26,11 @@ impl Greeter for MyGreeter {
     ) -> Result<Response<HelloReply>, Status> {
         println!("Got a request: {:?}", request);
 
+        let conn_info = request
+            .connect_info_from_incoming::<unix::UnixStream>()
+            .unwrap();
+        dbg!(&conn_info);
+
         let reply = hello_world::HelloReply {
             message: format!("Hello {}!", request.into_inner().name),
         };
@@ -64,6 +69,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 mod unix {
     use std::{
         pin::Pin,
+        sync::Arc,
         task::{Context, Poll},
     };
 
@@ -73,7 +79,22 @@ mod unix {
     #[derive(Debug)]
     pub struct UnixStream(pub tokio::net::UnixStream);
 
-    impl Connected for UnixStream {}
+    impl Connected for UnixStream {
+        type ConnectInfo = UdsConnectInfo;
+
+        fn connect_info(&self) -> Self::ConnectInfo {
+            UdsConnectInfo {
+                peer_addr: self.0.peer_addr().ok().map(Arc::new),
+                peer_cred: self.0.peer_cred().ok(),
+            }
+        }
+    }
+
+    #[derive(Clone, Debug)]
+    pub struct UdsConnectInfo {
+        pub peer_addr: Option<Arc<tokio::net::unix::SocketAddr>>,
+        pub peer_cred: Option<tokio::net::unix::UCred>,
+    }
 
     impl AsyncRead for UnixStream {
         fn poll_read(

tests/integration_tests/tests/connect_info.rs

@@ -0,0 +1,52 @@
+use futures_util::FutureExt;
+use integration_tests::pb::{test_client, test_server, Input, Output};
+use std::time::Duration;
+use tokio::sync::oneshot;
+use tonic::{
+    transport::{Endpoint, Server},
+    Request, Response, Status,
+};
+
+#[tokio::test]
+async fn getting_connect_info() {
+    struct Svc;
+
+    #[tonic::async_trait]
+    impl test_server::Test for Svc {
+        async fn unary_call(&self, req: Request<Input>) -> Result<Response<Output>, Status> {
+            match req.metadata().get("user-agent") {
+                Some(_) => Ok(Response::new(Output {})),
+                None => Err(Status::internal("user-agent header is missing")),
+            }
+        }
+    }
+
+    let svc = test_server::TestServer::new(Svc);
+
+    let (tx, rx) = oneshot::channel::<()>();
+
+    let jh = tokio::spawn(async move {
+        Server::builder()
+            .add_service(svc)
+            .serve_with_shutdown("127.0.0.1:1400".parse().unwrap(), rx.map(drop))
+            .await
+            .unwrap();
+    });
+
+    tokio::time::sleep(Duration::from_millis(100)).await;
+
+    let channel = Endpoint::from_static("http://127.0.0.1:1400")
+        .user_agent("my-client")
+        .expect("valid user agent")
+        .connect()
+        .await
+        .unwrap();
+
+    let mut client = test_client::TestClient::new(channel);
+
+    client.unary_call(Input {}).await.unwrap();
+
+    tx.send(()).unwrap();
+
+    jh.await.unwrap();
+}

tonic/src/request.rs

@@ -1,6 +1,8 @@
 use crate::metadata::{MetadataMap, MetadataValue};
+#[cfg(all(feature = "transport", feature = "tls"))]
+use crate::transport::server::TlsConnectInfo;
 #[cfg(feature = "transport")]
-use crate::transport::Certificate;
+use crate::transport::{server::TcpConnectInfo, Certificate};
 use crate::Extensions;
 use futures_core::Stream;
 #[cfg(feature = "transport")]
@@ -15,13 +17,6 @@ pub struct Request<T> {
     extensions: Extensions,
 }
 
-#[derive(Clone)]
-pub(crate) struct ConnectionInfo {
-    pub(crate) remote_addr: Option<SocketAddr>,
-    #[cfg(feature = "transport")]
-    pub(crate) peer_certs: Option<Arc<Vec<Certificate>>>,
-}
-
 /// Trait implemented by RPC request types.
 ///
 /// Types implementing this trait can be used as arguments to client RPC
@@ -203,7 +198,28 @@ impl<T> Request<T> {
     /// does not implement `Connected`. This currently,
     /// only works on the server side.
     pub fn remote_addr(&self) -> Option<SocketAddr> {
-        self.get::<ConnectionInfo>()?.remote_addr
+        #[cfg(feature = "transport")]
+        {
+            #[cfg(feature = "tls")]
+            {
+                self.get::<TcpConnectInfo>()
+                    .and_then(|i| i.remote_addr())
+                    .or_else(|| {
+                        self.get::<TlsConnectInfo<TcpConnectInfo>>()
+                            .and_then(|i| i.get_ref().remote_addr())
+                    })
+            }
+
+            #[cfg(not(feature = "tls"))]
+            {
+                self.get::<TcpConnectInfo>().and_then(|i| i.remote_addr())
+            }
+        }
+
+        #[cfg(not(feature = "transport"))]
+        {
+            None
+        }
     }
 
     /// Get the peer certificates of the connected client.
@@ -215,9 +231,19 @@ impl<T> Request<T> {
     #[cfg(feature = "transport")]
     #[cfg_attr(docsrs, doc(cfg(feature = "transport")))]
     pub fn peer_certs(&self) -> Option<Arc<Vec<Certificate>>> {
-        self.get::<ConnectionInfo>()?.peer_certs.clone()
+        #[cfg(feature = "tls")]
+        {
+            self.get::<TlsConnectInfo<TcpConnectInfo>>()
+                .and_then(|i| i.peer_certs())
+        }
+
+        #[cfg(not(feature = "tls"))]
+        {
+            None
+        }
     }
 
+    #[allow(dead_code)]
     pub(crate) fn get<I: Send + Sync + 'static>(&self) -> Option<&I> {
         self.extensions.get::<I>()
     }
@@ -308,6 +334,15 @@ impl<T> Request<T> {
     pub fn extensions_mut(&mut self) -> &mut Extensions {
         &mut self.extensions
     }
+
+    /// TODO(david)
+    #[cfg(feature = "transport")]
+    pub fn connect_info_from_incoming<C>(&self) -> Option<&C::ConnectInfo>
+    where
+        C: crate::transport::server::Connected,
+    {
+        self.get()
+    }
 }
 
 impl<T> IntoRequest<T> for T {

tonic/src/transport/server/conn.rs

@@ -1,7 +1,11 @@
-use crate::transport::Certificate;
 use hyper::server::conn::AddrStream;
 use std::net::SocketAddr;
 use tokio::net::TcpStream;
+
+#[cfg(feature = "tls")]
+use crate::transport::Certificate;
+#[cfg(feature = "tls")]
+use std::sync::Arc;
 #[cfg(feature = "tls")]
 use tokio_rustls::{rustls::Session, server::TlsStream};
 
@@ -11,48 +15,90 @@ use tokio_rustls::{rustls::Session, server::TlsStream};
 /// custom IO types that can still provide the same connection
 /// metadata.
 pub trait Connected {
-    /// Return the remote address this IO resource is connected too.
-    fn remote_addr(&self) -> Option<SocketAddr> {
-        None
-    }
+    /// TODO(david)
+    // all these bounds are necessary to set this as a request extension
+    type ConnectInfo: Clone + Send + Sync + 'static;
+
+    /// TODO(david)
+    fn connect_info(&self) -> Self::ConnectInfo;
+}
+
+/// TODO(david)
+#[derive(Debug, Clone)]
+pub struct TcpConnectInfo {
+    remote_addr: Option<SocketAddr>,
+}
 
-    /// Return the set of connected peer TLS certificates.
-    fn peer_certs(&self) -> Option<Vec<Certificate>> {
-        None
+impl TcpConnectInfo {
+    /// TODO(david)
+    pub fn remote_addr(&self) -> Option<SocketAddr> {
+        self.remote_addr
     }
 }
 
 impl Connected for AddrStream {
-    fn remote_addr(&self) -> Option<SocketAddr> {
-        Some(self.remote_addr())
+    type ConnectInfo = TcpConnectInfo;
+
+    fn connect_info(&self) -> Self::ConnectInfo {
+        TcpConnectInfo {
+            remote_addr: Some(self.remote_addr()),
+        }
     }
 }
 
 impl Connected for TcpStream {
-    fn remote_addr(&self) -> Option<SocketAddr> {
-        self.peer_addr().ok()
+    type ConnectInfo = TcpConnectInfo;
+
+    fn connect_info(&self) -> Self::ConnectInfo {
+        TcpConnectInfo {
+            remote_addr: self.peer_addr().ok(),
+        }
     }
 }
 
 #[cfg(feature = "tls")]
-impl<T: Connected> Connected for TlsStream<T> {
-    fn remote_addr(&self) -> Option<SocketAddr> {
-        let (inner, _) = self.get_ref();
+impl<T> Connected for TlsStream<T>
+where
+    T: Connected,
+{
+    type ConnectInfo = TlsConnectInfo<T::ConnectInfo>;
 
-        inner.remote_addr()
-    }
-
-    fn peer_certs(&self) -> Option<Vec<Certificate>> {
-        let (_, session) = self.get_ref();
+    fn connect_info(&self) -> Self::ConnectInfo {
+        let (inner, session) = self.get_ref();
+        let inner = inner.connect_info();
 
-        if let Some(certs) = session.get_peer_certificates() {
+        let certs = if let Some(certs) = session.get_peer_certificates() {
             let certs = certs
                 .into_iter()
                 .map(|c| Certificate::from_pem(c.0))
                 .collect();
-            Some(certs)
+            Some(Arc::new(certs))
         } else {
             None
-        }
+        };
+
+        TlsConnectInfo { inner, certs }
+    }
+}
+
+/// TODO(david)
+#[cfg(feature = "tls")]
+#[derive(Debug, Clone)]
+pub struct TlsConnectInfo<T> {
+    inner: T,
+    certs: Option<Arc<Vec<Certificate>>>,
+}
+
+/// TODO(david)
+#[cfg(feature = "tls")]
+impl<T> TlsConnectInfo<T> {
+    /// TODO(david)
+    pub fn get_ref(&self) -> &T {
+        &self.inner
+    }
+
+    /// TODO(david)
+    pub fn peer_certs(&self) -> Option<Arc<Vec<Certificate>>> {
+        self.certs.clone()
     }
 }

tonic/src/transport/server/incoming.rs

@@ -18,18 +18,16 @@ use tokio::io::{AsyncRead, AsyncWrite};
 pub(crate) fn tcp_incoming<IO, IE>(
     incoming: impl Stream<Item = Result<IO, IE>>,
     _server: Server,
-) -> impl Stream<Item = Result<ServerIo, crate::Error>>
+) -> impl Stream<Item = Result<ServerIo<IO>, crate::Error>>
 where
     IO: AsyncRead + AsyncWrite + Connected + Unpin + Send + 'static,
     IE: Into<crate::Error>,
 {
     async_stream::try_stream! {
         futures_util::pin_mut!(incoming);
 
-
         while let Some(stream) = incoming.try_next().await? {
-
-            yield ServerIo::new(stream);
+            yield ServerIo::new_io(stream);
         }
     }
 }
@@ -38,7 +36,7 @@ where
 pub(crate) fn tcp_incoming<IO, IE>(
     incoming: impl Stream<Item = Result<IO, IE>>,
     server: Server,
-) -> impl Stream<Item = Result<ServerIo, crate::Error>>
+) -> impl Stream<Item = Result<ServerIo<IO>, crate::Error>>
 where
     IO: AsyncRead + AsyncWrite + Connected + Unpin + Send + 'static,
     IE: Into<crate::Error>,
@@ -57,12 +55,12 @@ where
 
                         let accept = tokio::spawn(async move {
                             let io = tls.accept(stream).await?;
-                            Ok(ServerIo::new(io))
+                            Ok(ServerIo::new_tls_io(io))
                         });
 
                         tasks.push(accept);
                     } else {
-                        yield ServerIo::new(stream);
+                        yield ServerIo::new_io(stream);
                     }
                 }
 
@@ -86,7 +84,7 @@ where
 async fn select<IO, IE>(
     incoming: &mut (impl Stream<Item = Result<IO, IE>> + Unpin),
     tasks: &mut futures_util::stream::futures_unordered::FuturesUnordered<
-        tokio::task::JoinHandle<Result<ServerIo, crate::Error>>,
+        tokio::task::JoinHandle<Result<ServerIo<IO>, crate::Error>>,
     >,
 ) -> SelectOutput<IO>
 where
@@ -124,7 +122,7 @@ where
 #[cfg(feature = "tls")]
 enum SelectOutput<A> {
     Incoming(A),
-    Io(ServerIo),
+    Io(ServerIo<A>),
     Err(crate::Error),
     Done,
 }