tokio_websockets/
tls.rs

1//! Wrapper types for TLS functionality, abstracting over [`rustls`] and
2//! [`native-tls`] connector and stream types.
3//!
4//! [`native-tls`]: tokio_native_tls::native_tls
5//! [`rustls`]: tokio_rustls::rustls
6
7#[cfg(any(
8    feature = "rustls-webpki-roots",
9    feature = "rustls-native-roots",
10    feature = "rustls-platform-verifier"
11))]
12use std::sync::Arc;
13use std::{
14    fmt::{Debug, Formatter, Result as FmtResult},
15    io,
16    pin::Pin,
17    task::{Context, Poll},
18};
19
20#[cfg(any(
21    feature = "rustls-webpki-roots",
22    feature = "rustls-native-roots",
23    feature = "rustls-platform-verifier",
24    feature = "rustls-bring-your-own-connector"
25))]
26use rustls_pki_types::ServerName;
27#[cfg(feature = "rustls-platform-verifier")]
28use rustls_platform_verifier::BuilderVerifierExt;
29use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
30#[cfg(any(
31    feature = "rustls-native-roots",
32    feature = "rustls-webpki-roots",
33    feature = "rustls-platform-verifier"
34))]
35use tokio_rustls::rustls::ClientConfig;
36
37use crate::Error;
38
39/// A reusable TLS connector for wrapping streams.
40pub enum Connector {
41    /// Plain (non-TLS) connector.
42    Plain,
43    /// [`native-tls`] TLS connector.
44    ///
45    /// [`native-tls`]: tokio_native_tls::native_tls
46    #[cfg(feature = "native-tls")]
47    NativeTls(tokio_native_tls::TlsConnector),
48    /// [`rustls`] TLS connector.
49    ///
50    /// [`rustls`]: tokio_rustls::rustls
51    #[cfg(any(
52        feature = "rustls-native-roots",
53        feature = "rustls-webpki-roots",
54        feature = "rustls-platform-verifier",
55        feature = "rustls-bring-your-own-connector"
56    ))]
57    Rustls(tokio_rustls::TlsConnector),
58}
59
60impl Debug for Connector {
61    fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
62        match self {
63            Self::Plain => f.write_str("Connector::Plain"),
64            #[cfg(feature = "native-tls")]
65            Self::NativeTls(connector) => connector.fmt(f),
66            #[cfg(any(
67                feature = "rustls-native-roots",
68                feature = "rustls-webpki-roots",
69                feature = "rustls-platform-verifier",
70                feature = "rustls-bring-your-own-connector"
71            ))]
72            Self::Rustls(_) => f.write_str("Connector::Rustls"),
73        }
74    }
75}
76
77/// A stream that might be protected with TLS.
78#[allow(clippy::large_enum_variant)] // Only one or two of these will be used
79#[derive(Debug)]
80pub enum MaybeTlsStream<S> {
81    /// Unencrypted socket stream.
82    Plain(S),
83    /// Encrypted socket stream using [`native-tls`].
84    ///
85    /// [`native-tls`]: tokio_native_tls::native_tls
86    #[cfg(feature = "native-tls")]
87    NativeTls(tokio_native_tls::TlsStream<S>),
88    /// Encrypted socket stream using [`rustls`].
89    ///
90    /// [`rustls`]: tokio_rustls::rustls
91    #[cfg(any(
92        feature = "rustls-native-roots",
93        feature = "rustls-webpki-roots",
94        feature = "rustls-platform-verifier",
95        feature = "rustls-bring-your-own-connector"
96    ))]
97    Rustls(tokio_rustls::client::TlsStream<S>),
98}
99
100impl<S: AsyncRead + AsyncWrite + Unpin> AsyncRead for MaybeTlsStream<S> {
101    fn poll_read(
102        self: Pin<&mut Self>,
103        cx: &mut Context<'_>,
104        buf: &mut ReadBuf<'_>,
105    ) -> Poll<io::Result<()>> {
106        match self.get_mut() {
107            Self::Plain(s) => Pin::new(s).poll_read(cx, buf),
108            #[cfg(feature = "native-tls")]
109            Self::NativeTls(s) => Pin::new(s).poll_read(cx, buf),
110            #[cfg(any(
111                feature = "rustls-native-roots",
112                feature = "rustls-webpki-roots",
113                feature = "rustls-platform-verifier",
114                feature = "rustls-bring-your-own-connector"
115            ))]
116            Self::Rustls(s) => Pin::new(s).poll_read(cx, buf),
117        }
118    }
119}
120
121impl<S: AsyncRead + AsyncWrite + Unpin> AsyncWrite for MaybeTlsStream<S> {
122    fn poll_write(
123        self: Pin<&mut Self>,
124        cx: &mut Context<'_>,
125        buf: &[u8],
126    ) -> Poll<Result<usize, io::Error>> {
127        match self.get_mut() {
128            Self::Plain(s) => Pin::new(s).poll_write(cx, buf),
129            #[cfg(feature = "native-tls")]
130            Self::NativeTls(s) => Pin::new(s).poll_write(cx, buf),
131            #[cfg(any(
132                feature = "rustls-native-roots",
133                feature = "rustls-webpki-roots",
134                feature = "rustls-platform-verifier",
135                feature = "rustls-bring-your-own-connector"
136            ))]
137            Self::Rustls(s) => Pin::new(s).poll_write(cx, buf),
138        }
139    }
140
141    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
142        match self.get_mut() {
143            Self::Plain(s) => Pin::new(s).poll_flush(cx),
144            #[cfg(feature = "native-tls")]
145            Self::NativeTls(s) => Pin::new(s).poll_flush(cx),
146            #[cfg(any(
147                feature = "rustls-native-roots",
148                feature = "rustls-webpki-roots",
149                feature = "rustls-platform-verifier",
150                feature = "rustls-bring-your-own-connector"
151            ))]
152            Self::Rustls(s) => Pin::new(s).poll_flush(cx),
153        }
154    }
155
156    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
157        match self.get_mut() {
158            Self::Plain(s) => Pin::new(s).poll_shutdown(cx),
159            #[cfg(feature = "native-tls")]
160            Self::NativeTls(s) => Pin::new(s).poll_shutdown(cx),
161            #[cfg(any(
162                feature = "rustls-native-roots",
163                feature = "rustls-webpki-roots",
164                feature = "rustls-platform-verifier",
165                feature = "rustls-bring-your-own-connector"
166            ))]
167            Self::Rustls(s) => Pin::new(s).poll_shutdown(cx),
168        }
169    }
170
171    fn poll_write_vectored(
172        self: Pin<&mut Self>,
173        cx: &mut Context<'_>,
174        bufs: &[io::IoSlice<'_>],
175    ) -> Poll<Result<usize, io::Error>> {
176        match self.get_mut() {
177            Self::Plain(s) => Pin::new(s).poll_write_vectored(cx, bufs),
178            #[cfg(feature = "native-tls")]
179            Self::NativeTls(s) => Pin::new(s).poll_write_vectored(cx, bufs),
180            #[cfg(any(
181                feature = "rustls-native-roots",
182                feature = "rustls-webpki-roots",
183                feature = "rustls-platform-verifier",
184                feature = "rustls-bring-your-own-connector"
185            ))]
186            Self::Rustls(s) => Pin::new(s).poll_write_vectored(cx, bufs),
187        }
188    }
189
190    fn is_write_vectored(&self) -> bool {
191        match self {
192            Self::Plain(s) => s.is_write_vectored(),
193            #[cfg(feature = "native-tls")]
194            Self::NativeTls(s) => s.is_write_vectored(),
195            #[cfg(any(
196                feature = "rustls-native-roots",
197                feature = "rustls-webpki-roots",
198                feature = "rustls-platform-verifier",
199                feature = "rustls-bring-your-own-connector"
200            ))]
201            Self::Rustls(s) => s.is_write_vectored(),
202        }
203    }
204}
205
206impl Connector {
207    /// Creates a new `Connector` with the underlying TLS library specified in
208    /// the feature flags.
209    ///
210    /// # Errors
211    ///
212    /// This method returns an [`Error`] when creating the underlying TLS
213    /// connector fails.
214    ///
215    /// # Panics
216    ///
217    /// Panics if no rustls crypto provider is installed and rustls is unable
218    /// to determine a default crypto provider from its feature flags.
219    pub fn new() -> Result<Self, Error> {
220        #[cfg(not(any(
221            feature = "native-tls",
222            feature = "rustls-webpki-roots",
223            feature = "rustls-native-roots",
224            feature = "rustls-platform-verifier"
225        )))]
226        {
227            Ok(Self::Plain)
228        }
229        #[cfg(all(
230            feature = "native-tls",
231            not(any(
232                feature = "rustls-webpki-roots",
233                feature = "rustls-native-roots",
234                feature = "rustls-platform-verifier"
235            ))
236        ))]
237        {
238            Ok(Self::NativeTls(
239                tokio_native_tls::native_tls::TlsConnector::new()?.into(),
240            ))
241        }
242        #[cfg(any(
243            feature = "rustls-webpki-roots",
244            feature = "rustls-native-roots",
245            feature = "rustls-platform-verifier"
246        ))]
247        {
248            let config_builder = ClientConfig::builder();
249
250            #[cfg(feature = "rustls-platform-verifier")]
251            let config_builder = config_builder.with_platform_verifier()?;
252
253            #[cfg(not(feature = "rustls-platform-verifier"))]
254            let config_builder = {
255                let mut roots = tokio_rustls::rustls::RootCertStore::empty();
256
257                #[cfg(feature = "rustls-native-roots")]
258                {
259                    #[cfg_attr(feature = "rustls-webpki-roots", allow(unused))]
260                    let rustls_native_certs::CertificateResult { certs, errors, .. } =
261                        rustls_native_certs::load_native_certs();
262
263                    // Not finding any native roots is not fatal if webpki roots are enabled
264                    #[cfg(not(feature = "rustls-webpki-roots"))]
265                    if certs.is_empty() {
266                        return Err(Error::NoNativeRootCertificatesFound(errors));
267                    }
268
269                    for cert in certs {
270                        roots.add(cert)?;
271                    }
272                }
273
274                #[cfg(feature = "rustls-webpki-roots")]
275                {
276                    roots.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
277                };
278
279                config_builder.with_root_certificates(roots)
280            };
281
282            Ok(Self::Rustls(tokio_rustls::TlsConnector::from(Arc::new(
283                config_builder.with_no_client_auth(),
284            ))))
285        }
286    }
287
288    /// Wraps a given stream with a layer of TLS.
289    ///
290    /// # Errors
291    ///
292    /// This method returns an [`Error`] if the TLS handshake fails.
293    #[cfg_attr(
294        not(any(
295            feature = "native-tls",
296            feature = "rustls-webpki-roots",
297            feature = "rustls-native-roots",
298            feature = "rustls-platform-verifier",
299            feature = "rustls-bring-your-own-connector"
300        )),
301        allow(unused_variables, clippy::unused_async)
302    )]
303    pub async fn wrap<S: AsyncRead + AsyncWrite + Unpin>(
304        &self,
305        domain: &str,
306        stream: S,
307    ) -> Result<MaybeTlsStream<S>, Error> {
308        match self {
309            Self::Plain => Ok(MaybeTlsStream::Plain(stream)),
310            #[cfg(feature = "native-tls")]
311            Self::NativeTls(connector) => Ok(MaybeTlsStream::NativeTls(
312                connector.connect(domain, stream).await?,
313            )),
314            #[cfg(any(
315                feature = "rustls-webpki-roots",
316                feature = "rustls-native-roots",
317                feature = "rustls-platform-verifier",
318                feature = "rustls-bring-your-own-connector"
319            ))]
320            Self::Rustls(connector) => Ok(MaybeTlsStream::Rustls(
321                connector
322                    .connect(ServerName::try_from(domain)?.to_owned(), stream)
323                    .await?,
324            )),
325        }
326    }
327}
tokio_websockets/tls.rs

tokio_websockets/
tls.rs
