async_nats_tokio_rustls_deps/
lib.rs

1//! This is a fork of [tokio-ruslts](https://crates.io/crates/tokio-rustls) created out of need
2//! for features merged into main, but not yet released.
3//! It can be treated as alpha/nightly.
4//!
5//! Asynchronous TLS/SSL streams for Tokio using [Rustls](https://github.com/rustls/rustls).
6//!
7//! # Why do I need to call `poll_flush`?
8//!
9//! Most TLS implementations will have an internal buffer to improve throughput,
10//! and rustls is no exception.
11//!
12//! When we write data to `TlsStream`, we always write rustls buffer first,
13//! then take out rustls encrypted data packet, and write it to data channel (like TcpStream).
14//! When data channel is pending, some data may remain in rustls buffer.
15//!
16//! `tokio-rustls` To keep it simple and correct, [TlsStream] will behave like `BufWriter`.
17//! For `TlsStream<TcpStream>`, this means that data written by `poll_write` is not guaranteed to be written to `TcpStream`.
18//! You must call `poll_flush` to ensure that it is written to `TcpStream`.
19//!
20//! You should call `poll_flush` at the appropriate time,
21//! such as when a period of `poll_write` write is complete and there is no more data to write.
22//!
23//! ## Why don't we write during `poll_read`?
24//!
25//! We did this in the early days of `tokio-rustls`, but it caused some bugs.
26//! We can solve these bugs through some solutions, but this will cause performance degradation (reverse false wakeup).
27//!
28//! And reverse write will also prevent us implement full duplex in the future.
29//!
30//! see <https://github.com/tokio-rs/tls/issues/40>
31//!
32//! ## Why can't we handle it like `native-tls`?
33//!
34//! When data channel returns to pending, `native-tls` will falsely report the number of bytes it consumes.
35//! This means that if data written by `poll_write` is not actually written to data channel, it will not return `Ready`.
36//! Thus avoiding the call of `poll_flush`.
37//!
38//! but which does not conform to convention of `AsyncWrite` trait.
39//! This means that if you give inconsistent data in two `poll_write`, it may cause unexpected behavior.
40//!
41//! see <https://github.com/tokio-rs/tls/issues/41>
42
43macro_rules! ready {
44    ( $e:expr ) => {
45        match $e {
46            std::task::Poll::Ready(t) => t,
47            std::task::Poll::Pending => return std::task::Poll::Pending,
48        }
49    };
50}
51
52pub mod client;
53mod common;
54pub mod server;
55
56use common::{MidHandshake, Stream, TlsState};
57use rustls::{ClientConfig, ClientConnection, CommonState, ServerConfig, ServerConnection};
58use std::future::Future;
59use std::io;
60#[cfg(unix)]
61use std::os::unix::io::{AsRawFd, RawFd};
62#[cfg(windows)]
63use std::os::windows::io::{AsRawSocket, RawSocket};
64use std::pin::Pin;
65use std::sync::Arc;
66use std::task::{Context, Poll};
67use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
68
69pub use rustls;
70pub use webpki;
71
72/// A wrapper around a `rustls::ClientConfig`, providing an async `connect` method.
73#[derive(Clone)]
74pub struct TlsConnector {
75    inner: Arc<ClientConfig>,
76    #[cfg(feature = "early-data")]
77    early_data: bool,
78}
79
80/// A wrapper around a `rustls::ServerConfig`, providing an async `accept` method.
81#[derive(Clone)]
82pub struct TlsAcceptor {
83    inner: Arc<ServerConfig>,
84}
85
86impl From<Arc<ClientConfig>> for TlsConnector {
87    fn from(inner: Arc<ClientConfig>) -> TlsConnector {
88        TlsConnector {
89            inner,
90            #[cfg(feature = "early-data")]
91            early_data: false,
92        }
93    }
94}
95
96impl From<Arc<ServerConfig>> for TlsAcceptor {
97    fn from(inner: Arc<ServerConfig>) -> TlsAcceptor {
98        TlsAcceptor { inner }
99    }
100}
101
102impl TlsConnector {
103    /// Enable 0-RTT.
104    ///
105    /// If you want to use 0-RTT,
106    /// You must also set `ClientConfig.enable_early_data` to `true`.
107    #[cfg(feature = "early-data")]
108    pub fn early_data(mut self, flag: bool) -> TlsConnector {
109        self.early_data = flag;
110        self
111    }
112
113    #[inline]
114    pub fn connect<IO>(&self, domain: rustls::ServerName, stream: IO) -> Connect<IO>
115    where
116        IO: AsyncRead + AsyncWrite + Unpin,
117    {
118        self.connect_with(domain, stream, |_| ())
119    }
120
121    pub fn connect_with<IO, F>(&self, domain: rustls::ServerName, stream: IO, f: F) -> Connect<IO>
122    where
123        IO: AsyncRead + AsyncWrite + Unpin,
124        F: FnOnce(&mut ClientConnection),
125    {
126        let mut session = match ClientConnection::new(self.inner.clone(), domain) {
127            Ok(session) => session,
128            Err(error) => {
129                return Connect(MidHandshake::Error {
130                    io: stream,
131                    // TODO(eliza): should this really return an `io::Error`?
132                    // Probably not...
133                    error: io::Error::new(io::ErrorKind::Other, error),
134                });
135            }
136        };
137        f(&mut session);
138
139        Connect(MidHandshake::Handshaking(client::TlsStream {
140            io: stream,
141
142            #[cfg(not(feature = "early-data"))]
143            state: TlsState::Stream,
144
145            #[cfg(feature = "early-data")]
146            state: if self.early_data && session.early_data().is_some() {
147                TlsState::EarlyData(0, Vec::new())
148            } else {
149                TlsState::Stream
150            },
151
152            #[cfg(feature = "early-data")]
153            early_waker: None,
154
155            session,
156        }))
157    }
158}
159
160impl TlsAcceptor {
161    #[inline]
162    pub fn accept<IO>(&self, stream: IO) -> Accept<IO>
163    where
164        IO: AsyncRead + AsyncWrite + Unpin,
165    {
166        self.accept_with(stream, |_| ())
167    }
168
169    pub fn accept_with<IO, F>(&self, stream: IO, f: F) -> Accept<IO>
170    where
171        IO: AsyncRead + AsyncWrite + Unpin,
172        F: FnOnce(&mut ServerConnection),
173    {
174        let mut session = match ServerConnection::new(self.inner.clone()) {
175            Ok(session) => session,
176            Err(error) => {
177                return Accept(MidHandshake::Error {
178                    io: stream,
179                    // TODO(eliza): should this really return an `io::Error`?
180                    // Probably not...
181                    error: io::Error::new(io::ErrorKind::Other, error),
182                });
183            }
184        };
185        f(&mut session);
186
187        Accept(MidHandshake::Handshaking(server::TlsStream {
188            session,
189            io: stream,
190            state: TlsState::Stream,
191        }))
192    }
193}
194
195pub struct LazyConfigAcceptor<IO> {
196    acceptor: rustls::server::Acceptor,
197    io: Option<IO>,
198}
199
200impl<IO> LazyConfigAcceptor<IO>
201where
202    IO: AsyncRead + AsyncWrite + Unpin,
203{
204    #[inline]
205    pub fn new(acceptor: rustls::server::Acceptor, io: IO) -> Self {
206        Self {
207            acceptor,
208            io: Some(io),
209        }
210    }
211}
212
213impl<IO> Future for LazyConfigAcceptor<IO>
214where
215    IO: AsyncRead + AsyncWrite + Unpin,
216{
217    type Output = Result<StartHandshake<IO>, io::Error>;
218
219    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
220        let this = self.get_mut();
221        loop {
222            let io = match this.io.as_mut() {
223                Some(io) => io,
224                None => {
225                    return Poll::Ready(Err(io::Error::new(
226                        io::ErrorKind::Other,
227                        "acceptor cannot be polled after acceptance",
228                    )))
229                }
230            };
231
232            let mut reader = common::SyncReadAdapter { io, cx };
233            match this.acceptor.read_tls(&mut reader) {
234                Ok(0) => return Err(io::ErrorKind::UnexpectedEof.into()).into(),
235                Ok(_) => {}
236                Err(e) if e.kind() == io::ErrorKind::WouldBlock => return Poll::Pending,
237                Err(e) => return Err(e).into(),
238            }
239
240            match this.acceptor.accept() {
241                Ok(Some(accepted)) => {
242                    let io = this.io.take().unwrap();
243                    return Poll::Ready(Ok(StartHandshake { accepted, io }));
244                }
245                Ok(None) => continue,
246                Err(err) => {
247                    return Poll::Ready(Err(io::Error::new(io::ErrorKind::InvalidInput, err)))
248                }
249            }
250        }
251    }
252}
253
254pub struct StartHandshake<IO> {
255    accepted: rustls::server::Accepted,
256    io: IO,
257}
258
259impl<IO> StartHandshake<IO>
260where
261    IO: AsyncRead + AsyncWrite + Unpin,
262{
263    pub fn client_hello(&self) -> rustls::server::ClientHello<'_> {
264        self.accepted.client_hello()
265    }
266
267    pub fn into_stream(self, config: Arc<ServerConfig>) -> Accept<IO> {
268        self.into_stream_with(config, |_| ())
269    }
270
271    pub fn into_stream_with<F>(self, config: Arc<ServerConfig>, f: F) -> Accept<IO>
272    where
273        F: FnOnce(&mut ServerConnection),
274    {
275        let mut conn = match self.accepted.into_connection(config) {
276            Ok(conn) => conn,
277            Err(error) => {
278                return Accept(MidHandshake::Error {
279                    io: self.io,
280                    // TODO(eliza): should this really return an `io::Error`?
281                    // Probably not...
282                    error: io::Error::new(io::ErrorKind::Other, error),
283                });
284            }
285        };
286        f(&mut conn);
287
288        Accept(MidHandshake::Handshaking(server::TlsStream {
289            session: conn,
290            io: self.io,
291            state: TlsState::Stream,
292        }))
293    }
294}
295
296/// Future returned from `TlsConnector::connect` which will resolve
297/// once the connection handshake has finished.
298pub struct Connect<IO>(MidHandshake<client::TlsStream<IO>>);
299
300/// Future returned from `TlsAcceptor::accept` which will resolve
301/// once the accept handshake has finished.
302pub struct Accept<IO>(MidHandshake<server::TlsStream<IO>>);
303
304/// Like [Connect], but returns `IO` on failure.
305pub struct FallibleConnect<IO>(MidHandshake<client::TlsStream<IO>>);
306
307/// Like [Accept], but returns `IO` on failure.
308pub struct FallibleAccept<IO>(MidHandshake<server::TlsStream<IO>>);
309
310impl<IO> Connect<IO> {
311    #[inline]
312    pub fn into_fallible(self) -> FallibleConnect<IO> {
313        FallibleConnect(self.0)
314    }
315
316    pub fn get_ref(&self) -> Option<&IO> {
317        match &self.0 {
318            MidHandshake::Handshaking(sess) => Some(sess.get_ref().0),
319            MidHandshake::Error { io, .. } => Some(io),
320            MidHandshake::End => None,
321        }
322    }
323
324    pub fn get_mut(&mut self) -> Option<&mut IO> {
325        match &mut self.0 {
326            MidHandshake::Handshaking(sess) => Some(sess.get_mut().0),
327            MidHandshake::Error { io, .. } => Some(io),
328            MidHandshake::End => None,
329        }
330    }
331}
332
333impl<IO> Accept<IO> {
334    #[inline]
335    pub fn into_fallible(self) -> FallibleAccept<IO> {
336        FallibleAccept(self.0)
337    }
338
339    pub fn get_ref(&self) -> Option<&IO> {
340        match &self.0 {
341            MidHandshake::Handshaking(sess) => Some(sess.get_ref().0),
342            MidHandshake::Error { io, .. } => Some(io),
343            MidHandshake::End => None,
344        }
345    }
346
347    pub fn get_mut(&mut self) -> Option<&mut IO> {
348        match &mut self.0 {
349            MidHandshake::Handshaking(sess) => Some(sess.get_mut().0),
350            MidHandshake::Error { io, .. } => Some(io),
351            MidHandshake::End => None,
352        }
353    }
354}
355
356impl<IO: AsyncRead + AsyncWrite + Unpin> Future for Connect<IO> {
357    type Output = io::Result<client::TlsStream<IO>>;
358
359    #[inline]
360    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
361        Pin::new(&mut self.0).poll(cx).map_err(|(err, _)| err)
362    }
363}
364
365impl<IO: AsyncRead + AsyncWrite + Unpin> Future for Accept<IO> {
366    type Output = io::Result<server::TlsStream<IO>>;
367
368    #[inline]
369    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
370        Pin::new(&mut self.0).poll(cx).map_err(|(err, _)| err)
371    }
372}
373
374impl<IO: AsyncRead + AsyncWrite + Unpin> Future for FallibleConnect<IO> {
375    type Output = Result<client::TlsStream<IO>, (io::Error, IO)>;
376
377    #[inline]
378    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
379        Pin::new(&mut self.0).poll(cx)
380    }
381}
382
383impl<IO: AsyncRead + AsyncWrite + Unpin> Future for FallibleAccept<IO> {
384    type Output = Result<server::TlsStream<IO>, (io::Error, IO)>;
385
386    #[inline]
387    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
388        Pin::new(&mut self.0).poll(cx)
389    }
390}
391
392/// Unified TLS stream type
393///
394/// This abstracts over the inner `client::TlsStream` and `server::TlsStream`, so you can use
395/// a single type to keep both client- and server-initiated TLS-encrypted connections.
396#[allow(clippy::large_enum_variant)] // https://github.com/rust-lang/rust-clippy/issues/9798
397#[derive(Debug)]
398pub enum TlsStream<T> {
399    Client(client::TlsStream<T>),
400    Server(server::TlsStream<T>),
401}
402
403impl<T> TlsStream<T> {
404    pub fn get_ref(&self) -> (&T, &CommonState) {
405        use TlsStream::*;
406        match self {
407            Client(io) => {
408                let (io, session) = io.get_ref();
409                (io, session)
410            }
411            Server(io) => {
412                let (io, session) = io.get_ref();
413                (io, session)
414            }
415        }
416    }
417
418    pub fn get_mut(&mut self) -> (&mut T, &mut CommonState) {
419        use TlsStream::*;
420        match self {
421            Client(io) => {
422                let (io, session) = io.get_mut();
423                (io, &mut *session)
424            }
425            Server(io) => {
426                let (io, session) = io.get_mut();
427                (io, &mut *session)
428            }
429        }
430    }
431}
432
433impl<T> From<client::TlsStream<T>> for TlsStream<T> {
434    fn from(s: client::TlsStream<T>) -> Self {
435        Self::Client(s)
436    }
437}
438
439impl<T> From<server::TlsStream<T>> for TlsStream<T> {
440    fn from(s: server::TlsStream<T>) -> Self {
441        Self::Server(s)
442    }
443}
444
445#[cfg(unix)]
446impl<S> AsRawFd for TlsStream<S>
447where
448    S: AsRawFd,
449{
450    fn as_raw_fd(&self) -> RawFd {
451        self.get_ref().0.as_raw_fd()
452    }
453}
454
455#[cfg(windows)]
456impl<S> AsRawSocket for TlsStream<S>
457where
458    S: AsRawSocket,
459{
460    fn as_raw_socket(&self) -> RawSocket {
461        self.get_ref().0.as_raw_socket()
462    }
463}
464
465impl<T> AsyncRead for TlsStream<T>
466where
467    T: AsyncRead + AsyncWrite + Unpin,
468{
469    #[inline]
470    fn poll_read(
471        self: Pin<&mut Self>,
472        cx: &mut Context<'_>,
473        buf: &mut ReadBuf<'_>,
474    ) -> Poll<io::Result<()>> {
475        match self.get_mut() {
476            TlsStream::Client(x) => Pin::new(x).poll_read(cx, buf),
477            TlsStream::Server(x) => Pin::new(x).poll_read(cx, buf),
478        }
479    }
480}
481
482impl<T> AsyncWrite for TlsStream<T>
483where
484    T: AsyncRead + AsyncWrite + Unpin,
485{
486    #[inline]
487    fn poll_write(
488        self: Pin<&mut Self>,
489        cx: &mut Context<'_>,
490        buf: &[u8],
491    ) -> Poll<io::Result<usize>> {
492        match self.get_mut() {
493            TlsStream::Client(x) => Pin::new(x).poll_write(cx, buf),
494            TlsStream::Server(x) => Pin::new(x).poll_write(cx, buf),
495        }
496    }
497
498    #[inline]
499    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
500        match self.get_mut() {
501            TlsStream::Client(x) => Pin::new(x).poll_flush(cx),
502            TlsStream::Server(x) => Pin::new(x).poll_flush(cx),
503        }
504    }
505
506    #[inline]
507    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
508        match self.get_mut() {
509            TlsStream::Client(x) => Pin::new(x).poll_shutdown(cx),
510            TlsStream::Server(x) => Pin::new(x).poll_shutdown(cx),
511        }
512    }
513}
async_nats_tokio_rustls_deps/lib.rs

async_nats_tokio_rustls_deps/
lib.rs
