diff --git a/crates/test-programs/src/bin/p3_http_forward_body.rs b/crates/test-programs/src/bin/p3_http_forward_body.rs new file mode 100644 index 000000000000..1d467693e101 --- /dev/null +++ b/crates/test-programs/src/bin/p3_http_forward_body.rs @@ -0,0 +1,34 @@ +use test_programs::p3::{ + service::exports::wasi::http::handler::Guest as Handler, + wasi::http::types::{ErrorCode, Fields, Request, Response}, + wit_future, +}; + +struct Component; + +test_programs::p3::service::export!(Component); + +impl Handler for Component { + /// Forward the request body as the response body, optionally reading from + /// it or replacing its trailers first, depending on the `x-mode` header. + async fn handle(request: Request) -> Result { + let mode = request + .get_headers() + .get("x-mode") + .pop() + .unwrap_or_default(); + let (mut body, mut trailers) = Request::consume_body(request, wit_future::new(|| Ok(())).1); + match &mode[..] { + b"read-0" => _ = body.read(Vec::new()).await, + b"read-1" => _ = body.read(Vec::with_capacity(1)).await, + b"trailers-error" => { + trailers = wit_future::new(|| Err(ErrorCode::InternalError(None))).1; + } + _ => {} + } + Ok(Response::new(Fields::new(), Some(body), trailers).0) + } +} + +// Unused function; required since this file is built as a `bin`: +fn main() {} diff --git a/crates/wasi-http/src/p3/body.rs b/crates/wasi-http/src/p3/body.rs index f6cb899c6a87..a9f5f39a4630 100644 --- a/crates/wasi-http/src/p3/body.rs +++ b/crates/wasi-http/src/p3/body.rs @@ -597,13 +597,67 @@ where Poll::Ready(Ok(StreamResult::Dropped)) } - fn try_into(me: Pin>, ty: TypeId) -> Result, Pin>> { + fn try_into( + me: Pin>, + buffer: Bytes, + ty: TypeId, + ) -> Result, (Pin>, Bytes)> { if ty == TypeId::of::() { - let me = Pin::into_inner(me); + let mut me = Pin::into_inner(me); + if !buffer.is_empty() { + // Data buffered by a previous read has already been taken from + // `body`, so it must be yielded before the rest of `body`. + let body = mem::take(&mut me.body); + me.body = UnsyncBoxBody::new(BodyWithPrefix { + prefix: buffer, + body, + }); + } Ok(me) } else { - Err(me) + Err((me, buffer)) + } + } +} + +/// A wrapper around [http_body::Body], which yields `prefix` before `body` +struct BodyWithPrefix { + prefix: Bytes, + body: T, +} + +impl http_body::Body for BodyWithPrefix +where + T: http_body::Body + Unpin, +{ + type Data = Bytes; + type Error = T::Error; + + fn poll_frame( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + let this = self.get_mut(); + if !this.prefix.is_empty() { + let prefix = std::mem::take(&mut this.prefix); + return Poll::Ready(Some(Ok(http_body::Frame::data(prefix)))); + } + Pin::new(&mut this.body).poll_frame(cx) + } + + fn is_end_stream(&self) -> bool { + self.prefix.is_empty() && self.body.is_end_stream() + } + + fn size_hint(&self) -> http_body::SizeHint { + let n = self.prefix.len() as u64; + let hint = self.body.size_hint(); + let mut res = http_body::SizeHint::new(); + res.set_lower(hint.lower().saturating_add(n)); + if let Some(upper) = hint.upper() { + res.set_upper(upper + n); } + res } } diff --git a/crates/wasi-http/tests/all/p3/mod.rs b/crates/wasi-http/tests/all/p3/mod.rs index 981cd4924993..9f436f5c4f8a 100644 --- a/crates/wasi-http/tests/all/p3/mod.rs +++ b/crates/wasi-http/tests/all/p3/mod.rs @@ -1257,3 +1257,62 @@ async fn p3_http_named_imports() -> Result<()> { ) .await } + +/// A guest forwarding a host body it has partly read, or with its own +/// trailers, as its response body. +mod p3_http_forward_body { + use super::*; + + async fn forward(mode: &str) -> Result> { + let mut trailers = http::HeaderMap::new(); + trailers.insert("x-host", HeaderValue::from_static("host")); + // Response body consists of "hello world" two frames followed by a + // trailer. + let frames = [ + http_body::Frame::data(Bytes::from_static(b"hello ")), + http_body::Frame::data(Bytes::from_static(b"world")), + http_body::Frame::trailers(trailers), + ]; + let req = http::Request::builder() + .uri("http://localhost/") + .header("x-mode", mode) + .body(http_body_util::StreamBody::new(futures::stream::iter( + frames.map(Ok::<_, ErrorCode>), + )))?; + let (tx, _) = oneshot::channel(); + let res = run_http(P3_HTTP_FORWARD_BODY_COMPONENT, req, tx) + .await? + .unwrap(); + Ok(res.into_body()) + } + + fn assert_host_trailers(body: &Collected) { + let trailers = body.trailers().expect("missing trailers"); + assert_eq!(trailers.len(), 1, "{trailers:?}"); + assert_eq!(trailers["x-host"], "host"); + } + + #[test_log::test(tokio::test(flavor = "multi_thread"))] + async fn read_0() -> Result<()> { + let body = forward("read-0").await?; + assert_host_trailers(&body); + assert_eq!(body.to_bytes(), "hello world"); + Ok(()) + } + + #[test_log::test(tokio::test(flavor = "multi_thread"))] + async fn read_1() -> Result<()> { + let body = forward("read-1").await?; + assert_host_trailers(&body); + assert_eq!(body.to_bytes(), "ello world"); + Ok(()) + } + + #[ignore = "Enable when https://github.com/bytecodealliance/wasmtime/issues/15026 has been fixed."] + #[test_log::test(tokio::test(flavor = "multi_thread"))] + async fn trailers_error() -> Result<()> { + let res = forward("trailers-error").await; + assert!(res.is_err(), "{res:?}"); + Ok(()) + } +} diff --git a/crates/wasmtime/src/runtime/component/concurrent/futures_and_streams.rs b/crates/wasmtime/src/runtime/component/concurrent/futures_and_streams.rs index c0de956b6718..a25bdf3f438a 100644 --- a/crates/wasmtime/src/runtime/component/concurrent/futures_and_streams.rs +++ b/crates/wasmtime/src/runtime/component/concurrent/futures_and_streams.rs @@ -735,10 +735,20 @@ pub trait StreamProducer: Send + 'static { /// Attempt to convert the specified object into a `Box` which may /// be downcast to the specified type. /// + /// `buffer` contains any items which this producer previously passed to + /// [`Destination::set_buffer`] but which have not yet been delivered to + /// the reader. If the implementation returns `Ok(_)`, it must ensure those + /// items are retained by the returned value, since they are otherwise + /// lost. + /// /// The implementation must ensure that, if it returns `Ok(_)`, a downcast /// to the specified type is guaranteed to succeed. - fn try_into(me: Pin>, _ty: TypeId) -> Result, Pin>> { - Err(me) + fn try_into( + me: Pin>, + buffer: Self::Buffer, + _ty: TypeId, + ) -> Result, (Pin>, Self::Buffer)> { + Err((me, buffer)) } } @@ -2826,9 +2836,9 @@ impl StoreContextMut<'_, T> { }); let try_into = Box::new(move |ty| { let (mine, buffer) = producer.try_lock().ok()?.take()?; - match P::try_into(mine, ty) { + match P::try_into(mine, buffer, ty) { Ok(value) => Some(value), - Err(mine) => { + Err((mine, buffer)) => { *producer.try_lock().ok()? = Some((mine, buffer)); None }