From 34817963f1338d1855826cab0176f91eb7bc6ac3 Mon Sep 17 00:00:00 2001 From: Adam Bratschi-Kaye Date: Fri, 28 Aug 2026 20:21:01 +0000 Subject: [PATCH 1/2] add regression test --- .../src/bin/p3_http_drop_transmit.rs | 62 +++++++++++++++++++ crates/wasi-http/tests/all/p3/mod.rs | 6 ++ 2 files changed, 68 insertions(+) create mode 100644 crates/test-programs/src/bin/p3_http_drop_transmit.rs diff --git a/crates/test-programs/src/bin/p3_http_drop_transmit.rs b/crates/test-programs/src/bin/p3_http_drop_transmit.rs new file mode 100644 index 000000000000..e8ae3f69111c --- /dev/null +++ b/crates/test-programs/src/bin/p3_http_drop_transmit.rs @@ -0,0 +1,62 @@ +use anyhow::Context as _; +use futures::join; +use test_programs::p3::wasi::http::client; +use test_programs::p3::wasi::http::types::{ + Headers, Method, Request, RequestOptions, Response, Scheme, +}; +use test_programs::p3::{wit_future, wit_stream}; + +struct Component; + +test_programs::p3::export!(Component); + +impl test_programs::p3::exports::wasi::cli::run::Guest for Component { + async fn run() -> Result<(), ()> { + const LEN: usize = 1 << 20; + let body = vec![1; LEN]; + + let addr = test_programs::p3::wasi::cli::environment::get_environment() + .into_iter() + .find_map(|(k, v)| k.eq("HTTP_SERVER").then_some(v)) + .unwrap(); + + let headers = Headers::from_list(&[]).unwrap(); + let (mut contents_tx, contents_rx) = wit_stream::new(); + let (trailers_tx, trailers_rx) = wit_future::new(|| Ok(None)); + drop(trailers_tx); + let options = RequestOptions::new(); + let (request, transmit) = + Request::new(headers, Some(contents_rx), trailers_rx, Some(options)); + request.set_method(&Method::Post).unwrap(); + request.set_scheme(Some(&Scheme::Http)).unwrap(); + request.set_authority(Some(&addr)).unwrap(); + request.set_path_with_query(Some("/post")).unwrap(); + + drop(transmit); + + let ((), echoed) = join!( + async { + let remaining = contents_tx.write_all((&body[..]).into()).await; + assert!(remaining.is_empty()); + drop(contents_tx); + }, + async { + let response = client::send(request).await.context("send failed").unwrap(); + let status = response.get_status_code(); + assert_eq!(status, 200); + let (_, result_rx) = wit_future::new(|| Ok(())); + let (body_rx, _trailers_rx) = Response::consume_body(response, result_rx); + body_rx.collect().await + }, + ); + + assert_eq!( + echoed.len(), + LEN, + "response body was truncated after dropping the transmit future" + ); + Ok(()) + } +} + +fn main() {} diff --git a/crates/wasi-http/tests/all/p3/mod.rs b/crates/wasi-http/tests/all/p3/mod.rs index 1fc17440f944..3b30567446b9 100644 --- a/crates/wasi-http/tests/all/p3/mod.rs +++ b/crates/wasi-http/tests/all/p3/mod.rs @@ -946,3 +946,9 @@ async fn p3_http_outbound_request_chunk_size() -> Result<()> { let server = Server::http1(1)?; run_cli(P3_HTTP_OUTBOUND_REQUEST_CHUNK_SIZE_COMPONENT, &server).await } + +#[test_log::test(tokio::test(flavor = "multi_thread"))] +async fn p3_http_drop_transmit() -> Result<()> { + let server = Server::http1(1)?; + run_cli(P3_HTTP_DROP_TRANSMIT_COMPONENT, &server).await +} From 044b069463a1d81d6196e8a5ccfc5f8ee7c4f12e Mon Sep 17 00:00:00 2001 From: Adam Bratschi-Kaye Date: Fri, 28 Aug 2026 20:12:30 +0000 Subject: [PATCH 2/2] Revert "Remove `BodyWithState`" This reverts commit e13936abdab766f67f5063da092f9330db12c4ad. --- crates/wasi-http/src/p3/body.rs | 43 +++++++++++++++++++++++++ crates/wasi-http/src/p3/host/handler.rs | 29 +++++++++++------ 2 files changed, 63 insertions(+), 9 deletions(-) diff --git a/crates/wasi-http/src/p3/body.rs b/crates/wasi-http/src/p3/body.rs index 89b43380f128..4712aa9905f7 100644 --- a/crates/wasi-http/src/p3/body.rs +++ b/crates/wasi-http/src/p3/body.rs @@ -585,6 +585,39 @@ where } } +/// A wrapper around [http_body::Body], which allows attaching arbitrary state to it +pub(crate) struct BodyWithState { + body: T, + _state: U, +} + +impl http_body::Body for BodyWithState +where + T: http_body::Body + Unpin, + U: Unpin, +{ + type Data = T::Data; + type Error = T::Error; + + #[inline] + fn poll_frame( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + Pin::new(&mut self.get_mut().body).poll_frame(cx) + } + + #[inline] + fn is_end_stream(&self) -> bool { + self.body.is_end_stream() + } + + #[inline] + fn size_hint(&self) -> http_body::SizeHint { + self.body.size_hint() + } +} + /// A wrapper around [http_body::Body], which validates `Content-Length` pub(crate) struct BodyWithContentLength { body: T, @@ -667,6 +700,16 @@ where } pub(crate) trait BodyExt { + fn with_state(self, state: T) -> BodyWithState + where + Self: Sized, + { + BodyWithState { + body: self, + _state: state, + } + } + fn with_content_length( self, limit: u64, diff --git a/crates/wasi-http/src/p3/host/handler.rs b/crates/wasi-http/src/p3/host/handler.rs index d514eb5121bb..3b3fc4251685 100644 --- a/crates/wasi-http/src/p3/host/handler.rs +++ b/crates/wasi-http/src/p3/host/handler.rs @@ -1,10 +1,12 @@ use crate::FieldMap; use crate::p3::bindings::http::client::{Host, HostWithStore}; use crate::p3::bindings::http::types::{Request, Response}; -use crate::p3::body::Body; +use crate::p3::body::{Body, BodyExt as _}; use crate::p3::{HttpError, HttpResult}; use crate::{Error, WasiHttp, WasiHttpCtxView}; use core::task::{Context, Poll, Waker}; +use http_body_util::BodyExt as _; +use std::sync::Arc; use tokio::sync::oneshot; use tokio::task::{self, JoinHandle}; use tracing::debug; @@ -26,7 +28,7 @@ const DROPPED_FUTURE_ERROR: &str = async fn io_task_result( rx: oneshot::Receiver<( - Option, + Option>, oneshot::Receiver>, )>, ) -> Result<(), Error> { @@ -41,7 +43,7 @@ async fn io_task_result( fn send_dummy_io( result: Result<(), Error>, io_result_tx: oneshot::Sender<( - Option, + Option>, oneshot::Receiver>, )>, ) { @@ -54,7 +56,7 @@ fn send_dummy_io_err( store: &Accessor, e: Error, io_result_tx: oneshot::Sender<( - Option, + Option>, oneshot::Receiver>, )>, ) -> HttpError { @@ -68,6 +70,10 @@ impl HostWithStore for WasiHttp { store: &Accessor, req: Resource, ) -> HttpResult> { + // A handle to the I/O task, if spawned, will be sent on this channel + // and kept as part of request body state + let (io_task_tx, io_task_rx) = oneshot::channel(); + // A handle to the I/O task, if spawned, will be sent on this channel // along with the result receiver let (io_result_tx, io_result_rx) = oneshot::channel(); @@ -85,7 +91,9 @@ impl HostWithStore for WasiHttp { let (req, options) = req.into_http_with_getter(&mut store, io_task_result(io_result_rx), getter)?; HttpResult::Ok(store.get().hooks.send_request( - req, + // Attach a reference to the io task to the body so that it + // isn't cancelled if the body is dropped. + req.map(|body| body.with_state(io_task_rx).boxed_unsync()), options.as_deref().copied(), Box::new(async { // Forward the response processing result to `WasiHttpCtx` implementation @@ -134,13 +142,16 @@ impl HostWithStore for WasiHttp { Poll::Pending => { // I/O driver still needs to be polled, spawn a task and send handles to it let (tx, rx) = oneshot::channel(); - let io = AbortOnDropJoinHandle(task::spawn(async move { + let io = Arc::new(AbortOnDropJoinHandle(task::spawn(async move { let res = io.await; debug!(?res, "`send_request` I/O future finished"); _ = tx.send(res); - })); - _ = io_result_tx.send((Some(io), rx)); - body + }))); + _ = io_result_tx.send((Some(Arc::clone(&io)), rx)); + _ = io_task_tx.send(Arc::clone(&io)); + // Attach a reference to the io task to the body so that it + // isn't cancelled if the body is dropped. + body.with_state(io).boxed_unsync() } }; store.with(|mut store| {