From 216d8bfa0538e5023c85c4d62a386780ee200395 Mon Sep 17 00:00:00 2001 From: iadev09 <166385752+iadev09@users.noreply.github.com> Date: Sat, 29 Aug 2026 00:23:39 +0300 Subject: [PATCH] Port tokio-rustls to rustls 0.24 --- .github/workflows/CI.yml | 4 +- Cargo.lock | 127 ++++++++----- Cargo.toml | 19 +- examples/client.rs | 23 ++- examples/server.rs | 34 +++- src/client.rs | 157 ++++++++++++---- src/common/handshake.rs | 125 +++++++++---- src/common/mod.rs | 381 ++++++++++++++++++++++++++++---------- src/common/test_stream.rs | 224 ++++++++++++++-------- src/lib.rs | 17 +- src/server.rs | 264 +++++++++++++++++++------- tests/badssl.rs | 17 +- tests/early-data.rs | 147 ++++++++++++--- tests/test.rs | 43 ++++- tests/utils.rs | 38 +++- 15 files changed, 1199 insertions(+), 421 deletions(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index f93217ea..01808701 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -133,7 +133,7 @@ jobs: - uses: dtolnay/rust-toolchain@master with: - toolchain: "1.71" + toolchain: "1.85" - run: cargo check --locked --lib --features "aws_lc_rs, aws-lc-rs, early-data, fips, logging, ring, tls12" @@ -148,7 +148,7 @@ jobs: - uses: dtolnay/rust-toolchain@master with: - toolchain: "1.75" + toolchain: "1.85" - run: cargo check --locked --lib --all-features diff --git a/Cargo.lock b/Cargo.lock index 0607cf6c..ae689734 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1,6 +1,6 @@ # This file is automatically @generated by Cargo. # It is not intended for manual editing. -version = 3 +version = 4 [[package]] name = "aho-corasick" @@ -104,23 +104,24 @@ checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" [[package]] name = "aws-lc-fips-sys" -version = "0.13.14" +version = "0.14.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d3d619165468401dec3caa3366ebffbcb83f2f31883e5b3932f8e2dec2ddc568" +checksum = "118303cd75f63d1933a90c2ceb7e697281ac6acbdbcc490b46419f25a527ab90" dependencies = [ "bindgen", "cc", "cmake", "dunce", "fs_extra", + "pkg-config", "regex", ] [[package]] name = "aws-lc-rs" -version = "1.17.0" +version = "1.18.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5ec2f1fc3ec205783a5da9a7e6c1509cc69dedf09a1949e412c1e18469326d00" +checksum = "ce2b2dcc879c3bae0d371e77c99f2238400ef24ec001394befa67b6e543add9e" dependencies = [ "aws-lc-fips-sys", "aws-lc-sys", @@ -129,14 +130,15 @@ dependencies = [ [[package]] name = "aws-lc-sys" -version = "0.41.0" +version = "0.44.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1a2f9779ce85b93ab6170dd940ad0169b5766ff848247aff13bb788b832fe3f4" +checksum = "f09fae7be8bb3174e05c6afdb34199e6dc0c7c04ba9fa237b1967adfbde27483" dependencies = [ "cc", "cmake", "dunce", "fs_extra", + "pkg-config", ] [[package]] @@ -165,15 +167,6 @@ dependencies = [ "syn", ] -[[package]] -name = "bit-vec" -version = "0.9.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b71798fca2c1fe1086445a7258a4bc81e6e49dcd24c8d0dd9a1e57395b603f51" -dependencies = [ - "serde", -] - [[package]] name = "bitflags" version = "2.13.0" @@ -276,9 +269,9 @@ dependencies = [ [[package]] name = "deranged" -version = "0.4.0" +version = "0.5.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9c9e6a11ca8224451684bc0d7d5a7adbf8f2fd6887261a1cfc3c0432f9d4068e" +checksum = "7cd812cc2bc1d69d4764bd80df88b4317eaef9e773c75226407d9bc0876b211c" dependencies = [ "powerfmt", ] @@ -576,6 +569,12 @@ version = "0.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" +[[package]] +name = "pkg-config" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f6b464fbc74e149a392436b17d523f769e057cb6877f6a5c4618bc6f11800548" + [[package]] name = "powerfmt" version = "0.2.0" @@ -618,9 +617,9 @@ checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" [[package]] name = "rcgen" -version = "0.14.8" +version = "0.14.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "57f6d249aad744e274e682777a50283a225a32705394ee6d5fcc01efa25e4055" +checksum = "10b99e0098aa4082912d4c649628623db6aba77335e4f4569ff5083a6448b32e" dependencies = [ "pem", "ring", @@ -699,40 +698,66 @@ dependencies = [ [[package]] name = "rustls" -version = "0.23.43" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0283386ce02abc0151e1761d08802dfe86c173b0b494af5cbc086574e453da06" +version = "0.24.0-dev.1" +source = "git+https://github.com/rustls/rustls.git?branch=main#3925f65934364edafe8d6b20707d9e5e6183648e" dependencies = [ - "aws-lc-rs", "brotli", "brotli-decompressor", - "log", "once_cell", - "ring", "rustls-pki-types", "rustls-webpki", "subtle", + "tracing", "zeroize", "zlib-rs", ] +[[package]] +name = "rustls-aws-lc-rs" +version = "0.1.0-dev.1" +source = "git+https://github.com/rustls/rustls.git?branch=main#3925f65934364edafe8d6b20707d9e5e6183648e" +dependencies = [ + "aws-lc-rs", + "rustls", + "rustls-pki-types", + "subtle", + "zeroize", +] + [[package]] name = "rustls-pki-types" -version = "1.14.1" +version = "1.15.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "30a7197ae7eb376e574fe940d068c30fe0462554a3ddbe4eca7838e049c937a9" +checksum = "2f4925028c7eb5d1fcdaf196971378ed9d2c1c4efc7dc5d011256f76c99c0a96" dependencies = [ "zeroize", ] +[[package]] +name = "rustls-ring" +version = "0.1.0-dev.1" +source = "git+https://github.com/rustls/rustls.git?branch=main#3925f65934364edafe8d6b20707d9e5e6183648e" +dependencies = [ + "ring", + "rustls", + "rustls-pki-types", + "subtle", +] + +[[package]] +name = "rustls-util" +version = "0.1.0-dev.1" +source = "git+https://github.com/rustls/rustls.git?branch=main#3925f65934364edafe8d6b20707d9e5e6183648e" +dependencies = [ + "rustls", +] + [[package]] name = "rustls-webpki" -version = "0.103.13" +version = "0.104.0-alpha.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "61c429a8649f110dddef65e2a5ad240f747e85f7758a6bccc7e5777bd33f756e" +checksum = "bea702cca24d344fc70973022bf7eb920c224e318466eb49784272337dd24b1a" dependencies = [ - "aws-lc-rs", - "ring", "rustls-pki-types", "untrusted", ] @@ -867,9 +892,9 @@ dependencies = [ [[package]] name = "time" -version = "0.3.41" +version = "0.3.44" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8a7619e19bc266e0f9c5e6686659d394bc57973859340060a69221e57dbc0c40" +checksum = "91e7d9e3bb61134e77bde20dd4825b97c010155709965fedf0f49bb138e52a9d" dependencies = [ "deranged", "itoa", @@ -882,15 +907,15 @@ dependencies = [ [[package]] name = "time-core" -version = "0.1.4" +version = "0.1.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c9e9a38711f559d9e3ce1cdb06dd7c5b8ea546bc90052da6d06bb76da74bb07c" +checksum = "40868e7c1d2f0b8d73e4a8c7f0ff63af4f6d19be117e90bd73eb1d62cf831c6b" [[package]] name = "time-macros" -version = "0.2.22" +version = "0.2.24" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3526739392ec93fd8b359c8e98514cb3e8e021beb4e5f597b00a0221f8ed8a49" +checksum = "30cfb0125f12d9c277f35663a0a33f8c30190f4e4574868a330595412d34ebf3" dependencies = [ "num-conv", "time-core", @@ -926,17 +951,36 @@ dependencies = [ [[package]] name = "tokio-rustls" -version = "0.26.4" +version = "0.27.0" dependencies = [ "argh", "futures-util", "lazy_static", "rcgen", "rustls", + "rustls-aws-lc-rs", + "rustls-ring", + "rustls-util", "tokio", "webpki-roots", ] +[[package]] +name = "tracing" +version = "0.1.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" +dependencies = [ + "pin-project-lite", + "tracing-core", +] + +[[package]] +name = "tracing-core" +version = "0.1.36" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" + [[package]] name = "unicode-ident" version = "1.0.24" @@ -1087,11 +1131,10 @@ dependencies = [ [[package]] name = "yasna" -version = "0.6.0" +version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b5f6765e852b9b4dc8e2a76843e4d64d1cea8e79bcde0b6901aea8e7c7f08282" +checksum = "e17bb3549cc1321ae1296b9cdc2698e2b6cb1992adfa19a8c72e5b7a738f44cd" dependencies = [ - "bit-vec", "time", ] diff --git a/Cargo.toml b/Cargo.toml index d917a5cb..f420f230 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "tokio-rustls" -version = "0.26.4" +version = "0.27.0" license = "MIT OR Apache-2.0" repository = "https://github.com/rustls/tokio-rustls" homepage = "https://github.com/rustls/tokio-rustls" @@ -9,23 +9,25 @@ readme = "README.md" description = "Asynchronous TLS/SSL streams for Tokio using Rustls." categories = ["asynchronous", "cryptography", "network-programming"] edition = "2021" -rust-version = "1.71" +rust-version = "1.85" exclude = ["/.github", "/examples", "/scripts", "/tests/"] [dependencies] -rustls = { version = "0.23.27", default-features = false, features = ["std"] } +rustls = { git = "https://github.com/rustls/rustls.git", branch = "main", package = "rustls", version = "0.24.0-dev.1", default-features = false, features = ["webpki"] } +rustls-aws-lc-rs = { git = "https://github.com/rustls/rustls.git", branch = "main", package = "rustls-aws-lc-rs", version = "0.1.0-dev.1", default-features = false, features = ["aws-lc-sys", "std"], optional = true } +rustls-ring = { git = "https://github.com/rustls/rustls.git", branch = "main", package = "rustls-ring", version = "0.1.0-dev.1", default-features = false, features = ["std"], optional = true } tokio = "1.0" [features] default = ["logging", "tls12", "aws_lc_rs"] -aws_lc_rs = ["rustls/aws_lc_rs"] +aws_lc_rs = ["dep:rustls-aws-lc-rs"] aws-lc-rs = ["aws_lc_rs"] # Alias because Cargo features commonly use `-` brotli = ["rustls/brotli"] early-data = [] -fips = ["rustls/fips"] -logging = ["rustls/logging"] -ring = ["rustls/ring"] -tls12 = ["rustls/tls12"] +fips = ["aws_lc_rs", "rustls-aws-lc-rs/fips"] +logging = ["rustls/tracing"] +ring = ["dep:rustls-ring"] +tls12 = [] zlib = ["rustls/zlib"] [dev-dependencies] @@ -33,6 +35,7 @@ argh = "0.1.1" futures-util = "0.3.1" lazy_static = "1.1" rcgen = { version = "0.14", features = ["pem"] } +rustls-util = { git = "https://github.com/rustls/rustls.git", branch = "main", package = "rustls-util", version = "0.1.0-dev.1", default-features = false } tokio = { version = "1.0", features = ["full"] } webpki-roots = "1" diff --git a/examples/client.rs b/examples/client.rs index 4c2c0ca4..2ea0013a 100644 --- a/examples/client.rs +++ b/examples/client.rs @@ -26,7 +26,7 @@ use tokio::time::timeout; use tokio_rustls::TlsConnector; use tokio_rustls::rustls::pki_types::pem::PemObject; use tokio_rustls::rustls::pki_types::{CertificateDer, ServerName}; -use tokio_rustls::rustls::{ClientConfig, RootCertStore}; +use tokio_rustls::rustls::{self, ClientConfig, RootCertStore}; #[tokio::main] async fn main() -> Result<(), Box> { @@ -49,9 +49,9 @@ async fn main() -> Result<(), Box> { } let connector = TlsConnector::from(Arc::new( - ClientConfig::builder() + ClientConfig::builder(provider()) .with_root_certificates(root_cert_store) - .with_no_client_auth(), + .with_no_client_auth()?, )); // TCP connect, bounded by the connect timeout. @@ -163,3 +163,20 @@ async fn with_timeout( )), } } + +fn provider() -> Arc { + #[cfg(feature = "aws_lc_rs")] + { + Arc::new(rustls_aws_lc_rs::DEFAULT_PROVIDER.clone()) + } + + #[cfg(all(not(feature = "aws_lc_rs"), feature = "ring"))] + { + Arc::new(rustls_ring::DEFAULT_PROVIDER.clone()) + } + + #[cfg(not(any(feature = "aws_lc_rs", feature = "ring")))] + { + panic!("enable either the `aws_lc_rs` or `ring` feature") + } +} diff --git a/examples/server.rs b/examples/server.rs index 5997715c..3a0bf3b5 100644 --- a/examples/server.rs +++ b/examples/server.rs @@ -20,10 +20,10 @@ use argh::FromArgs; use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; use tokio::time::{Instant, timeout, timeout_at}; -use tokio_rustls::rustls::ServerConfig; +use tokio_rustls::rustls::crypto::Identity; use tokio_rustls::rustls::pki_types::pem::PemObject; use tokio_rustls::rustls::pki_types::{CertificateDer, PrivateKeyDer}; -use tokio_rustls::rustls::server::Acceptor; +use tokio_rustls::rustls::{self, ServerConfig}; use tokio_rustls::server::TlsStream; use tokio_rustls::{LazyConfigAcceptor, TlsAcceptor}; @@ -41,13 +41,12 @@ async fn main() -> Result<(), Box> { let io_timeout = Duration::from_secs(options.io_timeout); let shutdown_timeout = Duration::from_secs(options.shutdown_timeout); + let certs = CertificateDer::pem_file_iter(&options.cert)?.collect::, _>>()?; + let identity = Arc::new(Identity::from_cert_chain(certs)?); let config = Arc::new( - ServerConfig::builder() + ServerConfig::builder(provider()) .with_no_client_auth() - .with_single_cert( - CertificateDer::pem_file_iter(&options.cert)?.collect::>()?, - PrivateKeyDer::from_pem_file(&options.key)?, - )?, + .with_single_cert(identity, PrivateKeyDer::from_pem_file(&options.key)?)?, ); let acceptor = TlsAcceptor::from(config.clone()); let listener = TcpListener::bind(&addr).await?; @@ -133,7 +132,7 @@ async fn accept_lazy( ) -> io::Result> { let deadline = Instant::now() + handshake_timeout; - let acceptor = LazyConfigAcceptor::new(Acceptor::default(), stream); + let acceptor = LazyConfigAcceptor::new(stream); tokio::pin!(acceptor); // Read the ClientHello, respecting the deadline. @@ -156,7 +155,7 @@ async fn accept_lazy( // The ClientHello is now available, e.g. for SNI-based config selection. if let Some(sni) = start.client_hello().server_name() { - println!("ClientHello with SNI: {sni}"); + println!("ClientHello with SNI: {}", sni.as_ref()); } // Complete the handshake, respecting the deadline. @@ -221,3 +220,20 @@ async fn with_timeout( )), } } + +fn provider() -> Arc { + #[cfg(feature = "aws_lc_rs")] + { + Arc::new(rustls_aws_lc_rs::DEFAULT_PROVIDER.clone()) + } + + #[cfg(all(not(feature = "aws_lc_rs"), feature = "ring"))] + { + Arc::new(rustls_ring::DEFAULT_PROVIDER.clone()) + } + + #[cfg(not(any(feature = "aws_lc_rs", feature = "ring")))] + { + panic!("enable either the `aws_lc_rs` or `ring` feature") + } +} diff --git a/src/client.rs b/src/client.rs index bb01cf94..5654f74b 100644 --- a/src/client.rs +++ b/src/client.rs @@ -1,5 +1,5 @@ use std::future::Future; -use std::io::{self, BufRead as _}; +use std::io; #[cfg(unix)] use std::os::unix::io::{AsRawFd, RawFd}; #[cfg(windows)] @@ -10,11 +10,16 @@ use std::sync::Arc; use std::task::Waker; use std::task::{Context, Poll}; +#[cfg(feature = "early-data")] +use rustls::crypto::cipher::OutboundPlain; +use rustls::enums::ApplicationProtocol; use rustls::pki_types::ServerName; -use rustls::{ClientConfig, ClientConnection}; +use rustls::{ClientConfig, ClientConnection, Connection as _, VecInput}; use tokio::io::{AsyncBufRead, AsyncRead, AsyncWrite, ReadBuf}; -use crate::common::{IoSession, MidHandshake, Stream, TlsState}; +#[cfg(feature = "early-data")] +use crate::common::DEFAULT_BUFFER_LIMIT; +use crate::common::{CursorVec, IoSession, IoSessionParts, MidHandshake, Stream, TlsState}; /// A wrapper around a `rustls::ClientConfig`, providing an async `connect` method. #[derive(Clone)] @@ -78,15 +83,26 @@ impl TlsConnector { IO: AsyncRead + AsyncWrite + Unpin, F: FnOnce(&mut ClientConnection), { - let alpn = alpn_protocols.unwrap_or_else(|| self.inner.alpn_protocols.clone()); - let mut session = match ClientConnection::new_with_alpn(self.inner.clone(), domain, alpn) { + let builder = self.inner.connect(domain); + let builder = if let Some(alpn_protocols) = alpn_protocols { + builder.with_alpn( + alpn_protocols + .into_iter() + .map(ApplicationProtocol::from) + .collect(), + ) + } else { + builder + }; + let mut tls = Vec::new(); + let mut session = match builder.build(&mut tls) { Ok(session) => session, Err(error) => { return Connect(MidHandshake::Error { io: stream, // TODO(eliza): should this really return an `io::Error`? // Probably not... - error: io::Error::new(io::ErrorKind::Other, error), + error: io::Error::other(error), }); } }; @@ -111,6 +127,9 @@ impl TlsConnector { early_waker: None, session, + input: VecInput::default(), + tls: CursorVec::from_vec(tls), + plaintext: CursorVec::default(), })) } @@ -234,6 +253,9 @@ pub struct FallibleConnect(MidHandshake>); pub struct TlsStream { pub(crate) io: IO, pub(crate) session: ClientConnection, + pub(crate) input: VecInput, + pub(crate) tls: CursorVec, + pub(crate) plaintext: CursorVec, pub(crate) state: TlsState, pub(crate) need_flush: bool, @@ -252,6 +274,9 @@ impl TlsStream { (&mut self.io, &mut self.session) } + /// Extract the transport and rustls connection. + /// + /// Any buffered TLS or plaintext data is discarded. #[inline] pub fn into_inner(self) -> (IO, ClientConnection) { (self.io, self.session) @@ -288,18 +313,21 @@ impl IoSession for TlsStream { } #[inline] - fn get_mut(&mut self) -> (&mut TlsState, &mut Self::Io, &mut Self::Session, &mut bool) { - ( - &mut self.state, - &mut self.io, - &mut self.session, - &mut self.need_flush, - ) + fn get_mut(&mut self) -> IoSessionParts<'_, Self::Io, Self::Session> { + IoSessionParts { + state: &mut self.state, + io: &mut self.io, + session: &mut self.session, + input: &mut self.input, + tls: &mut self.tls, + plaintext: &mut self.plaintext, + need_flush: &mut self.need_flush, + } } #[inline] - fn into_io(self) -> Self::Io { - self.io + fn into_io_tls(self) -> (Self::Io, CursorVec) { + (self.io, self.tls) } } @@ -356,8 +384,14 @@ where } TlsState::Stream | TlsState::WriteShutdown => { let this = self.get_mut(); - let stream = - Stream::new(&mut this.io, &mut this.session).set_eof(!this.state.readable()); + let stream = Stream::new( + &mut this.io, + &mut this.session, + &mut this.input, + &mut this.tls, + &mut this.plaintext, + ) + .set_eof(!this.state.readable()); match stream.poll_fill_buf(cx) { Poll::Ready(Ok(buf)) => { @@ -379,7 +413,7 @@ where } fn consume(mut self: Pin<&mut Self>, amt: usize) { - self.session.reader().consume(amt); + self.plaintext.consume(amt); } } @@ -395,9 +429,15 @@ where buf: &[u8], ) -> Poll> { let this = self.get_mut(); - let mut stream = Stream::new(&mut this.io, &mut this.session) - .set_eof(!this.state.readable()) - .set_need_flush(this.need_flush); + let mut stream = Stream::new( + &mut this.io, + &mut this.session, + &mut this.input, + &mut this.tls, + &mut this.plaintext, + ) + .set_eof(!this.state.readable()) + .set_need_flush(this.need_flush); #[cfg(feature = "early-data")] { @@ -411,7 +451,10 @@ where )?; match written { Poll::Ready(0) => {} - Poll::Ready(written) => return Poll::Ready(Ok(written)), + Poll::Ready(written) => { + this.need_flush = stream.need_flush; + return Poll::Ready(Ok(written)); + } Poll::Pending => { this.need_flush = stream.need_flush; return Poll::Pending; @@ -430,9 +473,15 @@ where bufs: &[io::IoSlice<'_>], ) -> Poll> { let this = self.get_mut(); - let mut stream = Stream::new(&mut this.io, &mut this.session) - .set_eof(!this.state.readable()) - .set_need_flush(this.need_flush); + let mut stream = Stream::new( + &mut this.io, + &mut this.session, + &mut this.input, + &mut this.tls, + &mut this.plaintext, + ) + .set_eof(!this.state.readable()) + .set_need_flush(this.need_flush); #[cfg(feature = "early-data")] { @@ -445,7 +494,10 @@ where )?; match written { Poll::Ready(0) => {} - Poll::Ready(written) => return Poll::Ready(Ok(written)), + Poll::Ready(written) => { + this.need_flush = stream.need_flush; + return Poll::Ready(Ok(written)); + } Poll::Pending => { this.need_flush = stream.need_flush; return Poll::Pending; @@ -463,9 +515,15 @@ where fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let this = self.get_mut(); - let mut stream = Stream::new(&mut this.io, &mut this.session) - .set_eof(!this.state.readable()) - .set_need_flush(this.need_flush); + let mut stream = Stream::new( + &mut this.io, + &mut this.session, + &mut this.input, + &mut this.tls, + &mut this.plaintext, + ) + .set_eof(!this.state.readable()) + .set_need_flush(this.need_flush); #[cfg(feature = "early-data")] { @@ -495,13 +553,20 @@ where } if self.state.writeable() { - self.session.send_close_notify(); + let this = self.as_mut().get_mut(); + this.session.send_close_notify(this.tls.append_vec()); self.state.shutdown_write(); } let this = self.get_mut(); - let mut stream = - Stream::new(&mut this.io, &mut this.session).set_eof(!this.state.readable()); + let mut stream = Stream::new( + &mut this.io, + &mut this.session, + &mut this.input, + &mut this.tls, + &mut this.plaintext, + ) + .set_eof(!this.state.readable()); stream.as_mut_pin().poll_shutdown(cx) } } @@ -518,24 +583,32 @@ where IO: AsyncRead + AsyncWrite + Unpin, { if let TlsState::EarlyData(pos, data) = state { - use std::io::Write; + while !stream.tls.is_empty() { + match stream.write_io(cx) { + Poll::Ready(Ok(0)) | Poll::Pending => return Poll::Pending, + Poll::Ready(Ok(_)) => {} + Poll::Ready(Err(err)) => return Poll::Ready(Err(err)), + } + } // write early data if let Some(mut early_data) = stream.session.early_data() { let mut written = 0; + let mut available = DEFAULT_BUFFER_LIMIT; for buf in bufs { - if buf.is_empty() { + if buf.is_empty() || available == 0 { continue; } - let len = match early_data.write(buf) { - Ok(0) => break, - Ok(n) => n, - Err(err) => return Poll::Ready(Err(err)), - }; + let buf = &buf[..buf.len().min(available)]; + let len = early_data.write_tls(OutboundPlain::from(buf), stream.tls.append_vec()); + if len == 0 { + break; + } written += len; + available -= len; data.extend_from_slice(&buf[..len]); if len < buf.len() { @@ -544,6 +617,12 @@ where } if written != 0 { + while !stream.tls.is_empty() { + match stream.write_io(cx) { + Poll::Ready(Ok(0)) | Poll::Pending | Poll::Ready(Err(_)) => break, + Poll::Ready(Ok(_)) => {} + } + } return Poll::Ready(Ok(written)); } } diff --git a/src/common/handshake.rs b/src/common/handshake.rs index a772c00b..d76e4e46 100644 --- a/src/common/handshake.rs +++ b/src/common/handshake.rs @@ -1,22 +1,31 @@ use std::future::Future; -use std::ops::{Deref, DerefMut}; +use std::io::Write as _; use std::pin::Pin; use std::task::{Context, Poll}; use std::{io, mem}; -use rustls::server::AcceptedAlert; -use rustls::{ConnectionCommon, SideData}; +use rustls::{Connection as RustlsConnection, VecInput}; use tokio::io::{AsyncRead, AsyncWrite}; -use crate::common::{Stream, SyncWriteAdapter, TlsState}; +use crate::common::{CursorVec, Stream, SyncWriteAdapter, TlsState}; + +pub(crate) struct IoSessionParts<'a, IO, Session> { + pub(crate) state: &'a mut TlsState, + pub(crate) io: &'a mut IO, + pub(crate) session: &'a mut Session, + pub(crate) input: &'a mut VecInput, + pub(crate) tls: &'a mut CursorVec, + pub(crate) plaintext: &'a mut CursorVec, + pub(crate) need_flush: &'a mut bool, +} pub(crate) trait IoSession { type Io; type Session; fn skip_handshake(&self) -> bool; - fn get_mut(&mut self) -> (&mut TlsState, &mut Self::Io, &mut Self::Session, &mut bool); - fn into_io(self) -> Self::Io; + fn get_mut(&mut self) -> IoSessionParts<'_, Self::Io, Self::Session>; + fn into_io_tls(self) -> (Self::Io, CursorVec); } pub(crate) enum MidHandshake { @@ -24,7 +33,7 @@ pub(crate) enum MidHandshake { End, SendAlert { io: IS::Io, - alert: AcceptedAlert, + alert: CursorVec, error: io::Error, }, Error { @@ -33,12 +42,11 @@ pub(crate) enum MidHandshake { }, } -impl Future for MidHandshake +impl Future for MidHandshake where IS: IoSession + Unpin, IS::Io: AsyncRead + AsyncWrite + Unpin, - IS::Session: DerefMut + Deref> + Unpin, - SD: SideData, + IS::Session: RustlsConnection + Unpin, { type Output = Result; @@ -52,14 +60,32 @@ where mut alert, error, } => loop { - match alert.write(&mut SyncWriteAdapter { io: &mut io, cx }) { - Err(e) if e.kind() == io::ErrorKind::WouldBlock => { + if !alert.is_empty() { + let available = alert.len(); + match (SyncWriteAdapter { io: &mut io, cx }).write(alert.pending()) { + Err(e) if e.kind() == io::ErrorKind::WouldBlock => { + *this = Self::SendAlert { io, error, alert }; + return Poll::Pending; + } + Err(_) | Ok(0) => return Poll::Ready(Err((error, io))), + Ok(written) if written > available => { + alert.clear(); + return Poll::Ready(Err((error, io))); + } + Ok(written) => { + alert.consume(written); + continue; + } + } + } + + match Pin::new(&mut io).poll_flush(cx) { + Poll::Pending => { *this = Self::SendAlert { io, error, alert }; return Poll::Pending; } - Err(_) | Ok(0) => return Poll::Ready(Err((error, io))), - Ok(_) => {} - }; + Poll::Ready(_) => return Poll::Ready(Err((error, io))), + } }, // Starting the handshake returned an error; fail the future immediately. Self::Error { io, error } => return Poll::Ready(Err((error, io))), @@ -67,30 +93,57 @@ where }; if !stream.skip_handshake() { - let (state, io, session, need_flush) = stream.get_mut(); - let mut tls_stream = Stream::new(io, session) - .set_eof(!state.readable()) - .set_need_flush(*need_flush); - - macro_rules! try_poll { - ( $e:expr ) => { - match $e { - Poll::Ready(Ok(x)) => x, - Poll::Ready(Err(err)) => return Poll::Ready(Err((err, stream.into_io()))), - Poll::Pending => { - *need_flush = tls_stream.need_flush; - *this = MidHandshake::Handshaking(stream); - return Poll::Pending; + let IoSessionParts { + state, + io, + session, + input, + tls, + plaintext, + need_flush, + } = stream.get_mut(); + let handshake_error = 'handshake: { + let mut tls_stream = Stream::new(io, session, input, tls, plaintext) + .set_eof(!state.readable()) + .set_need_flush(*need_flush); + + macro_rules! try_poll { + ( $e:expr ) => { + match $e { + Poll::Ready(Ok(x)) => x, + Poll::Ready(Err(err)) => { + break 'handshake Some((err, tls_stream.need_flush)); + } + Poll::Pending => { + *need_flush = tls_stream.need_flush; + *this = MidHandshake::Handshaking(stream); + return Poll::Pending; + } } - } - }; - } + }; + } - while tls_stream.session.is_handshaking() { - try_poll!(tls_stream.handshake(cx)); - } + while tls_stream.session.is_handshaking() { + try_poll!(tls_stream.handshake(cx)); + } + + try_poll!(Pin::new(&mut tls_stream).poll_flush(cx)); + None + }; - try_poll!(Pin::new(&mut tls_stream).poll_flush(cx)); + if let Some((err, need_flush)) = handshake_error { + let (io, alert) = stream.into_io_tls(); + if alert.is_empty() && !need_flush { + return Poll::Ready(Err((err, io))); + } + *this = Self::SendAlert { + io, + alert, + error: err, + }; + cx.waker().wake_by_ref(); + return Poll::Pending; + } } Poll::Ready(Ok(stream)) diff --git a/src/common/mod.rs b/src/common/mod.rs index ede69d9b..91ed7651 100644 --- a/src/common/mod.rs +++ b/src/common/mod.rs @@ -1,13 +1,96 @@ -use std::io::{self, BufRead as _, IoSlice, Read, Write}; -use std::ops::{Deref, DerefMut}; +use std::io::{self, IoSlice, Read, Write}; use std::pin::Pin; use std::task::{Context, Poll}; -use rustls::{ConnectionCommon, SideData}; +use rustls::crypto::cipher::OutboundPlain; +use rustls::{Connection as RustlsConnection, VecInput}; use tokio::io::{AsyncBufRead, AsyncRead, AsyncWrite, ReadBuf}; +pub(crate) const DEFAULT_BUFFER_LIMIT: usize = 64 * 1024; + mod handshake; -pub(crate) use handshake::{IoSession, MidHandshake}; +pub(crate) use handshake::{IoSession, IoSessionParts, MidHandshake}; + +/// A `Vec` with an advancing read cursor. +/// +/// Rustls appends output to a `Vec`, while Tokio may consume that output a +/// little at a time. Keeping a cursor avoids shifting the remaining bytes on +/// every partial read or write. +#[derive(Default)] +pub(crate) struct CursorVec { + bytes: Vec, + start: usize, +} + +impl std::fmt::Debug for CursorVec { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + std::fmt::Debug::fmt(self.pending(), f) + } +} + +impl CursorVec { + pub(crate) fn from_vec(bytes: Vec) -> Self { + Self { bytes, start: 0 } + } + + pub(crate) fn is_empty(&self) -> bool { + self.len() == 0 + } + + pub(crate) fn len(&self) -> usize { + self.bytes.len() - self.start + } + + pub(crate) fn pending(&self) -> &[u8] { + &self.bytes[self.start..] + } + + pub(crate) fn consume(&mut self, amount: usize) { + assert!(amount <= self.len()); + self.start += amount; + + if self.start == self.bytes.len() { + self.bytes.clear(); + self.start = 0; + } + } + + pub(crate) fn clear(&mut self) { + self.bytes.clear(); + self.start = 0; + } + + /// Returns the backing vector for APIs that only append to it. + /// + /// The vector may still contain a consumed prefix, so callers must not + /// inspect or replace its existing contents. + /// + /// Compact only after at least half the stored bytes have been consumed. + /// This bounds the retained prefix across appends while making compaction + /// amortized linear. + pub(crate) fn append_vec(&mut self) -> &mut Vec { + let remaining = self.len(); + if remaining == 0 { + self.clear(); + } else if self.start >= remaining { + self.bytes.copy_within(self.start.., 0); + self.bytes.truncate(remaining); + self.start = 0; + } + + &mut self.bytes + } + + #[cfg(test)] + fn into_pending_vec(mut self) -> Vec { + let remaining = self.len(); + if self.start != 0 { + self.bytes.copy_within(self.start.., 0); + self.bytes.truncate(remaining); + } + self.bytes + } +} #[derive(Debug)] pub(crate) enum TlsState { @@ -62,19 +145,30 @@ impl TlsState { pub(crate) struct Stream<'a, IO, C> { pub(crate) io: &'a mut IO, pub(crate) session: &'a mut C, + pub(crate) input: &'a mut VecInput, + pub(crate) tls: &'a mut CursorVec, + pub(crate) plaintext: &'a mut CursorVec, pub(crate) eof: bool, pub(crate) need_flush: bool, } -impl<'a, IO: AsyncRead + AsyncWrite + Unpin, C, SD> Stream<'a, IO, C> +impl<'a, IO: AsyncRead + AsyncWrite + Unpin, C> Stream<'a, IO, C> where - C: DerefMut + Deref>, - SD: SideData, + C: RustlsConnection, { - pub(crate) fn new(io: &'a mut IO, session: &'a mut C) -> Self { + pub(crate) fn new( + io: &'a mut IO, + session: &'a mut C, + input: &'a mut VecInput, + tls: &'a mut CursorVec, + plaintext: &'a mut CursorVec, + ) -> Self { Stream { io, session, + input, + tls, + plaintext, // The state so far is only used to detect EOF, so both Stream // and EarlyData states should be acceptable. eof: false, @@ -100,30 +194,51 @@ where pub(crate) fn read_io(&mut self, cx: &mut Context) -> Poll> { let mut reader = SyncReadAdapter { io: self.io, cx }; - let n = match self.session.read_tls(&mut reader) { + let n = match self.input.read(&mut reader) { Ok(n) => n, Err(ref err) if err.kind() == io::ErrorKind::WouldBlock => return Poll::Pending, Err(err) => return Poll::Ready(Err(err)), }; - self.session.process_new_packets().map_err(|err| { + let processed = self + .session + .process_new_packets(self.input, self.tls.append_vec()) + .handle_all(self.plaintext.append_vec()); + + if let Err(err) = processed { // In case we have an alert to send describing this error, // try a last-gasp write -- but don't predate the primary // error. let _ = self.write_io(cx); - io::Error::new(io::ErrorKind::InvalidData, err) - })?; + return Poll::Ready(Err(io::Error::new(io::ErrorKind::InvalidData, err))); + } Poll::Ready(Ok(n)) } pub(crate) fn write_io(&mut self, cx: &mut Context) -> Poll> { - let mut writer = SyncWriteAdapter { io: self.io, cx }; + if self.tls.is_empty() { + return Poll::Ready(Ok(0)); + } - match self.session.write_tls(&mut writer) { - Err(ref err) if err.kind() == io::ErrorKind::WouldBlock => Poll::Pending, - result => Poll::Ready(result), + let available = self.tls.len(); + match Pin::new(&mut self.io).poll_write(cx, self.tls.pending()) { + Poll::Ready(Ok(n)) if n > available => { + // The amount actually written is unknowable after a broken + // AsyncWrite implementation violates its contract. + self.tls.clear(); + Poll::Ready(Err(io::Error::other(format!( + "illegal poll_write() return value ({n} > {available})" + )))) + } + Poll::Ready(Ok(n)) => { + self.tls.consume(n); + self.need_flush |= n != 0; + Poll::Ready(Ok(n)) + } + Poll::Ready(Err(err)) => Poll::Ready(Err(err)), + Poll::Pending => Poll::Pending, } } @@ -135,7 +250,7 @@ where let mut write_would_block = false; let mut read_would_block = false; - while self.session.wants_write() { + while !self.tls.is_empty() { match self.write_io(cx) { Poll::Ready(Ok(0)) => return Poll::Ready(Err(io::ErrorKind::WriteZero.into())), Poll::Ready(Ok(n)) => { @@ -188,16 +303,18 @@ where } } - pub(crate) fn poll_fill_buf(mut self, cx: &mut Context<'_>) -> Poll> - where - SD: 'a, - { + pub(crate) fn poll_fill_buf(mut self, cx: &mut Context<'_>) -> Poll> { + if !self.plaintext.is_empty() { + return Poll::Ready(Ok(self.plaintext.pending())); + } + let mut io_pending = false; // read a packet - while !self.eof && self.session.wants_read() { + while self.plaintext.is_empty() && !self.eof && self.session.wants_read() { match self.read_io(cx) { Poll::Ready(Ok(0)) => { + self.eof = true; break; } Poll::Ready(Ok(_)) => (), @@ -209,32 +326,39 @@ where } } - match self.session.reader().into_first_chunk() { - Ok(buf) => { - // Note that this could be empty (i.e. EOF) if a `CloseNotify` has been - // received and there is no more buffered data. - Poll::Ready(Ok(buf)) - } - Err(e) if e.kind() == io::ErrorKind::WouldBlock => { - if !io_pending { - // If `wants_read()` is satisfied, rustls will not return `WouldBlock`. - // But if it does, we can try again. - // If the rustls state is abnormal, it may cause a cyclic wakeup. - // but tokio's cooperative budget will prevent infinite wakeup. - cx.waker().wake_by_ref(); - } + if !self.plaintext.is_empty() { + return Poll::Ready(Ok(self.plaintext.pending())); + } - Poll::Pending - } - Err(e) => Poll::Ready(Err(e)), + if self.eof { + return if self.session.wants_read() { + Poll::Ready(Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "peer closed connection without sending TLS close_notify", + ))) + } else { + Poll::Ready(Ok(self.plaintext.pending())) + }; + } + + if !self.session.wants_read() { + // A close_notify with no remaining plaintext is a clean EOF. + return Poll::Ready(Ok(self.plaintext.pending())); + } + + if !io_pending { + // No operation registered this task's waker. Arrange another poll in + // case the connection made internal progress without yielding data. + cx.waker().wake_by_ref(); } + + Poll::Pending } } -impl<'a, IO: AsyncRead + AsyncWrite + Unpin, C, SD> AsyncRead for Stream<'a, IO, C> +impl<'a, IO: AsyncRead + AsyncWrite + Unpin, C> AsyncRead for Stream<'a, IO, C> where - C: DerefMut + Deref>, - SD: SideData + 'a, + C: RustlsConnection + 'a, { fn poll_read( mut self: Pin<&mut Self>, @@ -244,15 +368,14 @@ where let data = ready!(self.as_mut().poll_fill_buf(cx))?; let amount = buf.remaining().min(data.len()); buf.put_slice(&data[..amount]); - self.session.reader().consume(amount); + self.plaintext.consume(amount); Poll::Ready(Ok(())) } } -impl<'a, IO: AsyncRead + AsyncWrite + Unpin, C, SD> AsyncBufRead for Stream<'a, IO, C> +impl<'a, IO: AsyncRead + AsyncWrite + Unpin, C> AsyncBufRead for Stream<'a, IO, C> where - C: DerefMut + Deref>, - SD: SideData + 'a, + C: RustlsConnection + 'a, { fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let this = self.get_mut(); @@ -260,55 +383,62 @@ where // reborrow io: this.io, session: this.session, - ..*this + input: this.input, + tls: this.tls, + plaintext: this.plaintext, + eof: this.eof, + need_flush: this.need_flush, } .poll_fill_buf(cx) } fn consume(mut self: Pin<&mut Self>, amt: usize) { - self.session.reader().consume(amt); + self.plaintext.consume(amt); } } -impl AsyncWrite for Stream<'_, IO, C> +impl AsyncWrite for Stream<'_, IO, C> where - C: DerefMut + Deref>, - SD: SideData, + C: RustlsConnection, { fn poll_write( mut self: Pin<&mut Self>, cx: &mut Context, buf: &[u8], ) -> Poll> { - let mut pos = 0; + if buf.is_empty() { + return Poll::Ready(Ok(0)); + } - while pos != buf.len() { - let mut would_block = false; + // Do not grow the caller-owned TLS output buffer without bound. Pending + // output must make progress before we accept more plaintext. + while !self.tls.is_empty() { + match self.write_io(cx) { + Poll::Ready(Ok(0)) | Poll::Pending => return Poll::Pending, + Poll::Ready(Ok(_)) => {} + Poll::Ready(Err(err)) => return Poll::Ready(Err(err)), + } + } - match self.session.writer().write(&buf[pos..]) { - Ok(n) => pos += n, - Err(err) => return Poll::Ready(Err(err)), - }; + let len = buf.len().min(DEFAULT_BUFFER_LIMIT); + let this = &mut *self; + if let Err(err) = this + .session + .write_tls((&buf[..len]).into(), this.tls.append_vec()) + { + return Poll::Ready(Err(io::Error::new(io::ErrorKind::InvalidData, err))); + } - while self.session.wants_write() { - match self.write_io(cx) { - Poll::Ready(Ok(0)) | Poll::Pending => { - would_block = true; - break; - } - Poll::Ready(Ok(_)) => (), - Poll::Ready(Err(err)) => return Poll::Ready(Err(err)), - } + // The plaintext has been accepted even when the socket now blocks. The + // generated TLS bytes remain in `tls` for the next poll. + while !self.tls.is_empty() { + match self.write_io(cx) { + Poll::Ready(Ok(0)) | Poll::Pending | Poll::Ready(Err(_)) => break, + Poll::Ready(Ok(_)) => {} } - - return match (pos, would_block) { - (0, true) => Poll::Pending, - (n, true) => Poll::Ready(Ok(n)), - (_, false) => continue, - }; } - Poll::Ready(Ok(pos)) + Poll::Ready(Ok(len)) } fn poll_write_vectored( @@ -320,27 +450,40 @@ where return Poll::Ready(Ok(0)); } - loop { - let mut would_block = false; - let written = self.session.writer().write_vectored(bufs)?; - - while self.session.wants_write() { - match self.write_io(cx) { - Poll::Ready(Ok(0)) | Poll::Pending => { - would_block = true; - break; - } - Poll::Ready(Ok(_)) => (), - Poll::Ready(Err(err)) => return Poll::Ready(Err(err)), - } + while !self.tls.is_empty() { + match self.write_io(cx) { + Poll::Ready(Ok(0)) | Poll::Pending => return Poll::Pending, + Poll::Ready(Ok(_)) => {} + Poll::Ready(Err(err)) => return Poll::Ready(Err(err)), } + } - return match (written, would_block) { - (0, true) => Poll::Pending, - (0, false) => continue, - (n, _) => Poll::Ready(Ok(n)), - }; + let mut available = DEFAULT_BUFFER_LIMIT; + let slices: Vec<&[u8]> = bufs + .iter() + .map(|buf| { + let take = buf.len().min(available); + available -= take; + &buf[..take] + }) + .collect(); + let written = slices.iter().map(|buf| buf.len()).sum(); + let this = &mut *self; + if let Err(err) = this + .session + .write_tls(OutboundPlain::new(&slices), this.tls.append_vec()) + { + return Poll::Ready(Err(io::Error::new(io::ErrorKind::InvalidData, err))); + } + + while !self.tls.is_empty() { + match self.write_io(cx) { + Poll::Ready(Ok(0)) | Poll::Pending | Poll::Ready(Err(_)) => break, + Poll::Ready(Ok(_)) => {} + } } + + Poll::Ready(Ok(written)) } #[inline] @@ -349,17 +492,23 @@ where } fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll> { - self.session.writer().flush()?; - while self.session.wants_write() { + while !self.tls.is_empty() { if ready!(self.write_io(cx))? == 0 { return Poll::Ready(Err(io::ErrorKind::WriteZero.into())); } } - Pin::new(&mut self.io).poll_flush(cx) + + match Pin::new(&mut self.io).poll_flush(cx) { + Poll::Ready(Ok(())) => { + self.need_flush = false; + Poll::Ready(Ok(())) + } + poll => poll, + } } fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - while self.session.wants_write() { + while !self.tls.is_empty() { if ready!(self.write_io(cx))? == 0 { return Poll::Ready(Err(io::ErrorKind::WriteZero.into())); } @@ -433,5 +582,45 @@ impl Write for SyncWriteAdapter<'_, '_, T> { } } +#[cfg(test)] +mod cursor_vec_tests { + use super::CursorVec; + + #[test] + fn preserves_pending_bytes_across_compaction() { + let mut buffer = CursorVec::from_vec(b"abcdef".to_vec()); + buffer.consume(2); + buffer.append_vec().extend_from_slice(b"gh"); + buffer.consume(3); + buffer.append_vec().extend_from_slice(b"ij"); + + assert_eq!(buffer.pending(), b"fghij"); + assert_eq!(buffer.into_pending_vec(), b"fghij"); + } + + #[test] + fn repeated_small_consumption_and_append_keeps_storage_bounded() { + let mut buffer = CursorVec::from_vec(vec![0, 1]); + + for byte in 2..=u8::MAX { + buffer.consume(1); + buffer.append_vec().push(byte); + + assert_eq!(buffer.len(), 2); + assert!(buffer.bytes.len() <= 2 * buffer.len()); + } + + assert_eq!(buffer.pending(), &[u8::MAX - 1, u8::MAX]); + } + + #[test] + fn debug_only_shows_pending_bytes() { + let mut buffer = CursorVec::from_vec(vec![1, 2, 3]); + buffer.consume(2); + + assert_eq!(format!("{buffer:?}"), "[3]"); + } +} + #[cfg(test)] mod test_stream; diff --git a/src/common/test_stream.rs b/src/common/test_stream.rs index 46c8a74b..047115f4 100644 --- a/src/common/test_stream.rs +++ b/src/common/test_stream.rs @@ -1,6 +1,6 @@ #![cfg(any(feature = "aws_lc_rs", feature = "ring"))] -use std::io::{self, Cursor, Read, Write}; +use std::io::{self, Cursor}; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; @@ -8,55 +8,74 @@ use std::task::{Context, Poll}; use futures_util::future::poll_fn; use futures_util::task::noop_waker_ref; use rustls::pki_types::ServerName; -use rustls::{ClientConnection, Connection, ServerConnection}; +use rustls::{ClientConnection, Connection as _, ServerConnection, VecInput}; use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, ReadBuf}; -use super::Stream; +use super::{CursorVec, Stream}; -struct Good<'a>(&'a mut Connection); +struct ConnectionState { + session: C, + input: VecInput, + tls: CursorVec, + plaintext: CursorVec, +} + +impl ConnectionState { + fn new(session: C, tls: Vec) -> Self { + Self { + session, + input: VecInput::default(), + tls: CursorVec::from_vec(tls), + plaintext: CursorVec::default(), + } + } +} + +struct Good<'a>(&'a mut ConnectionState); impl AsyncRead for Good<'_> { fn poll_read( - mut self: Pin<&mut Self>, + self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { - let mut buf2 = buf.initialize_unfilled(); - - Poll::Ready(match self.0.write_tls(buf2.by_ref()) { - Ok(n) => { - buf.advance(n); - Ok(()) - } - Err(err) => Err(err), - }) + let this = self.get_mut(); + let n = buf.remaining().min(this.0.tls.len()); + if n == 0 { + return Poll::Pending; + } + + buf.put_slice(&this.0.tls.pending()[..n]); + this.0.tls.consume(n); + Poll::Ready(Ok(())) } } impl AsyncWrite for Good<'_> { fn poll_write( - mut self: Pin<&mut Self>, + self: Pin<&mut Self>, _cx: &mut Context<'_>, - mut buf: &[u8], + buf: &[u8], ) -> Poll> { - let len = self.0.read_tls(buf.by_ref())?; - self.0 - .process_new_packets() + let this = self.get_mut(); + let mut reader = buf; + let len = this.0.input.read(&mut reader)?; + this.0 + .session + .process_new_packets(&mut this.0.input, this.0.tls.append_vec()) + .handle_all(this.0.plaintext.append_vec()) .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?; Poll::Ready(Ok(len)) } - fn poll_flush(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { - self.0 - .process_new_packets() - .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?; + fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { Poll::Ready(Ok(())) } - fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - self.0.send_close_notify(); - dbg!("sent close notify"); - self.poll_flush(cx) + fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + this.0.session.send_close_notify(this.0.tls.append_vec()); + Poll::Ready(Ok(())) } } @@ -172,16 +191,24 @@ async fn stream_good_bufread() -> io::Result<()> { async fn stream_good_impl(vectored: bool, bufread: bool) -> io::Result<()> { const FILE: &[u8] = include_bytes!("../../README.md"); - let (server, mut client) = make_pair(); - let mut server = Connection::from(server); + let (mut server, mut client) = make_pair(); poll_fn(|cx| do_handshake(&mut client, &mut server, cx)).await?; - io::copy(&mut Cursor::new(FILE), &mut server.writer())?; - server.send_close_notify(); + server + .session + .write_tls(FILE.into(), server.tls.append_vec()) + .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?; + server.session.send_close_notify(server.tls.append_vec()); { let mut good = Good(&mut server); - let mut stream = Stream::new(&mut good, &mut client); + let mut stream = Stream::new( + &mut good, + &mut client.session, + &mut client.input, + &mut client.tls, + &mut client.plaintext, + ); let mut buf = Vec::new(); if bufread { @@ -192,38 +219,33 @@ async fn stream_good_impl(vectored: bool, bufread: bool) -> io::Result<()> { assert_eq!(buf, FILE); dbg!(utils::write(&mut stream, b"Hello World!", vectored).await)?; - stream.session.send_close_notify(); + stream.session.send_close_notify(stream.tls.append_vec()); dbg!(stream.shutdown().await)?; } - let mut buf = String::new(); - dbg!(server.process_new_packets()).map_err(|e| io::Error::new(io::ErrorKind::Other, e))?; - dbg!(server.reader().read_to_string(&mut buf))?; - assert_eq!(buf, "Hello World!"); + assert_eq!(server.plaintext.pending(), b"Hello World!"); Ok(()) as io::Result<()> } #[tokio::test] async fn stream_bad() -> io::Result<()> { - let (server, mut client) = make_pair(); - let mut server = Connection::from(server); + let (mut server, mut client) = make_pair(); poll_fn(|cx| do_handshake(&mut client, &mut server, cx)).await?; - client.set_buffer_limit(Some(1024)); let mut bad = Pending; - let mut stream = Stream::new(&mut bad, &mut client); - assert_eq!( - poll_fn(|cx| stream.as_mut_pin().poll_write(cx, &[0x42; 8])).await?, - 8 + let mut stream = Stream::new( + &mut bad, + &mut client.session, + &mut client.input, + &mut client.tls, + &mut client.plaintext, ); assert_eq!( poll_fn(|cx| stream.as_mut_pin().poll_write(cx, &[0x42; 8])).await?, 8 ); - let r = poll_fn(|cx| stream.as_mut_pin().poll_write(cx, &[0x00; 1024])).await?; // fill buffer - assert!(r < 1024); let mut cx = Context::from_waker(noop_waker_ref()); let ret = stream.as_mut_pin().poll_write(&mut cx, &[0x01]); @@ -234,12 +256,17 @@ async fn stream_bad() -> io::Result<()> { #[tokio::test] async fn stream_handshake() -> io::Result<()> { - let (server, mut client) = make_pair(); - let mut server = Connection::from(server); + let (mut server, mut client) = make_pair(); { let mut good = Good(&mut server); - let mut stream = Stream::new(&mut good, &mut client); + let mut stream = Stream::new( + &mut good, + &mut client.session, + &mut client.input, + &mut client.tls, + &mut client.plaintext, + ); let (r, w) = poll_fn(|cx| stream.handshake(cx)).await?; assert!(r > 0); @@ -248,8 +275,8 @@ async fn stream_handshake() -> io::Result<()> { poll_fn(|cx| stream.handshake(cx)).await?; // finish server handshake } - assert!(!server.is_handshaking()); - assert!(!client.is_handshaking()); + assert!(!server.session.is_handshaking()); + assert!(!client.session.is_handshaking()); Ok(()) as io::Result<()> } @@ -258,12 +285,17 @@ async fn stream_handshake() -> io::Result<()> { async fn stream_buffered_handshake() -> io::Result<()> { use tokio::io::BufWriter; - let (server, mut client) = make_pair(); - let mut server = Connection::from(server); + let (mut server, mut client) = make_pair(); { let mut good = BufWriter::new(Good(&mut server)); - let mut stream = Stream::new(&mut good, &mut client); + let mut stream = Stream::new( + &mut good, + &mut client.session, + &mut client.input, + &mut client.tls, + &mut client.plaintext, + ); let (r, w) = poll_fn(|cx| stream.handshake(cx)).await?; assert!(r > 0); @@ -272,8 +304,8 @@ async fn stream_buffered_handshake() -> io::Result<()> { poll_fn(|cx| stream.handshake(cx)).await?; // finish server handshake } - assert!(!server.is_handshaking()); - assert!(!client.is_handshaking()); + assert!(!server.session.is_handshaking()); + assert!(!client.session.is_handshaking()); Ok(()) as io::Result<()> } @@ -283,7 +315,13 @@ async fn stream_handshake_eof() -> io::Result<()> { let (_, mut client) = make_pair(); let mut bad = Expected(Cursor::new(Vec::new())); - let mut stream = Stream::new(&mut bad, &mut client); + let mut stream = Stream::new( + &mut bad, + &mut client.session, + &mut client.input, + &mut client.tls, + &mut client.plaintext, + ); let mut cx = Context::from_waker(noop_waker_ref()); let r = stream.handshake(&mut cx); @@ -300,7 +338,13 @@ async fn stream_handshake_write_eof() -> io::Result<()> { let (_, mut client) = make_pair(); let mut io = Eof; - let mut stream = Stream::new(&mut io, &mut client); + let mut stream = Stream::new( + &mut io, + &mut client.session, + &mut client.input, + &mut client.tls, + &mut client.plaintext, + ); let mut cx = Context::from_waker(noop_waker_ref()); let r = stream.handshake(&mut cx); @@ -318,7 +362,13 @@ async fn stream_handshake_regression_issues_77() -> io::Result<()> { let (_, mut client) = make_pair(); let mut bad = Expected(Cursor::new(b"\x15\x03\x01\x00\x02\x02\x00".to_vec())); - let mut stream = Stream::new(&mut bad, &mut client); + let mut stream = Stream::new( + &mut bad, + &mut client.session, + &mut client.input, + &mut client.tls, + &mut client.plaintext, + ); let mut cx = Context::from_waker(noop_waker_ref()); let r = stream.handshake(&mut cx); @@ -332,12 +382,17 @@ async fn stream_handshake_regression_issues_77() -> io::Result<()> { #[tokio::test] async fn stream_eof() -> io::Result<()> { - let (server, mut client) = make_pair(); - let mut server = Connection::from(server); + let (mut server, mut client) = make_pair(); poll_fn(|cx| do_handshake(&mut client, &mut server, cx)).await?; let mut bad = Expected(Cursor::new(Vec::new())); - let mut stream = Stream::new(&mut bad, &mut client); + let mut stream = Stream::new( + &mut bad, + &mut client.session, + &mut client.input, + &mut client.tls, + &mut client.plaintext, + ); let mut buf = Vec::new(); let result = stream.read_to_end(&mut buf).await; @@ -351,12 +406,17 @@ async fn stream_eof() -> io::Result<()> { #[tokio::test] async fn stream_write_zero() -> io::Result<()> { - let (server, mut client) = make_pair(); - let mut server = Connection::from(server); + let (mut server, mut client) = make_pair(); poll_fn(|cx| do_handshake(&mut client, &mut server, cx)).await?; let mut io = Eof; - let mut stream = Stream::new(&mut io, &mut client); + let mut stream = Stream::new( + &mut io, + &mut client.session, + &mut client.input, + &mut client.tls, + &mut client.plaintext, + ); stream.write_all(b"1").await.unwrap(); let result = stream.flush().await; @@ -368,29 +428,45 @@ async fn stream_write_zero() -> io::Result<()> { Ok(()) as io::Result<()> } -fn make_pair() -> (ServerConnection, ClientConnection) { +fn make_pair() -> ( + ConnectionState, + ConnectionState, +) { let (sconfig, cconfig) = utils::make_configs(); let server = ServerConnection::new(Arc::new(sconfig)).unwrap(); - let domain = ServerName::try_from("foobar.com").unwrap(); - let client = ClientConnection::new(Arc::new(cconfig), domain).unwrap(); - - (server, client) + let domain = ServerName::try_from("foobar.com").unwrap().to_owned(); + let mut client_tls = Vec::new(); + let client = Arc::new(cconfig) + .connect(domain) + .build(&mut client_tls) + .unwrap(); + + ( + ConnectionState::new(server, Vec::new()), + ConnectionState::new(client, client_tls), + ) } fn do_handshake( - client: &mut ClientConnection, - server: &mut Connection, + client: &mut ConnectionState, + server: &mut ConnectionState, cx: &mut Context<'_>, ) -> Poll> { let mut good = Good(server); - let mut stream = Stream::new(&mut good, client); + let mut stream = Stream::new( + &mut good, + &mut client.session, + &mut client.input, + &mut client.tls, + &mut client.plaintext, + ); while stream.session.is_handshaking() { ready!(stream.handshake(cx))?; } - while stream.session.wants_write() { + while !stream.tls.is_empty() { ready!(stream.write_io(cx))?; } diff --git a/src/lib.rs b/src/lib.rs index 8c6fd795..4b387a7e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -47,7 +47,7 @@ use std::pin::Pin; use std::task::{Context, Poll}; pub use rustls; -use rustls::CommonState; +use rustls::ConnectionOutputs; use tokio::io::{AsyncBufRead, AsyncRead, AsyncWrite, ReadBuf}; macro_rules! ready { @@ -77,7 +77,7 @@ pub enum TlsStream { } impl TlsStream { - pub fn get_ref(&self) -> (&T, &CommonState) { + pub fn get_ref(&self) -> (&T, &ConnectionOutputs) { use TlsStream::*; match self { Client(io) => { @@ -91,21 +91,28 @@ impl TlsStream { } } - pub fn get_mut(&mut self) -> (&mut T, &mut CommonState) { + pub fn get_mut(&mut self) -> (&mut T, TlsConnectionMut<'_>) { use TlsStream::*; match self { Client(io) => { let (io, session) = io.get_mut(); - (io, &mut *session) + (io, TlsConnectionMut::Client(session)) } Server(io) => { let (io, session) = io.get_mut(); - (io, &mut *session) + (io, TlsConnectionMut::Server(session)) } } } } +/// Mutable access to either side's rustls connection. +#[derive(Debug)] +pub enum TlsConnectionMut<'a> { + Client(&'a mut rustls::ClientConnection), + Server(&'a mut rustls::ServerConnection), +} + impl From> for TlsStream { fn from(s: client::TlsStream) -> Self { Self::Client(s) diff --git a/src/server.rs b/src/server.rs index 2b068984..d64bdcc3 100644 --- a/src/server.rs +++ b/src/server.rs @@ -1,5 +1,5 @@ use std::future::Future; -use std::io::{self, BufRead as _}; +use std::io::{self, Write as _}; #[cfg(unix)] use std::os::unix::io::{AsRawFd, RawFd}; #[cfg(windows)] @@ -8,11 +8,14 @@ use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; -use rustls::server::AcceptedAlert; -use rustls::{ServerConfig, ServerConnection}; +use rustls::server::{NeedsInput, ServerHandshake}; +use rustls::{Connection as _, ServerConfig, ServerConnection, VecInput}; use tokio::io::{AsyncBufRead, AsyncRead, AsyncWrite, ReadBuf}; -use crate::common::{IoSession, MidHandshake, Stream, SyncReadAdapter, SyncWriteAdapter, TlsState}; +use crate::common::{ + CursorVec, IoSession, IoSessionParts, MidHandshake, Stream, SyncReadAdapter, SyncWriteAdapter, + TlsState, +}; /// A wrapper around a `rustls::ServerConfig`, providing an async `accept` method. #[derive(Clone)] @@ -47,7 +50,7 @@ impl TlsAcceptor { /// bound the handshake time. /// /// The `f` handler is given a mutable reference to a [`ServerConnection`][] that can be used - /// to configure the connection before the handshake, for example, adjusting the buffer limit. + /// to configure the connection before the handshake. /// /// Because no data has been read from `stream` yet when `f` is called ClientHello /// dependent state (like early data) is not yet available. @@ -65,7 +68,7 @@ impl TlsAcceptor { io: stream, // TODO(eliza): should this really return an `io::Error`? // Probably not... - error: io::Error::new(io::ErrorKind::Other, error), + error: io::Error::other(error), }); } }; @@ -74,6 +77,9 @@ impl TlsAcceptor { Accept(MidHandshake::Handshaking(TlsStream { session, io: stream, + input: VecInput::default(), + tls: CursorVec::default(), + plaintext: CursorVec::default(), state: TlsState::Stream, need_flush: false, })) @@ -94,9 +100,12 @@ impl TlsAcceptor { /// [`ServerConfig`]: https://docs.rs/rustls/latest/rustls/server/struct.ServerConfig.html /// [`ClientHello`]: https://docs.rs/rustls/latest/rustls/server/struct.ClientHello.html pub struct LazyConfigAcceptor { - acceptor: rustls::server::Acceptor, + handshake: Option, io: Option, - alert: Option<(rustls::Error, AcceptedAlert)>, + input: VecInput, + tls: CursorVec, + error: Option, + flush_error: bool, } impl LazyConfigAcceptor @@ -120,11 +129,14 @@ where /// /// [`tokio::time::timeout_at`]: https://docs.rs/tokio/latest/tokio/time/fn.timeout_at.html #[inline] - pub fn new(acceptor: rustls::server::Acceptor, io: IO) -> Self { + pub fn new(io: IO) -> Self { Self { - acceptor, + handshake: Some(ServerHandshake::start()), io: Some(io), - alert: None, + input: VecInput::default(), + tls: CursorVec::default(), + error: None, + flush_error: false, } } @@ -145,7 +157,7 @@ where /// let listener = tokio::net::TcpListener::bind("127.0.0.1:4443").await.unwrap(); /// let (stream, _) = listener.accept().await.unwrap(); /// - /// let acceptor = tokio_rustls::LazyConfigAcceptor::new(rustls::server::Acceptor::default(), stream); + /// let acceptor = tokio_rustls::LazyConfigAcceptor::new(stream); /// tokio::pin!(acceptor); /// /// match acceptor.as_mut().await { @@ -186,45 +198,87 @@ where let io = match this.io.as_mut() { Some(io) => io, None => { - return Poll::Ready(Err(io::Error::new( - io::ErrorKind::Other, + return Poll::Ready(Err(io::Error::other( "acceptor cannot be polled after acceptance", ))); } }; - if let Some((err, mut alert)) = this.alert.take() { - match alert.write(&mut SyncWriteAdapter { io, cx }) { - Err(e) if e.kind() == io::ErrorKind::WouldBlock => { - this.alert = Some((err, alert)); - return Poll::Pending; + while !this.tls.is_empty() { + let available = this.tls.len(); + match (SyncWriteAdapter { io, cx }).write(this.tls.pending()) { + Err(e) if e.kind() == io::ErrorKind::WouldBlock => return Poll::Pending, + Err(e) => { + return Poll::Ready(Err(this.error.take().unwrap_or(e))); } - Ok(0) | Err(_) => { - return Poll::Ready(Err(io::Error::new(io::ErrorKind::InvalidData, err))); + Ok(0) => { + return Poll::Ready(Err(this + .error + .take() + .unwrap_or_else(|| io::ErrorKind::WriteZero.into()))); } - Ok(_) => { - this.alert = Some((err, alert)); - continue; + Ok(written) if written > available => { + this.tls.clear(); + return Poll::Ready(Err(this.error.take().unwrap_or_else(|| { + io::Error::other(format!( + "illegal poll_write() return value ({written} > {available})" + )) + }))); } - }; + Ok(written) => { + this.tls.consume(written); + } + } + } + + if this.error.is_some() { + if this.flush_error { + match Pin::new(io).poll_flush(cx) { + Poll::Pending => return Poll::Pending, + Poll::Ready(_) => {} + } + } + return Poll::Ready(Err(this.error.take().unwrap())); } + let Some(handshake) = this.handshake.take() else { + return Poll::Ready(Err(io::Error::other( + "acceptor cannot be polled after acceptance", + ))); + }; + let mut reader = SyncReadAdapter { io, cx }; - match this.acceptor.read_tls(&mut reader) { + match this.input.read(&mut reader) { Ok(0) => return Err(io::ErrorKind::UnexpectedEof.into()).into(), Ok(_) => {} - Err(e) if e.kind() == io::ErrorKind::WouldBlock => return Poll::Pending, + Err(e) if e.kind() == io::ErrorKind::WouldBlock => { + this.handshake = Some(handshake); + return Poll::Pending; + } Err(e) => return Err(e).into(), } - match this.acceptor.accept() { - Ok(Some(accepted)) => { + match handshake.process(&mut this.input, this.tls.append_vec()) { + Ok(ServerHandshake::Accepted(accepted)) => { let io = this.io.take().unwrap(); - return Poll::Ready(Ok(StartHandshake { accepted, io })); + return Poll::Ready(Ok(StartHandshake { + accepted, + io, + input: std::mem::take(&mut this.input), + tls: std::mem::take(&mut this.tls), + })); + } + Ok(ServerHandshake::NeedsInput(next)) => { + this.handshake = Some(next); + } + Ok(_) => { + return Poll::Ready(Err(io::Error::other( + "unexpected server handshake state before configuration", + ))); } - Ok(None) => {} - Err((err, alert)) => { - this.alert = Some((err, alert)); + Err(error) => { + this.flush_error = !this.tls.is_empty(); + this.error = Some(io::Error::new(io::ErrorKind::InvalidData, error)); } } } @@ -240,20 +294,24 @@ where #[non_exhaustive] #[derive(Debug)] pub struct StartHandshake { - pub accepted: rustls::server::Accepted, - pub io: IO, + accepted: rustls::server::Accepted, + io: IO, + input: VecInput, + tls: CursorVec, } impl StartHandshake where IO: AsyncRead + AsyncWrite + Unpin, { - /// Create a new object from an `IO` transport and prior TLS metadata. - pub fn from_parts(accepted: rustls::server::Accepted, transport: IO) -> Self { - Self { - accepted, - io: transport, - } + /// Returns a reference to the underlying transport. + pub fn get_ref(&self) -> &IO { + &self.io + } + + /// Returns a mutable reference to the underlying transport. + pub fn get_mut(&mut self) -> &mut IO { + &mut self.io } pub fn client_hello(&self) -> rustls::server::ClientHello<'_> { @@ -285,15 +343,34 @@ where where F: FnOnce(&mut ServerConnection), { - let mut conn = match self.accepted.into_connection(config) { - Ok(conn) => conn, - Err((error, alert)) => { - return Accept(MidHandshake::SendAlert { - io: self.io, - alert, - // TODO(eliza): should this really return an `io::Error`? - // Probably not... - error: io::Error::new(io::ErrorKind::InvalidData, error), + let Self { + accepted, + io, + input, + mut tls, + } = self; + let mut conn = match accepted.choose_config(config, tls.append_vec()) { + Ok(ServerHandshake::NeedsInput(next)) => next.into_buffered_connection(), + Ok(_) => { + return Accept(MidHandshake::Error { + io, + error: io::Error::other( + "unexpected server handshake state after choosing configuration", + ), + }); + } + Err(error) => { + // TODO(eliza): should this really return an `io::Error`? + // Probably not... + let error = io::Error::new(io::ErrorKind::InvalidData, error); + return Accept(if tls.is_empty() { + MidHandshake::Error { io, error } + } else { + MidHandshake::SendAlert { + io, + alert: tls, + error, + } }); } }; @@ -301,7 +378,10 @@ where Accept(MidHandshake::Handshaking(TlsStream { session: conn, - io: self.io, + io, + input, + tls, + plaintext: CursorVec::default(), state: TlsState::Stream, need_flush: false, })) @@ -364,6 +444,9 @@ impl Future for FallibleAccept { pub struct TlsStream { pub(crate) io: IO, pub(crate) session: ServerConnection, + pub(crate) input: VecInput, + pub(crate) tls: CursorVec, + pub(crate) plaintext: CursorVec, pub(crate) state: TlsState, pub(crate) need_flush: bool, } @@ -379,6 +462,9 @@ impl TlsStream { (&mut self.io, &mut self.session) } + /// Extract the transport and rustls connection. + /// + /// Any buffered TLS or plaintext data is discarded. #[inline] pub fn into_inner(self) -> (IO, ServerConnection) { (self.io, self.session) @@ -395,18 +481,21 @@ impl IoSession for TlsStream { } #[inline] - fn get_mut(&mut self) -> (&mut TlsState, &mut Self::Io, &mut Self::Session, &mut bool) { - ( - &mut self.state, - &mut self.io, - &mut self.session, - &mut self.need_flush, - ) + fn get_mut(&mut self) -> IoSessionParts<'_, Self::Io, Self::Session> { + IoSessionParts { + state: &mut self.state, + io: &mut self.io, + session: &mut self.session, + input: &mut self.input, + tls: &mut self.tls, + plaintext: &mut self.plaintext, + need_flush: &mut self.need_flush, + } } #[inline] - fn into_io(self) -> Self::Io { - self.io + fn into_io_tls(self) -> (Self::Io, CursorVec) { + (self.io, self.tls) } } @@ -435,8 +524,14 @@ where match self.state { TlsState::Stream | TlsState::WriteShutdown => { let this = self.get_mut(); - let stream = - Stream::new(&mut this.io, &mut this.session).set_eof(!this.state.readable()); + let stream = Stream::new( + &mut this.io, + &mut this.session, + &mut this.input, + &mut this.tls, + &mut this.plaintext, + ) + .set_eof(!this.state.readable()); match stream.poll_fill_buf(cx) { Poll::Ready(Ok(buf)) => { @@ -460,7 +555,7 @@ where } fn consume(mut self: Pin<&mut Self>, amt: usize) { - self.session.reader().consume(amt); + self.plaintext.consume(amt); } } @@ -476,8 +571,14 @@ where buf: &[u8], ) -> Poll> { let this = self.get_mut(); - let mut stream = - Stream::new(&mut this.io, &mut this.session).set_eof(!this.state.readable()); + let mut stream = Stream::new( + &mut this.io, + &mut this.session, + &mut this.input, + &mut this.tls, + &mut this.plaintext, + ) + .set_eof(!this.state.readable()); stream.as_mut_pin().poll_write(cx, buf) } @@ -489,8 +590,14 @@ where bufs: &[io::IoSlice<'_>], ) -> Poll> { let this = self.get_mut(); - let mut stream = - Stream::new(&mut this.io, &mut this.session).set_eof(!this.state.readable()); + let mut stream = Stream::new( + &mut this.io, + &mut this.session, + &mut this.input, + &mut this.tls, + &mut this.plaintext, + ) + .set_eof(!this.state.readable()); stream.as_mut_pin().poll_write_vectored(cx, bufs) } @@ -501,20 +608,33 @@ where fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let this = self.get_mut(); - let mut stream = - Stream::new(&mut this.io, &mut this.session).set_eof(!this.state.readable()); + let mut stream = Stream::new( + &mut this.io, + &mut this.session, + &mut this.input, + &mut this.tls, + &mut this.plaintext, + ) + .set_eof(!this.state.readable()); stream.as_mut_pin().poll_flush(cx) } fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { if self.state.writeable() { - self.session.send_close_notify(); + let this = self.as_mut().get_mut(); + this.session.send_close_notify(this.tls.append_vec()); self.state.shutdown_write(); } let this = self.get_mut(); - let mut stream = - Stream::new(&mut this.io, &mut this.session).set_eof(!this.state.readable()); + let mut stream = Stream::new( + &mut this.io, + &mut this.session, + &mut this.input, + &mut this.tls, + &mut this.plaintext, + ) + .set_eof(!this.state.readable()); stream.as_mut_pin().poll_shutdown(cx) } } diff --git a/tests/badssl.rs b/tests/badssl.rs index 3e5e7f88..4fd5ac00 100644 --- a/tests/badssl.rs +++ b/tests/badssl.rs @@ -46,9 +46,15 @@ async fn test_tls12_vectored() -> io::Result<()> { async fn test_tls12_impl(vectored: bool) -> io::Result<()> { let mut root_store = rustls::RootCertStore::empty(); root_store.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned()); - let config = rustls::ClientConfig::builder_with_protocol_versions(&[&rustls::version::TLS12]) + let mut provider = (*utils::provider()).clone(); + provider.tls13_cipher_suites = Vec::new().into(); + + let mut config = rustls::ClientConfig::builder(Arc::new(provider)) .with_root_certificates(root_store) - .with_no_client_auth(); + .with_no_client_auth() + .unwrap(); + // This interoperability endpoint does not negotiate RFC 7627 EMS. + config.require_ems = false; let config = Arc::new(config); let domain = "tls-v1-2.badssl.com"; @@ -83,9 +89,12 @@ async fn test_modern_vectored() -> io::Result<()> { async fn test_modern_impl(vectored: bool) -> io::Result<()> { let mut root_store = rustls::RootCertStore::empty(); root_store.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned()); - let config = rustls::ClientConfig::builder() + let mut config = rustls::ClientConfig::builder(utils::provider()) .with_root_certificates(root_store) - .with_no_client_auth(); + .with_no_client_auth() + .unwrap(); + // This interoperability endpoint does not negotiate RFC 7627 EMS. + config.require_ems = false; let config = Arc::new(config); let domain = "mozilla-modern.badssl.com"; diff --git a/tests/early-data.rs b/tests/early-data.rs index 3ec91deb..16a000f1 100644 --- a/tests/early-data.rs +++ b/tests/early-data.rs @@ -2,12 +2,16 @@ use std::io::{self, Read, Write}; use std::net::{SocketAddr, TcpListener}; +use std::pin::Pin; use std::sync::Arc; +use std::task::{Context, Poll}; use std::thread; +use futures_util::FutureExt as _; use rustls::pki_types::ServerName; -use rustls::{self, ClientConfig, ServerConnection, Stream}; -use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; +use rustls::{self, ClientConfig, Connection as _, ServerConnection, VecInput}; +use rustls_util::{Stream, complete_io}; +use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, ReadBuf}; use tokio::net::TcpStream; use tokio_rustls::TlsConnector; use tokio_rustls::client::TlsStream; @@ -53,8 +57,22 @@ async fn test_0rtt_impl( wrapper: impl Fn(TcpStream) -> S, vectored: bool, ) -> io::Result<()> { + let (client, addr) = start_server(8192)?; + + let (io, buf) = send(client.clone(), addr, &wrapper, b"hello", vectored).await?; + assert!(!io.get_ref().1.is_early_data_accepted()); + assert_eq!("LATE:hello", String::from_utf8_lossy(&buf)); + + let (io, buf) = send(client, addr, wrapper, b"world!", vectored).await?; + assert!(io.get_ref().1.is_early_data_accepted()); + assert_eq!("EARLY:world!LATE:", String::from_utf8_lossy(&buf)); + + Ok(()) +} + +fn start_server(max_early_data_size: u32) -> io::Result<(Arc, SocketAddr)> { let (mut server, mut client) = utils::make_configs(); - server.max_early_data_size = 8192; + server.max_early_data_size = max_early_data_size; let server = Arc::new(server); let listener = TcpListener::bind("127.0.0.1:0")?; @@ -66,28 +84,60 @@ async fn test_0rtt_impl( let server = Arc::clone(&server); thread::spawn(move || { let mut conn = ServerConnection::new(server).unwrap(); - conn.complete_io(&mut sock).unwrap(); + let mut input = VecInput::default(); + let mut received_plaintext = Vec::new(); + let mut output = Vec::new(); + complete_io( + &mut sock, + &mut input, + &mut received_plaintext, + &mut output, + &mut conn, + ) + .unwrap(); if let Some(mut early_data) = conn.early_data() { let mut buf = Vec::new(); early_data.read_to_end(&mut buf).unwrap(); - let mut stream = Stream::new(&mut conn, &mut sock); + let mut stream = Stream::new( + &mut input, + &mut received_plaintext, + &mut output, + &mut conn, + &mut sock, + ); stream.write_all(b"EARLY:").unwrap(); stream.write_all(&buf).unwrap(); } - let mut stream = Stream::new(&mut conn, &mut sock); - stream.write_all(b"LATE:").unwrap(); - loop { - let mut buf = [0; 1024]; - let n = stream.read(&mut buf).unwrap(); - if n == 0 { - conn.send_close_notify(); - conn.complete_io(&mut sock).unwrap(); - break; + { + let mut stream = Stream::new( + &mut input, + &mut received_plaintext, + &mut output, + &mut conn, + &mut sock, + ); + stream.write_all(b"LATE:").unwrap(); + loop { + let mut buf = [0; 1024]; + let n = stream.read(&mut buf).unwrap(); + if n == 0 { + break; + } + stream.write_all(&buf[..n]).unwrap(); } - stream.write_all(&buf[..n]).unwrap(); } + + conn.send_close_notify(&mut output); + complete_io( + &mut sock, + &mut input, + &mut received_plaintext, + &mut output, + &mut conn, + ) + .unwrap(); }); } }); @@ -96,16 +146,71 @@ async fn test_0rtt_impl( let client = Arc::new(client); let addr = SocketAddr::from(([127, 0, 0, 1], server_port)); - let (io, buf) = send(client.clone(), addr, &wrapper, b"hello", vectored).await?; - assert!(!io.get_ref().1.is_early_data_accepted()); - assert_eq!("LATE:hello", String::from_utf8_lossy(&buf)); + Ok((client, addr)) +} - let (io, buf) = send(client, addr, wrapper, b"world!", vectored).await?; - assert!(io.get_ref().1.is_early_data_accepted()); - assert_eq!("EARLY:world!LATE:", String::from_utf8_lossy(&buf)); +#[tokio::test] +async fn early_data_applies_output_backpressure() -> io::Result<()> { + let (client, addr) = start_server(256 * 1024)?; + let _ = send(client.clone(), addr, |stream| stream, b"prime", false).await?; + + let connector = TlsConnector::from(client).early_data(true); + let domain = ServerName::try_from("foobar.com").unwrap(); + let mut stream = connector.connect(domain, ControlledIo::default()).await?; + + assert_eq!(stream.write(b"warmup").await?, 6); + stream.get_mut().0.blocked = true; + + let data = vec![0u8; 128 * 1024]; + assert_eq!(stream.write(&data).await?, 64 * 1024); + assert!(stream.write(&data[64 * 1024..]).now_or_never().is_none()); Ok(()) } +#[derive(Debug, Default)] +struct ControlledIo { + blocked: bool, + written: usize, +} + +impl AsyncRead for ControlledIo { + fn poll_read( + self: Pin<&mut Self>, + _: &mut Context<'_>, + _: &mut ReadBuf<'_>, + ) -> Poll> { + Poll::Pending + } +} + +impl AsyncWrite for ControlledIo { + fn poll_write( + self: Pin<&mut Self>, + _: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + let this = self.get_mut(); + if this.blocked { + return Poll::Pending; + } + + this.written += buf.len(); + Poll::Ready(Ok(buf.len())) + } + + fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + if self.blocked { + Poll::Pending + } else { + Poll::Ready(Ok(())) + } + } + + fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } +} + // Include `utils` module include!("utils.rs"); diff --git a/tests/test.rs b/tests/test.rs index e64cb05b..d094a9e3 100644 --- a/tests/test.rs +++ b/tests/test.rs @@ -158,11 +158,14 @@ async fn test_lazy_config_acceptor() -> io::Result<()> { client.read_to_end(&mut buf).await.unwrap(); }); - let acceptor = LazyConfigAcceptor::new(rustls::server::Acceptor::default(), sstream); + let acceptor = LazyConfigAcceptor::new(sstream); let start = acceptor.await.unwrap(); let ch = start.client_hello(); - assert_eq!(ch.server_name(), Some("foobar.com")); + assert_eq!( + ch.server_name().map(|name| name.as_ref()), + Some("foobar.com") + ); assert_eq!( ch.alpn() .map(|protos| protos.collect::>()) @@ -183,7 +186,7 @@ async fn test_lazy_config_acceptor() -> io::Result<()> { #[tokio::test] async fn lazy_config_acceptor_eof() { let buf = Cursor::new(Vec::new()); - let acceptor = LazyConfigAcceptor::new(rustls::server::Acceptor::default(), buf); + let acceptor = LazyConfigAcceptor::new(buf); let accept_result = match time::timeout(Duration::from_secs(3), acceptor).await { Ok(res) => res, @@ -211,7 +214,7 @@ async fn lazy_config_acceptor_take_io() -> Result<(), rustls::Error> { tx.send(buf).unwrap(); }); - let acceptor = LazyConfigAcceptor::new(rustls::server::Acceptor::default(), sstream); + let acceptor = LazyConfigAcceptor::new(sstream); futures_util::pin_mut!(acceptor); if (acceptor.as_mut().await).is_ok() { panic!("Expected Err(err)"); @@ -268,7 +271,7 @@ async fn acceptor_alert() { tx.send(buf).unwrap(); }); - let accept = LazyConfigAcceptor::new(rustls::server::Acceptor::default(), sstream); + let accept = LazyConfigAcceptor::new(utils::FlushWrapper::new(sstream)); let Ok(Ok(start_handshake)) = time::timeout(Duration::from_secs(3), accept).await else { panic!("timeout"); @@ -288,6 +291,34 @@ async fn acceptor_alert() { assert_eq!(received, [0x15, 0x03, 0x03, 0x00, 0x02, 0x02, 0x46]); } +#[tokio::test] +async fn acceptor_alert_flushes_buffered_writer() { + let (sconfig, _) = utils::make_configs(); + let (mut client, server) = tokio::io::duplex(64); + + let peer = tokio::spawn(async move { + client.write_all(b"not tls").await.unwrap(); + client.shutdown().await.unwrap(); + let mut received = Vec::new(); + client.read_to_end(&mut received).await.unwrap(); + received + }); + + let acceptor = TlsAcceptor::from(Arc::new(sconfig)); + let (error, io) = acceptor + .accept(utils::FlushWrapper::new(server)) + .into_fallible() + .await + .unwrap_err(); + + assert_eq!(error.kind(), io::ErrorKind::InvalidData); + assert!(!io.has_pending_output()); + drop(io); + + let received = peer.await.unwrap(); + assert_eq!(received, b"\x15\x03\x03\x00\x02\x02\x32"); +} + #[tokio::test] async fn lazy_config_acceptor_alert() { // Intentionally small so that we have to call alert.write several times @@ -304,7 +335,7 @@ async fn lazy_config_acceptor_alert() { tx.send(buf).unwrap(); }); - let acceptor = LazyConfigAcceptor::new(rustls::server::Acceptor::default(), sstream); + let acceptor = LazyConfigAcceptor::new(utils::FlushWrapper::new(sstream)); let Ok(accept_result) = time::timeout(Duration::from_secs(3), acceptor).await else { panic!("timeout"); diff --git a/tests/utils.rs b/tests/utils.rs index da0cd4b3..adfc74a9 100644 --- a/tests/utils.rs +++ b/tests/utils.rs @@ -2,8 +2,10 @@ mod utils { use std::collections::VecDeque; use std::io::IoSlice; use std::pin::Pin; + use std::sync::Arc; use std::task::{Context, Poll}; + use rustls::crypto::{CryptoProvider, Identity}; use rustls::pki_types::pem::PemObject; use rustls::pki_types::{CertificateDer, PrivateKeyDer}; use rustls::{ClientConfig, RootCertStore, ServerConfig}; @@ -24,9 +26,12 @@ mod utils { .collect::, _>>() .unwrap(); let key = PrivateKeyDer::from_pem_slice(EE_KEY.as_bytes()).unwrap(); - let sconfig = ServerConfig::builder() + let identity = Arc::new(Identity::from_cert_chain(cert).unwrap()); + let provider = provider(); + + let sconfig = ServerConfig::builder(provider.clone()) .with_no_client_auth() - .with_single_cert(cert, key) + .with_single_cert(identity, key) .unwrap(); let mut client_root_cert_store = RootCertStore::empty(); @@ -34,13 +39,32 @@ mod utils { client_root_cert_store.add(root.unwrap()).unwrap(); } - let cconfig = ClientConfig::builder() + let cconfig = ClientConfig::builder(provider) .with_root_certificates(client_root_cert_store) - .with_no_client_auth(); + .with_no_client_auth() + .unwrap(); (sconfig, cconfig) } + #[allow(dead_code)] + pub(crate) fn provider() -> Arc { + #[cfg(feature = "aws_lc_rs")] + { + Arc::new(rustls_aws_lc_rs::DEFAULT_PROVIDER.clone()) + } + + #[cfg(all(not(feature = "aws_lc_rs"), feature = "ring"))] + { + Arc::new(rustls_ring::DEFAULT_PROVIDER.clone()) + } + + #[cfg(not(any(feature = "aws_lc_rs", feature = "ring")))] + { + panic!("enable either the `aws_lc_rs` or `ring` feature") + } + } + #[allow(dead_code)] pub(crate) async fn write( w: &mut W, @@ -72,6 +96,7 @@ mod utils { /// An IO wrapper that never flushes when writing, and always returns pending on first flush. /// /// This is used to test that rustls always flushes to completion during handshake. + #[derive(Debug)] pub(crate) struct FlushWrapper { stream: S, buf: VecDeque>, @@ -87,6 +112,11 @@ mod utils { queued: Vec::new(), } } + + #[allow(dead_code)] + pub(crate) fn has_pending_output(&self) -> bool { + !self.buf.is_empty() || !self.queued.is_empty() + } } impl AsyncRead for FlushWrapper {