feat(transport): Add remote_addr to Request on the server si… (#186)

examples/src/helloworld/server.rs

@@ -18,7 +18,7 @@ impl Greeter for MyGreeter {
  &self,
  request: Request<HelloRequest>,
  ) -> Result<Response<HelloReply>, Status> {
- println!("Got a request: {:?}", request);
+ println!("Got a request from {:?}", request.remote_addr());
 
  let reply = hello_world::HelloReply {
  message: format!("Hello {}!", request.into_inner().name).into(),

examples/src/uds/server.rs

@@ -1,6 +1,17 @@
-use std::path::Path;
-use tokio::net::UnixListener;
-use tonic::{transport::Server, Request, Response, Status};
+use futures::stream::TryStreamExt;
+use std::{
+ path::Path,
+ pin::Pin,
+ task::{Context, Poll},
+};
+use tokio::{
+ io::{AsyncRead, AsyncWrite},
+ net::UnixListener,
+};
+use tonic::{
+ transport::{server::Connected, Server},
+ Request, Response, Status,
+};
 
 pub mod hello_world {
  tonic::include_proto!("helloworld");
@@ -41,8 +52,41 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 
  Server::builder()
  .add_service(GreeterServer::new(greeter))
- .serve_with_incoming(uds.incoming())
+ .serve_with_incoming(uds.incoming().map_ok(UnixStream))
  .await?;
 
  Ok(())
 }
+
+#[derive(Debug)]
+struct UnixStream(tokio::net::UnixStream);
+
+impl Connected for UnixStream {}
+
+impl AsyncRead for UnixStream {
+ fn poll_read(
+ mut self: Pin<&mut Self>,
+ cx: &mut Context<'_>,
+ buf: &mut [u8],
+ ) -> Poll<std::io::Result<usize>> {
+ Pin::new(&mut self.0).poll_read(cx, buf)
+ }
+}
+
+impl AsyncWrite for UnixStream {
+ fn poll_write(
+ mut self: Pin<&mut Self>,
+ cx: &mut Context<'_>,
+ buf: &[u8],
+ ) -> Poll<std::io::Result<usize>> {
+ Pin::new(&mut self.0).poll_write(cx, buf)
+ }
+
+ fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
+ Pin::new(&mut self.0).poll_flush(cx)
+ }
+
+ fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
+ Pin::new(&mut self.0).poll_shutdown(cx)
+ }
+}

tonic/src/request.rs

@@ -1,11 +1,19 @@
 use crate::metadata::MetadataMap;
 use futures_core::Stream;
+use http::Extensions;
+use std::net::SocketAddr;
 
 /// A gRPC request and metadata from an RPC call.
 #[derive(Debug)]
 pub struct Request<T> {
  metadata: MetadataMap,
  message: T,
+ extensions: Extensions,
+}
+
+#[derive(Clone)]
+pub(crate) struct ConnectionInfo {
+ pub(crate) remote_addr: Option<SocketAddr>,
 }
 
 /// Trait implemented by RPC request types.
@@ -102,6 +110,7 @@ impl<T> Request<T> {
  Request {
  metadata: MetadataMap::new(),
  message,
+ extensions: Extensions::default(),
  }
  }
 
@@ -134,6 +143,7 @@ impl<T> Request<T> {
  Request {
  metadata: MetadataMap::from_headers(parts.headers),
  message,
+ extensions: parts.extensions,
  }
  }
 
@@ -150,6 +160,7 @@ impl<T> Request<T> {
  *request.method_mut() = http::Method::POST;
  *request.uri_mut() = uri;
  *request.headers_mut() = self.metadata.into_sanitized_headers();
+ *request.extensions_mut() = self.extensions;
 
  request
  }
@@ -164,8 +175,22 @@ impl<T> Request<T> {
  Request {
  metadata: self.metadata,
  message,
+ extensions: Extensions::default(),
  }
  }
+
+ /// 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.
+ pub fn remote_addr(&self) -> Option<SocketAddr> {
+ self.get::<ConnectionInfo>()?.remote_addr
+ }
+
+ pub(crate) fn get<I: Send + Sync + 'static>(&self) -> Option<&I> {
+ self.extensions.get::<I>()
+ }
 }
 
 impl<T> IntoRequest<T> for T {

tonic/src/transport/server/conn.rs

@@ -0,0 +1,30 @@
+use hyper::server::conn::AddrStream;
+use std::net::SocketAddr;
+#[cfg(feature = "tls")]
+use tokio_rustls::TlsStream;
+
+/// Trait that connected IO resources implement.
+///
+/// The goal for this trait is to allow users to implement
+/// 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
+ }
+}
+
+impl Connected for AddrStream {
+ fn remote_addr(&self) -> Option<SocketAddr> {
+ Some(self.remote_addr())
+ }
+}
+
+#[cfg(feature = "tls")]
+impl<T: Connected> Connected for TlsStream<T> {
+ fn remote_addr(&self) -> Option<SocketAddr> {
+ let (inner, _) = self.get_ref();
+ inner.remote_addr()
+ }
+}

tonic/src/transport/server/incoming.rs

@@ -1,5 +1,5 @@
-use super::Server;
-use crate::transport::service::BoxedIo;
+use super::{Connected, Server};
+use crate::transport::service::ServerIo;
 use futures_core::Stream;
 use futures_util::stream::TryStreamExt;
 use hyper::server::{
@@ -20,9 +20,9 @@ use tracing::error;
 pub(crate) fn tcp_incoming<IO, IE>(
  incoming: impl Stream<Item = Result<IO, IE>>,
  server: Server,
-) -> impl Stream<Item = Result<BoxedIo, crate::Error>>
+) -> impl Stream<Item = Result<ServerIo, crate::Error>>
 where
- IO: AsyncRead + AsyncWrite + Unpin + Send + 'static,
+ IO: AsyncRead + AsyncWrite + Connected + Unpin + Send + 'static,
  IE: Into<crate::Error>,
 {
  async_stream::try_stream! {
@@ -39,12 +39,12 @@ where
  continue
  },
  };
- yield BoxedIo::new(io);
+ yield ServerIo::new(io);
  continue;
  }
  }
 
- yield BoxedIo::new(stream);
+ yield ServerIo::new(stream);
  }
  }
 }

tonic/src/transport/server/mod.rs

@@ -1,9 +1,11 @@
 //! Server implementation and builder.
 
+mod conn;
 mod incoming;
 #[cfg(feature = "tls")]
 mod tls;
 
+pub use conn::Connected;
 #[cfg(feature = "tls")]
 pub use tls::ServerTlsConfig;
 
@@ -12,8 +14,8 @@ use super::service::TlsAcceptor;
 
 use incoming::TcpIncoming;
 
-use super::service::{layer_fn, Or, Routes, ServiceBuilderExt};
-use crate::body::BoxBody;
+use super::service::{layer_fn, Or, Routes, ServerIo, ServiceBuilderExt};
+use crate::{body::BoxBody, request::ConnectionInfo};
 use futures_core::Stream;
 use futures_util::{
  future::{self, MapErr},
@@ -252,7 +254,7 @@ impl Server {
  S::Future: Send + 'static,
  S::Error: Into<crate::Error> + Send,
  I: Stream<Item = Result<IO, IE>>,
- IO: AsyncRead + AsyncWrite + Unpin + Send + 'static,
+ IO: AsyncRead + AsyncWrite + Connected + Unpin + Send + 'static,
  IE: Into<crate::Error>,
  F: Future<Output = ()>,
  {
@@ -390,7 +392,7 @@ where
  pub async fn serve_with_incoming<I, IO, IE>(self, incoming: I) -> Result<(), super::Error>
  where
  I: Stream<Item = Result<IO, IE>>,
- IO: AsyncRead + AsyncWrite + Unpin + Send + 'static,
+ IO: AsyncRead + AsyncWrite + Connected + Unpin + Send + 'static,
  IE: Into<crate::Error>,
  {
  self.server
@@ -412,6 +414,7 @@ impl fmt::Debug for Server {
 struct Svc<S> {
  inner: S,
  span: Option<TraceInterceptor>,
+ conn_info: ConnectionInfo,
 }
 
 impl<S> Service<Request<Body>> for Svc<S>
@@ -427,13 +430,15 @@ where
  self.inner.poll_ready(cx).map_err(Into::into)
  }
 
- fn call(&mut self, req: Request<Body>) -> Self::Future {
+ fn call(&mut self, mut req: Request<Body>) -> Self::Future {
  let span = if let Some(trace_interceptor) = &self.span {
  trace_interceptor(req.headers())
  } else {
  tracing::Span::none()
  };
 
+ req.extensions_mut().insert(self.conn_info.clone());
+
  self.inner.call(req).instrument(span).map_err(|e| e.into())
  }
 }
@@ -452,7 +457,7 @@ struct MakeSvc<S> {
  span: Option<TraceInterceptor>,
 }
 
-impl<S, T> Service<T> for MakeSvc<S>
+impl<S> Service<&ServerIo> for MakeSvc<S>
 where
  S: Service<Request<Body>, Response = Response<BoxBody>> + Clone + Send + 'static,
  S::Future: Send + 'static,
@@ -467,7 +472,11 @@ where
  Ok(()).into()
  }
 
- fn call(&mut self, _: T) -> Self::Future {
+ fn call(&mut self, io: &ServerIo) -> Self::Future {
+ let conn_info = crate::request::ConnectionInfo {
+ remote_addr: io.remote_addr(),
+ };
+
  let interceptor = self.interceptor.clone();
  let svc = self.inner.clone();
  let concurrency_limit = self.concurrency_limit;
@@ -481,10 +490,18 @@ where
  .service(svc);
 
  let svc = if let Some(interceptor) = interceptor {
- let layered = interceptor.layer(BoxService::new(Svc { inner: svc, span }));
+ let layered = interceptor.layer(BoxService::new(Svc {
+ inner: svc,
+ span,
+ conn_info,
+ }));
  BoxService::new(layered)
  } else {
- BoxService::new(Svc { inner: svc, span })
+ BoxService::new(Svc {
+ inner: svc,
+ span,
+ conn_info,
+ })
  };
 
  Ok(svc)

tonic/src/transport/service/io.rs

@@ -1,5 +1,7 @@
-use hyper::client::connect::{Connected, Connection};
+use crate::transport::server::Connected;
+use hyper::client::connect::{Connected as HyperConnected, Connection};
 use std::io;
+use std::net::SocketAddr;
 use std::pin::Pin;
 use std::task::{Context, Poll};
 use tokio::io::{AsyncRead, AsyncWrite};
@@ -20,11 +22,13 @@ impl BoxedIo {
 }
 
 impl Connection for BoxedIo {
- fn connected(&self) -> Connected {
- Connected::new()
+ fn connected(&self) -> HyperConnected {
+ HyperConnected::new()
  }
 }
 
+impl Connected for BoxedIo {}
+
 impl AsyncRead for BoxedIo {
  fn poll_read(
  mut self: Pin<&mut Self>,
@@ -52,3 +56,50 @@ impl AsyncWrite for BoxedIo {
  Pin::new(&mut self.0).poll_shutdown(cx)
  }
 }
+
+pub(in crate::transport) trait ConnectedIo: Io + Connected {}
+
+impl<T> ConnectedIo for T where T: Io + Connected {}
+
+pub(crate) struct ServerIo(Pin<Box<dyn ConnectedIo>>);
+
+impl ServerIo {
+ pub(in crate::transport) fn new<I: ConnectedIo>(io: I) -> Self {
+ ServerIo(Box::pin(io))
+ }
+}
+
+impl Connected for ServerIo {
+ fn remote_addr(&self) -> Option<SocketAddr> {
+ let io = &*self.0;
+ io.remote_addr()
+ }
+}
+
+impl AsyncRead for ServerIo {
+ fn poll_read(
+ mut self: Pin<&mut Self>,
+ cx: &mut Context<'_>,
+ buf: &mut [u8],
+ ) -> Poll<io::Result<usize>> {
+ Pin::new(&mut self.0).poll_read(cx, buf)
+ }
+}
+
+impl AsyncWrite for ServerIo {
+ fn poll_write(
+ mut self: Pin<&mut Self>,
+ cx: &mut Context<'_>,
+ buf: &[u8],
+ ) -> Poll<io::Result<usize>> {
+ Pin::new(&mut self.0).poll_write(cx, buf)
+ }
+
+ fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+ Pin::new(&mut self.0).poll_flush(cx)
+ }
+
+ fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+ Pin::new(&mut self.0).poll_shutdown(cx)
+ }
+}

tonic/src/transport/service/mod.rs

@@ -14,7 +14,7 @@ pub(crate) use self::add_origin::AddOrigin;
 pub(crate) use self::connection::Connection;
 pub(crate) use self::connector::connector;
 pub(crate) use self::discover::ServiceList;
-pub(crate) use self::io::BoxedIo;
+pub(crate) use self::io::ServerIo;
 pub(crate) use self::layer::{layer_fn, ServiceBuilderExt};
 pub(crate) use self::router::{Or, Routes};
 #[cfg(feature = "tls")]