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