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