tokio_rustls/
client.rs

1use std::future::Future;
2use std::io::{self, BufRead as _};
3#[cfg(unix)]
4use std::os::unix::io::{AsRawFd, RawFd};
5#[cfg(windows)]
6use std::os::windows::io::{AsRawSocket, RawSocket};
7use std::pin::Pin;
8use std::sync::Arc;
9#[cfg(feature = "early-data")]
10use std::task::Waker;
11use std::task::{Context, Poll};
12
13use rustls::pki_types::ServerName;
14use rustls::{ClientConfig, ClientConnection};
15use tokio::io::{AsyncBufRead, AsyncRead, AsyncWrite, ReadBuf};
16
17use crate::common::{IoSession, MidHandshake, Stream, TlsState};
18
19/// A wrapper around a `rustls::ClientConfig`, providing an async `connect` method.
20#[derive(Clone)]
21pub struct TlsConnector {
22    inner: Arc<ClientConfig>,
23    #[cfg(feature = "early-data")]
24    early_data: bool,
25}
26
27impl TlsConnector {
28    /// Enable 0-RTT.
29    ///
30    /// If you want to use 0-RTT,
31    /// You must also set `ClientConfig.enable_early_data` to `true`.
32    #[cfg(feature = "early-data")]
33    pub fn early_data(mut self, flag: bool) -> Self {
34        self.early_data = flag;
35        self
36    }
37
38    /// Returns a future that performs a TLS handshake with `domain` using the `stream`.
39    ///
40    /// You likely want to wrap this in a timeout (for example with [`tokio::time::timeout`][])
41    /// to bound the handshake time.
42    ///
43    /// [`tokio::time::timeout`]: https://docs.rs/tokio/latest/tokio/time/fn.timeout.html
44    #[inline]
45    pub fn connect<IO>(&self, domain: ServerName<'static>, stream: IO) -> Connect<IO>
46    where
47        IO: AsyncRead + AsyncWrite + Unpin,
48    {
49        self.connect_impl(domain, stream, None, |_| ())
50    }
51
52    /// Similar to [`Self::connect()`], but calls `f` before performing the handshake.
53    ///
54    /// As with [`Self::connect()`] you likely want to wrap this in a timeout to
55    /// bound the handshake time.
56    ///
57    /// The `f` handler is given a mutable reference to a [`ClientConnection`][] that can
58    /// be used for tasks like writing early data.
59    ///
60    /// [`ClientConnection`]: https://docs.rs/rustls/latest/rustls/client/struct.ClientConnection.html
61    #[inline]
62    pub fn connect_with<IO, F>(&self, domain: ServerName<'static>, stream: IO, f: F) -> Connect<IO>
63    where
64        IO: AsyncRead + AsyncWrite + Unpin,
65        F: FnOnce(&mut ClientConnection),
66    {
67        self.connect_impl(domain, stream, None, f)
68    }
69
70    fn connect_impl<IO, F>(
71        &self,
72        domain: ServerName<'static>,
73        stream: IO,
74        alpn_protocols: Option<Vec<Vec<u8>>>,
75        f: F,
76    ) -> Connect<IO>
77    where
78        IO: AsyncRead + AsyncWrite + Unpin,
79        F: FnOnce(&mut ClientConnection),
80    {
81        let alpn = alpn_protocols.unwrap_or_else(|| self.inner.alpn_protocols.clone());
82        let mut session = match ClientConnection::new_with_alpn(self.inner.clone(), domain, alpn) {
83            Ok(session) => session,
84            Err(error) => {
85                return Connect(MidHandshake::Error {
86                    io: stream,
87                    // TODO(eliza): should this really return an `io::Error`?
88                    // Probably not...
89                    error: io::Error::new(io::ErrorKind::Other, error),
90                });
91            }
92        };
93        f(&mut session);
94
95        Connect(MidHandshake::Handshaking(TlsStream {
96            io: stream,
97
98            #[cfg(not(feature = "early-data"))]
99            state: TlsState::Stream,
100
101            #[cfg(feature = "early-data")]
102            state: if self.early_data && session.early_data().is_some() {
103                TlsState::EarlyData(0, Vec::new())
104            } else {
105                TlsState::Stream
106            },
107
108            need_flush: false,
109
110            #[cfg(feature = "early-data")]
111            early_waker: None,
112
113            session,
114        }))
115    }
116
117    pub fn with_alpn(&self, alpn_protocols: Vec<Vec<u8>>) -> TlsConnectorWithAlpn<'_> {
118        TlsConnectorWithAlpn {
119            inner: self,
120            alpn_protocols,
121        }
122    }
123
124    /// Get a read-only reference to underlying config
125    pub fn config(&self) -> &Arc<ClientConfig> {
126        &self.inner
127    }
128}
129
130impl From<Arc<ClientConfig>> for TlsConnector {
131    fn from(inner: Arc<ClientConfig>) -> Self {
132        Self {
133            inner,
134            #[cfg(feature = "early-data")]
135            early_data: false,
136        }
137    }
138}
139
140pub struct TlsConnectorWithAlpn<'c> {
141    inner: &'c TlsConnector,
142    alpn_protocols: Vec<Vec<u8>>,
143}
144
145impl TlsConnectorWithAlpn<'_> {
146    /// Returns a future that performs a TLS handshake with `domain` using the `stream`.
147    ///
148    /// You likely want to wrap this in a timeout (for example with [`tokio::time::timeout`][])
149    /// to bound the handshake time.
150    ///
151    /// [`tokio::time::timeout`]: https://docs.rs/tokio/latest/tokio/time/fn.timeout.html
152    #[inline]
153    pub fn connect<IO>(self, domain: ServerName<'static>, stream: IO) -> Connect<IO>
154    where
155        IO: AsyncRead + AsyncWrite + Unpin,
156    {
157        self.inner
158            .connect_impl(domain, stream, Some(self.alpn_protocols), |_| ())
159    }
160
161    /// Similar to [`Self::connect()`], but calls `f` before performing the handshake.
162    ///
163    /// As with [`Self::connect()`] you likely want to wrap this in a timeout to
164    /// bound the handshake time.
165    ///
166    /// The `f` handler is given a mutable reference to a [`ClientConnection`][] that can
167    /// be used for tasks like writing early data.
168    ///
169    /// [`ClientConnection`]: https://docs.rs/rustls/latest/rustls/client/struct.ClientConnection.html
170    #[inline]
171    pub fn connect_with<IO, F>(self, domain: ServerName<'static>, stream: IO, f: F) -> Connect<IO>
172    where
173        IO: AsyncRead + AsyncWrite + Unpin,
174        F: FnOnce(&mut ClientConnection),
175    {
176        self.inner
177            .connect_impl(domain, stream, Some(self.alpn_protocols), f)
178    }
179}
180
181/// Future returned from `TlsConnector::connect` which will resolve
182/// once the connection handshake has finished.
183pub struct Connect<IO>(MidHandshake<TlsStream<IO>>);
184
185impl<IO> Connect<IO> {
186    #[inline]
187    pub fn into_fallible(self) -> FallibleConnect<IO> {
188        FallibleConnect(self.0)
189    }
190
191    pub fn get_ref(&self) -> Option<&IO> {
192        match &self.0 {
193            MidHandshake::Handshaking(sess) => Some(sess.get_ref().0),
194            MidHandshake::SendAlert { io, .. } => Some(io),
195            MidHandshake::Error { io, .. } => Some(io),
196            MidHandshake::End => None,
197        }
198    }
199
200    pub fn get_mut(&mut self) -> Option<&mut IO> {
201        match &mut self.0 {
202            MidHandshake::Handshaking(sess) => Some(sess.get_mut().0),
203            MidHandshake::SendAlert { io, .. } => Some(io),
204            MidHandshake::Error { io, .. } => Some(io),
205            MidHandshake::End => None,
206        }
207    }
208}
209
210impl<IO: AsyncRead + AsyncWrite + Unpin> Future for Connect<IO> {
211    type Output = io::Result<TlsStream<IO>>;
212
213    #[inline]
214    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
215        Pin::new(&mut self.0).poll(cx).map_err(|(err, _)| err)
216    }
217}
218
219impl<IO: AsyncRead + AsyncWrite + Unpin> Future for FallibleConnect<IO> {
220    type Output = Result<TlsStream<IO>, (io::Error, IO)>;
221
222    #[inline]
223    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
224        Pin::new(&mut self.0).poll(cx)
225    }
226}
227
228/// Like [Connect], but returns `IO` on failure.
229pub struct FallibleConnect<IO>(MidHandshake<TlsStream<IO>>);
230
231/// A wrapper around an underlying raw stream which implements the TLS or SSL
232/// protocol.
233#[derive(Debug)]
234pub struct TlsStream<IO> {
235    pub(crate) io: IO,
236    pub(crate) session: ClientConnection,
237    pub(crate) state: TlsState,
238    pub(crate) need_flush: bool,
239
240    #[cfg(feature = "early-data")]
241    pub(crate) early_waker: Option<Waker>,
242}
243
244impl<IO> TlsStream<IO> {
245    #[inline]
246    pub fn get_ref(&self) -> (&IO, &ClientConnection) {
247        (&self.io, &self.session)
248    }
249
250    #[inline]
251    pub fn get_mut(&mut self) -> (&mut IO, &mut ClientConnection) {
252        (&mut self.io, &mut self.session)
253    }
254
255    #[inline]
256    pub fn into_inner(self) -> (IO, ClientConnection) {
257        (self.io, self.session)
258    }
259}
260
261#[cfg(unix)]
262impl<S> AsRawFd for TlsStream<S>
263where
264    S: AsRawFd,
265{
266    fn as_raw_fd(&self) -> RawFd {
267        self.get_ref().0.as_raw_fd()
268    }
269}
270
271#[cfg(windows)]
272impl<S> AsRawSocket for TlsStream<S>
273where
274    S: AsRawSocket,
275{
276    fn as_raw_socket(&self) -> RawSocket {
277        self.get_ref().0.as_raw_socket()
278    }
279}
280
281impl<IO> IoSession for TlsStream<IO> {
282    type Io = IO;
283    type Session = ClientConnection;
284
285    #[inline]
286    fn skip_handshake(&self) -> bool {
287        self.state.is_early_data()
288    }
289
290    #[inline]
291    fn get_mut(&mut self) -> (&mut TlsState, &mut Self::Io, &mut Self::Session, &mut bool) {
292        (
293            &mut self.state,
294            &mut self.io,
295            &mut self.session,
296            &mut self.need_flush,
297        )
298    }
299
300    #[inline]
301    fn into_io(self) -> Self::Io {
302        self.io
303    }
304}
305
306#[cfg(feature = "early-data")]
307impl<IO> TlsStream<IO>
308where
309    IO: AsyncRead + AsyncWrite + Unpin,
310{
311    fn poll_early_data(&mut self, cx: &mut Context<'_>) {
312        // In the EarlyData state, we have not really established a Tls connection.
313        // Before writing data through `AsyncWrite` and completing the tls handshake,
314        // we ignore read readiness and return to pending.
315        //
316        // In order to avoid event loss,
317        // we need to register a waker and wake it up after tls is connected.
318        if self
319            .early_waker
320            .as_ref()
321            .filter(|waker| cx.waker().will_wake(waker))
322            .is_none()
323        {
324            self.early_waker = Some(cx.waker().clone());
325        }
326    }
327}
328
329impl<IO> AsyncRead for TlsStream<IO>
330where
331    IO: AsyncRead + AsyncWrite + Unpin,
332{
333    fn poll_read(
334        mut self: Pin<&mut Self>,
335        cx: &mut Context<'_>,
336        buf: &mut ReadBuf<'_>,
337    ) -> Poll<io::Result<()>> {
338        let data = ready!(self.as_mut().poll_fill_buf(cx))?;
339        let len = data.len().min(buf.remaining());
340        if len == 0 {
341            return Poll::Ready(Ok(()));
342        }
343        buf.put_slice(&data[..len]);
344        self.as_mut().consume(len);
345
346        while buf.remaining() > 0 {
347            let data = match self.as_mut().poll_fill_buf(cx) {
348                Poll::Ready(Ok([])) => break,
349                Poll::Ready(Ok(data)) => data,
350                Poll::Ready(Err(_)) => break, // non-transient error gets re-emitted next poll
351                Poll::Pending => break,
352            };
353            let len = Ord::min(data.len(), buf.remaining());
354            buf.put_slice(&data[..len]);
355            self.as_mut().consume(len);
356        }
357        Poll::Ready(Ok(()))
358    }
359}
360
361impl<IO> AsyncBufRead for TlsStream<IO>
362where
363    IO: AsyncRead + AsyncWrite + Unpin,
364{
365    fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
366        match self.state {
367            #[cfg(feature = "early-data")]
368            TlsState::EarlyData(..) => {
369                self.get_mut().poll_early_data(cx);
370                Poll::Pending
371            }
372            TlsState::Stream | TlsState::WriteShutdown => {
373                let this = self.get_mut();
374                let stream =
375                    Stream::new(&mut this.io, &mut this.session).set_eof(!this.state.readable());
376
377                match stream.poll_fill_buf(cx) {
378                    Poll::Ready(Ok(buf)) => {
379                        if buf.is_empty() {
380                            this.state.shutdown_read();
381                        }
382
383                        Poll::Ready(Ok(buf))
384                    }
385                    Poll::Ready(Err(err)) if err.kind() == io::ErrorKind::ConnectionAborted => {
386                        this.state.shutdown_read();
387                        Poll::Ready(Err(err))
388                    }
389                    output => output,
390                }
391            }
392            TlsState::ReadShutdown | TlsState::FullyShutdown => Poll::Ready(Ok(&[])),
393        }
394    }
395
396    fn consume(mut self: Pin<&mut Self>, amt: usize) {
397        self.session.reader().consume(amt);
398    }
399}
400
401impl<IO> AsyncWrite for TlsStream<IO>
402where
403    IO: AsyncRead + AsyncWrite + Unpin,
404{
405    /// Note: that it does not guarantee the final data to be sent.
406    /// To be cautious, you must manually call `flush`.
407    fn poll_write(
408        self: Pin<&mut Self>,
409        cx: &mut Context<'_>,
410        buf: &[u8],
411    ) -> Poll<io::Result<usize>> {
412        let this = self.get_mut();
413        let mut stream = Stream::new(&mut this.io, &mut this.session)
414            .set_eof(!this.state.readable())
415            .set_need_flush(this.need_flush);
416
417        #[cfg(feature = "early-data")]
418        {
419            let bufs = [io::IoSlice::new(buf)];
420            let written = poll_handle_early_data(
421                &mut this.state,
422                &mut stream,
423                &mut this.early_waker,
424                cx,
425                &bufs,
426            )?;
427            match written {
428                Poll::Ready(0) => {}
429                Poll::Ready(written) => return Poll::Ready(Ok(written)),
430                Poll::Pending => {
431                    this.need_flush = stream.need_flush;
432                    return Poll::Pending;
433                }
434            }
435        }
436
437        stream.as_mut_pin().poll_write(cx, buf)
438    }
439
440    /// Note: that it does not guarantee the final data to be sent.
441    /// To be cautious, you must manually call `flush`.
442    fn poll_write_vectored(
443        self: Pin<&mut Self>,
444        cx: &mut Context<'_>,
445        bufs: &[io::IoSlice<'_>],
446    ) -> Poll<io::Result<usize>> {
447        let this = self.get_mut();
448        let mut stream = Stream::new(&mut this.io, &mut this.session)
449            .set_eof(!this.state.readable())
450            .set_need_flush(this.need_flush);
451
452        #[cfg(feature = "early-data")]
453        {
454            let written = poll_handle_early_data(
455                &mut this.state,
456                &mut stream,
457                &mut this.early_waker,
458                cx,
459                bufs,
460            )?;
461            match written {
462                Poll::Ready(0) => {}
463                Poll::Ready(written) => return Poll::Ready(Ok(written)),
464                Poll::Pending => {
465                    this.need_flush = stream.need_flush;
466                    return Poll::Pending;
467                }
468            }
469        }
470
471        stream.as_mut_pin().poll_write_vectored(cx, bufs)
472    }
473
474    #[inline]
475    fn is_write_vectored(&self) -> bool {
476        true
477    }
478
479    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
480        let this = self.get_mut();
481        let mut stream = Stream::new(&mut this.io, &mut this.session)
482            .set_eof(!this.state.readable())
483            .set_need_flush(this.need_flush);
484
485        #[cfg(feature = "early-data")]
486        {
487            let written = poll_handle_early_data(
488                &mut this.state,
489                &mut stream,
490                &mut this.early_waker,
491                cx,
492                &[],
493            )?;
494            if written.is_pending() {
495                this.need_flush = stream.need_flush;
496                return Poll::Pending;
497            }
498        }
499
500        stream.as_mut_pin().poll_flush(cx)
501    }
502
503    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
504        #[cfg(feature = "early-data")]
505        {
506            // complete handshake
507            if matches!(self.state, TlsState::EarlyData(..)) {
508                ready!(self.as_mut().poll_flush(cx))?;
509            }
510        }
511
512        if self.state.writeable() {
513            self.session.send_close_notify();
514            self.state.shutdown_write();
515        }
516
517        let this = self.get_mut();
518        let mut stream =
519            Stream::new(&mut this.io, &mut this.session).set_eof(!this.state.readable());
520        stream.as_mut_pin().poll_shutdown(cx)
521    }
522}
523
524#[cfg(feature = "early-data")]
525fn poll_handle_early_data<IO>(
526    state: &mut TlsState,
527    stream: &mut Stream<IO, ClientConnection>,
528    early_waker: &mut Option<Waker>,
529    cx: &mut Context<'_>,
530    bufs: &[io::IoSlice<'_>],
531) -> Poll<io::Result<usize>>
532where
533    IO: AsyncRead + AsyncWrite + Unpin,
534{
535    if let TlsState::EarlyData(pos, data) = state {
536        use std::io::Write;
537
538        // write early data
539        if let Some(mut early_data) = stream.session.early_data() {
540            let mut written = 0;
541
542            for buf in bufs {
543                if buf.is_empty() {
544                    continue;
545                }
546
547                let len = match early_data.write(buf) {
548                    Ok(0) => break,
549                    Ok(n) => n,
550                    Err(err) => return Poll::Ready(Err(err)),
551                };
552
553                written += len;
554                data.extend_from_slice(&buf[..len]);
555
556                if len < buf.len() {
557                    break;
558                }
559            }
560
561            if written != 0 {
562                return Poll::Ready(Ok(written));
563            }
564        }
565
566        // complete handshake
567        while stream.session.is_handshaking() {
568            ready!(stream.handshake(cx))?;
569        }
570
571        // write early data (fallback)
572        if !stream.session.is_early_data_accepted() {
573            while *pos < data.len() {
574                let len = ready!(stream.as_mut_pin().poll_write(cx, &data[*pos..]))?;
575                *pos += len;
576            }
577        }
578
579        // end
580        *state = TlsState::Stream;
581
582        if let Some(waker) = early_waker.take() {
583            waker.wake();
584        }
585    }
586
587    Poll::Ready(Ok(0))
588}
tokio_rustls/client.rs

tokio_rustls/
client.rs
