wasi-tls: enforce the check-write permit in the p2 output stream

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);
+}

crates/wasi-tls/src/p2/io.rs

@@ -265,7 +265,7 @@ where
 }
 
 enum WriteState<IO> {
-    Ready(IO),
+    Ready(IO, usize),
     Writing(AbortOnDropJoinHandle<io::Result<IO>>),
     Flushing(AbortOnDropJoinHandle<io::Result<IO>>),
     Closing(AbortOnDropJoinHandle<io::Result<()>>),
@@ -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<usize> {
         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(_))));
+    }
+}

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<dyn TlsProvider>,
+    path: &str,
+) -> Result<(Option<String>, 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<dyn TlsProvider>) -> 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<dyn TlsProvider>) -> Result<()> {
     run_test(
         provider,