tonic/transport/server/
conn.rs

1use hyper::server::conn::AddrStream;
2use std::net::SocketAddr;
3use tokio::net::TcpStream;
4
5#[cfg(feature = "tls")]
6use crate::transport::Certificate;
7#[cfg(feature = "tls")]
8use std::sync::Arc;
9#[cfg(feature = "tls")]
10use tokio_rustls::{rustls::Session, server::TlsStream};
11
12/// Trait that connected IO resources implement and use to produce info about the connection.
13///
14/// The goal for this trait is to allow users to implement
15/// custom IO types that can still provide the same connection
16/// metadata.
17///
18/// # Example
19///
20/// The `ConnectInfo` returned will be accessible through [request extensions][ext]:
21///
22/// ```
23/// use tonic::{Request, transport::server::Connected};
24///
25/// // A `Stream` that yields connections
26/// struct MyConnector {}
27///
28/// // Return metadata about the connection as `MyConnectInfo`
29/// impl Connected for MyConnector {
30///     type ConnectInfo = MyConnectInfo;
31///
32///     fn connect_info(&self) -> Self::ConnectInfo {
33///         MyConnectInfo {}
34///     }
35/// }
36///
37/// #[derive(Clone)]
38/// struct MyConnectInfo {
39///     // Metadata about your connection
40/// }
41///
42/// // The connect info can be accessed through request extensions:
43/// # fn foo(request: Request<()>) {
44/// let connect_info: &MyConnectInfo = request
45///     .extensions()
46///     .get::<MyConnectInfo>()
47///     .expect("bug in tonic");
48/// # }
49/// ```
50///
51/// [ext]: crate::Request::extensions
52pub trait Connected {
53    /// The connection info type the IO resources generates.
54    // all these bounds are necessary to set this as a request extension
55    type ConnectInfo: Clone + Send + Sync + 'static;
56
57    /// Create type holding information about the connection.
58    fn connect_info(&self) -> Self::ConnectInfo;
59}
60
61/// Connection info for standard TCP streams.
62///
63/// This type will be accessible through [request extensions][ext] if you're using the default
64/// non-TLS connector.
65///
66/// See [`Connected`] for more details.
67///
68/// [ext]: crate::Request::extensions
69#[derive(Debug, Clone)]
70pub struct TcpConnectInfo {
71    remote_addr: Option<SocketAddr>,
72}
73
74impl TcpConnectInfo {
75    /// Return the remote address the IO resource is connected too.
76    pub fn remote_addr(&self) -> Option<SocketAddr> {
77        self.remote_addr
78    }
79}
80
81impl Connected for AddrStream {
82    type ConnectInfo = TcpConnectInfo;
83
84    fn connect_info(&self) -> Self::ConnectInfo {
85        TcpConnectInfo {
86            remote_addr: Some(self.remote_addr()),
87        }
88    }
89}
90
91impl Connected for TcpStream {
92    type ConnectInfo = TcpConnectInfo;
93
94    fn connect_info(&self) -> Self::ConnectInfo {
95        TcpConnectInfo {
96            remote_addr: self.peer_addr().ok(),
97        }
98    }
99}
100
101impl Connected for tokio::io::DuplexStream {
102    type ConnectInfo = ();
103
104    fn connect_info(&self) -> Self::ConnectInfo {}
105}
106
107#[cfg(feature = "tls")]
108impl<T> Connected for TlsStream<T>
109where
110    T: Connected,
111{
112    type ConnectInfo = TlsConnectInfo<T::ConnectInfo>;
113
114    fn connect_info(&self) -> Self::ConnectInfo {
115        let (inner, session) = self.get_ref();
116        let inner = inner.connect_info();
117
118        let certs = if let Some(certs) = session.get_peer_certificates() {
119            let certs = certs
120                .into_iter()
121                .map(|c| Certificate::from_pem(c.0))
122                .collect();
123            Some(Arc::new(certs))
124        } else {
125            None
126        };
127
128        TlsConnectInfo { inner, certs }
129    }
130}
131
132/// Connection info for TLS streams.
133///
134/// This type will be accessible through [request extensions][ext] if you're using a TLS connector.
135///
136/// See [`Connected`] for more details.
137///
138/// [ext]: crate::Request::extensions
139#[cfg(feature = "tls")]
140#[cfg_attr(docsrs, doc(cfg(feature = "tls")))]
141#[derive(Debug, Clone)]
142pub struct TlsConnectInfo<T> {
143    inner: T,
144    certs: Option<Arc<Vec<Certificate>>>,
145}
146
147#[cfg(feature = "tls")]
148#[cfg_attr(docsrs, doc(cfg(feature = "tls")))]
149impl<T> TlsConnectInfo<T> {
150    /// Get a reference to the underlying connection info.
151    pub fn get_ref(&self) -> &T {
152        &self.inner
153    }
154
155    /// Get a mutable reference to the underlying connection info.
156    pub fn get_mut(&mut self) -> &mut T {
157        &mut self.inner
158    }
159
160    /// Return the set of connected peer TLS certificates.
161    pub fn peer_certs(&self) -> Option<Arc<Vec<Certificate>>> {
162        self.certs.clone()
163    }
164}