wasi-http: Ensure buffered body reads aren't dropped when guests pass a host-provided body

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<Response, ErrorCode> {
+        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() {}

crates/wasi-http/src/p3/body.rs

@@ -597,13 +597,67 @@ where
         Poll::Ready(Ok(StreamResult::Dropped))
     }
 
-    fn try_into(me: Pin<Box<Self>>, ty: TypeId) -> Result<Box<dyn Any>, Pin<Box<Self>>> {
+    fn try_into(
+        me: Pin<Box<Self>>,
+        buffer: Bytes,
+        ty: TypeId,
+    ) -> Result<Box<dyn Any>, (Pin<Box<Self>>, Bytes)> {
         if ty == TypeId::of::<Self>() {
-            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<T> {
+    prefix: Bytes,
+    body: T,
+}
+
+impl<T> http_body::Body for BodyWithPrefix<T>
+where
+    T: http_body::Body<Data = Bytes> + Unpin,
+{
+    type Data = Bytes;
+    type Error = T::Error;
+
+    fn poll_frame(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Result<http_body::Frame<Self::Data>, 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
     }
 }
 

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<Collected<Bytes>> {
+        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<Bytes>) {
+        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(())
+    }
+}

crates/wasmtime/src/runtime/component/concurrent/futures_and_streams.rs

@@ -735,10 +735,20 @@ pub trait StreamProducer<D>: Send + 'static {
     /// Attempt to convert the specified object into a `Box<dyn Any>` 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<Box<Self>>, _ty: TypeId) -> Result<Box<dyn Any>, Pin<Box<Self>>> {
-        Err(me)
+    fn try_into(
+        me: Pin<Box<Self>>,
+        buffer: Self::Buffer,
+        _ty: TypeId,
+    ) -> Result<Box<dyn Any>, (Pin<Box<Self>>, Self::Buffer)> {
+        Err((me, buffer))
     }
 }
 
@@ -2826,9 +2836,9 @@ impl<T> 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
                 }