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}