From b6c1d0c45036f5b3d0826985a781f97da4beaa77 Mon Sep 17 00:00:00 2001 From: Xia Chao Date: Sat, 31 Oct 2026 22:20:01 +0800 Subject: [PATCH] wasi-tls: enforce the check-write permit in the p2 output stream Writes over the permitted size now trap instead of being accepted, and a test program covers it. --- .../src/bin/p2_tls_write_too_much.rs | 69 +++++++++++++++++++ crates/wasi-tls/src/p2/io.rs | 49 ++++++++++--- crates/wasi-tls/tests/p2/mod.rs | 59 +++++++++++++++- 3 files changed, 167 insertions(+), 10 deletions(-) create mode 100644 crates/test-programs/src/bin/p2_tls_write_too_much.rs diff --git a/crates/test-programs/src/bin/p2_tls_write_too_much.rs b/crates/test-programs/src/bin/p2_tls_write_too_much.rs new file mode 100644 index 000000000000..7f2b6deab7a2 --- /dev/null +++ b/crates/test-programs/src/bin/p2_tls_write_too_much.rs @@ -0,0 +1,69 @@ +//! Writes more bytes to a TLS stream than `check-write` allowed. + +use anyhow::{Context, Result, anyhow}; +use test_programs::wasi::sockets::network::{IpAddress, IpSocketAddress, Network}; +use test_programs::wasi::sockets::tcp::TcpSocket; +use test_programs::wasi::tls::types::ClientHandshake; + +const PORT: u16 = 443; + +fn write_over_permit(domain: &str, ip: IpAddress) -> Result<()> { + let net = Network::default(); + + let socket = TcpSocket::new(ip.family()).unwrap(); + let (tcp_input, tcp_output) = socket + .blocking_connect(&net, IpSocketAddress::new(ip, PORT)) + .context("tcp connect failed")?; + + let (_connection, _tls_input, tls_output) = ClientHandshake::new(domain, tcp_input, tcp_output) + .blocking_finish() + .context("tls handshake failed")?; + + let permit = tls_output + .check_write() + .map_err(|e| anyhow!("check_write failed: {e:?}"))?; + eprintln!("check-write permit: {permit}"); + + let data = vec![0u8; usize::try_from(permit).unwrap() + 1]; + // The host is expected to trap here. Reaching the next line means the + // oversized write was accepted. + let _ = tls_output.write(&data); + eprintln!("write returned without trapping"); + Ok(()) +} + +fn try_live_endpoints(test: impl Fn(&str, IpAddress) -> Result<()>) { + const DOMAINS: &[&str] = &[ + "example.com", + "api.github.com", + "docs.wasmtime.dev", + "bytecodealliance.org", + "www.rust-lang.org", + ]; + + let net = Network::default(); + + for &domain in DOMAINS { + let result = (|| { + let ip = net + .permissive_blocking_resolve_addresses(domain)? + .first() + .map(|a| a.to_owned()) + .ok_or_else(|| anyhow!("DNS lookup failed."))?; + test(&domain, ip) + })(); + + match result { + Ok(()) => return, + Err(e) => { + eprintln!("test for {domain} failed: {e:#}"); + } + } + } + + panic!("all endpoints failed"); +} + +fn main() { + try_live_endpoints(write_over_permit); +} diff --git a/crates/wasi-tls/src/p2/io.rs b/crates/wasi-tls/src/p2/io.rs index 4542e15741ad..fd5b1ae0b378 100644 --- a/crates/wasi-tls/src/p2/io.rs +++ b/crates/wasi-tls/src/p2/io.rs @@ -265,7 +265,7 @@ where } enum WriteState { - Ready(IO), + Ready(IO, usize), Writing(AbortOnDropJoinHandle>), Flushing(AbortOnDropJoinHandle>), Closing(AbortOnDropJoinHandle>), @@ -279,21 +279,27 @@ where IO: AsyncWrite + Send + Unpin + 'static, { fn new(stream: IO) -> Self { - Self::Ready(stream) + Self::Ready(stream, 0) } fn write(&mut self, mut bytes: bytes::Bytes) -> StreamResult<()> { - let WriteState::Ready(_) = self else { + let WriteState::Ready(_, permit) = self else { return Err(StreamError::Trap(wasmtime::format_err!( "unpermitted: must call check_write first" ))); }; + if bytes.len() > *permit { + return Err(StreamError::Trap(wasmtime::format_err!( + "not permitted to write {} bytes", + bytes.len() + ))); + } if bytes.is_empty() { return Ok(()); } - let WriteState::Ready(mut stream) = std::mem::replace(self, WriteState::Closed) else { + let WriteState::Ready(mut stream, _) = std::mem::replace(self, WriteState::Closed) else { unreachable!() }; @@ -312,8 +318,8 @@ where fn flush(&mut self) -> StreamResult<()> { match self { // Immediately flush: - WriteState::Ready(_) => { - let WriteState::Ready(mut stream) = std::mem::replace(self, WriteState::Closed) + WriteState::Ready(_, _) => { + let WriteState::Ready(mut stream, _) = std::mem::replace(self, WriteState::Closed) else { unreachable!() }; @@ -344,7 +350,10 @@ where fn check_write(&mut self) -> StreamResult { match self { - WriteState::Ready(_) => Ok(READY_SIZE), + WriteState::Ready(_, permit) => { + *permit = READY_SIZE; + Ok(READY_SIZE) + } WriteState::Writing(_) => Ok(0), WriteState::Flushing(_) => Ok(0), WriteState::Closing(_) => Ok(0), @@ -362,7 +371,7 @@ where fn close(&mut self) { match std::mem::replace(self, WriteState::Closed) { // No write in progress, immediately shut down: - WriteState::Ready(mut stream) => { + WriteState::Ready(mut stream, _) => { *self = WriteState::Closing(wasmtime_wasi::runtime::spawn(async move { stream.shutdown().await })); @@ -395,7 +404,7 @@ where match self { WriteState::Writing(task) | WriteState::Flushing(task) => { *self = match task.await { - Ok(s) => WriteState::Ready(s), + Ok(s) => WriteState::Ready(s, 0), Err(e) => WriteState::Error(e), } } @@ -409,3 +418,25 @@ where } } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn write_requires_a_permit() { + let mut state = WriteState::new(tokio::io::sink()); + assert!(matches!( + state.write(bytes::Bytes::from_static(b"x")), + Err(StreamError::Trap(_)) + )); + } + + #[test] + fn write_rejects_bytes_over_the_permit() { + let mut state = WriteState::new(tokio::io::sink()); + let permit = state.check_write().unwrap(); + let too_many = bytes::Bytes::from(vec![0; permit + 1]); + assert!(matches!(state.write(too_many), Err(StreamError::Trap(_)))); + } +} diff --git a/crates/wasi-tls/tests/p2/mod.rs b/crates/wasi-tls/tests/p2/mod.rs index c5d5c3df005d..1e4abe29d50b 100644 --- a/crates/wasi-tls/tests/p2/mod.rs +++ b/crates/wasi-tls/tests/p2/mod.rs @@ -3,7 +3,9 @@ use wasmtime::{ component::{Component, Linker, ResourceTable}, format_err, }; -use wasmtime_wasi::{WasiCtx, WasiCtxView, WasiView, p2::bindings::Command}; +use wasmtime_wasi::{ + WasiCtx, WasiCtxView, WasiView, p2::bindings::Command, p2::pipe::MemoryOutputPipe, +}; use wasmtime_wasi_tls::{ TlsProvider, WasiTlsCtx, WasiTlsCtxBuilder, WasiTlsCtxView, WasiTlsView, p2, }; @@ -106,6 +108,61 @@ mod nativetls { test_programs_artifacts::foreach_p2_tls!(nativetls_test_case); } +async fn run_test_capture_stderr( + provider: Box, + path: &str, +) -> Result<(Option, String)> { + let stderr = MemoryOutputPipe::new(64 * 1024); + let ctx = Ctx { + table: ResourceTable::new(), + wasi_ctx: WasiCtx::builder() + .inherit_stdout() + .stderr(stderr.clone()) + .inherit_network() + .allow_ip_name_lookup(true) + .allow_tcp(true) + .allow_udp(true) + .build(), + wasi_tls_ctx: WasiTlsCtxBuilder::new().provider(provider).build(), + }; + + let engine = test_programs_artifacts::engine(|_config| {}); + let mut store = Store::new(&engine, ctx); + let component = Component::from_file(&engine, path)?; + + let mut linker = Linker::new(&engine); + wasmtime_wasi::p2::add_to_linker_async(&mut linker)?; + let mut opts = p2::LinkOptions::default(); + opts.tls(true); + wasmtime_wasi_tls::p2::add_to_linker(&mut linker, &opts)?; + + let command = Command::instantiate_async(&mut store, &component, &linker).await?; + let trapped = match command.wasi_cli_run().call_run(&mut store).await { + Ok(Ok(())) => None, + Ok(Err(())) => Some(String::from("command returned with failing exit status")), + Err(e) => Some(format!("{e:#}")), + }; + let captured = String::from_utf8_lossy(&stderr.contents()).into_owned(); + Ok((trapped, captured)) +} + +async fn p2_tls_write_too_much(provider: Box) -> Result<()> { + let (trapped, stderr) = run_test_capture_stderr( + provider, + test_programs_artifacts::P2_TLS_WRITE_TOO_MUCH_COMPONENT, + ) + .await?; + assert!( + stderr.contains("check-write permit:"), + "the guest never reached the write: {stderr}" + ); + assert!( + trapped.is_some(), + "the oversized write returned instead of trapping: {stderr}" + ); + Ok(()) +} + async fn p2_tls_sample_application(provider: Box) -> Result<()> { run_test( provider,