Repository: bytecodealliance/wasmtime
Pull request #7412 (2023-10-30 20:06:46)

Add `wasi_http_double_echo` test
This tests streaming outgoing request bodies, which we didn't previously have coverage for.


Patch:
diff --git a/crates/wasi-http/src/body.rs b/crates/wasi-http/src/body.rs
index 30d217d57..00a3cc50a 100644
--- a/crates/wasi-http/src/body.rs
+++ b/crates/wasi-http/src/body.rs
@@ -20,6 +20,7 @@ pub type HyperIncomingBody = BoxBody<Bytes, anyhow::Error>;
 /// want to do that unless the user asks to consume the body.
 pub struct HostIncomingBodyBuilder {
     pub body: HyperIncomingBody,
+    pub worker: Option<Arc<AbortOnDropJoinHandle<anyhow::Result<()>>>>,
     pub between_bytes_timeout: Duration,
 }
 
@@ -100,6 +101,10 @@ impl HostIncomingBodyBuilder {
                 // loop running to relieve backpressure, so we get to the trailers.
                 let _ = body_writer.send(Ok(data)).await;
             }
+
+            // Force the builder's worker to be owned by this task, ensuring that it's kept around
+            // until this task exits.
+            drop(self.worker);
         });
 
         HostIncomingBody {
diff --git a/crates/wasi-http/src/types.rs b/crates/wasi-http/src/types.rs
index 945965781..36070d5e4 100644
--- a/crates/wasi-http/src/types.rs
+++ b/crates/wasi-http/src/types.rs
@@ -9,6 +9,7 @@ use anyhow::Context;
 use http_body_util::BodyExt;
 use hyper::header::HeaderName;
 use std::any::Any;
+use std::sync::Arc;
 use std::time::Duration;
 use tokio::net::TcpStream;
 use tokio::time::timeout;
@@ -38,6 +39,7 @@ pub trait WasiHttpView: Send {
         let (parts, body) = req.into_parts();
         let body = HostIncomingBodyBuilder {
             body,
+            worker: None,
             // TODO: this needs to be plumbed through
             between_bytes_timeout: std::time::Duration::from_millis(600 * 1000),
         };
@@ -162,7 +164,7 @@ pub fn default_send_request(
 
         Ok(IncomingResponseInternal {
             resp,
-            worker,
+            worker: Arc::new(worker),
             between_bytes_timeout,
         })
     });
@@ -273,7 +275,7 @@ pub struct HostIncomingResponse {
     pub status: u16,
     pub headers: FieldMap,
     pub body: Option<HostIncomingBodyBuilder>,
-    pub worker: AbortOnDropJoinHandle<anyhow::Result<()>>,
+    pub worker: Arc<AbortOnDropJoinHandle<anyhow::Result<()>>>,
 }
 
 pub struct HostOutgoingResponse {
@@ -324,7 +326,7 @@ pub enum HostFields {
 
 pub struct IncomingResponseInternal {
     pub resp: hyper::Response<HyperIncomingBody>,
-    pub worker: AbortOnDropJoinHandle<anyhow::Result<()>>,
+    pub worker: Arc<AbortOnDropJoinHandle<anyhow::Result<()>>>,
     pub between_bytes_timeout: std::time::Duration,
 }
 
diff --git a/crates/wasi-http/src/types_impl.rs b/crates/wasi-http/src/types_impl.rs
index c08c23892..422ecde2c 100644
--- a/crates/wasi-http/src/types_impl.rs
+++ b/crates/wasi-http/src/types_impl.rs
@@ -13,7 +13,7 @@ use crate::{
 };
 use anyhow::Context;
 use hyper::header::HeaderName;
-use std::any::Any;
+use std::{any::Any, sync::Arc};
 use wasmtime::component::Resource;
 use wasmtime_wasi::preview2::{
     bindings::io::streams::{InputStream, OutputStream},
@@ -744,6 +744,7 @@ impl<T: WasiHttpView> crate::bindings::http::types::HostFutureIncomingResponse f
             headers: FieldMap::from(parts.headers),
             body: Some(HostIncomingBodyBuilder {
                 body,
+                worker: Some(Arc::clone(&resp.worker)),
                 between_bytes_timeout: resp.between_bytes_timeout,
             }),
             worker: resp.worker,


Test patch:
diff --git a/crates/test-programs/src/bin/api_proxy_streaming.rs b/crates/test-programs/src/bin/api_proxy_streaming.rs
index c2cde5c15..a5f59b4bb 100644
--- a/crates/test-programs/src/bin/api_proxy_streaming.rs
+++ b/crates/test-programs/src/bin/api_proxy_streaming.rs
@@ -1,9 +1,9 @@
 use anyhow::{bail, Result};
 use bindings::wasi::http::types::{
-    Fields, IncomingRequest, Method, OutgoingBody, OutgoingRequest, OutgoingResponse,
-    ResponseOutparam, Scheme,
+    Fields, IncomingRequest, IncomingResponse, Method, OutgoingBody, OutgoingRequest,
+    OutgoingResponse, ResponseOutparam, Scheme,
 };
-use futures::{stream, SinkExt, StreamExt, TryStreamExt};
+use futures::{future, stream, Future, SinkExt, StreamExt, TryStreamExt};
 use url::Url;
 
 mod bindings {
@@ -35,6 +35,9 @@ async fn handle_request(request: IncomingRequest, response_out: ResponseOutparam
 
     match (request.method(), request.path_with_query().as_deref()) {
         (Method::Get, Some("/hash-all")) => {
+            // Send outgoing GET requests to the specified URLs and stream the hashes of the response bodies as
+            // they arrive.
+
             let urls = headers.iter().filter_map(|(k, v)| {
                 (k == "url")
                     .then_some(v)
@@ -73,6 +76,8 @@ async fn handle_request(request: IncomingRequest, response_out: ResponseOutparam
         }
 
         (Method::Post, Some("/echo")) => {
+            // Echo the request body without buffering it.
+
             let response = OutgoingResponse::new(
                 200,
                 Fields::from_list(
@@ -108,16 +113,142 @@ async fn handle_request(request: IncomingRequest, response_out: ResponseOutparam
             }
         }
 
-        _ => {
-            let response = OutgoingResponse::new(405, Fields::new());
+        (Method::Post, Some("/double-echo")) => {
+            // Pipe the request body to an outgoing request and stream the response back to the client.
 
-            let body = response.body().expect("response should be writable");
+            if let Some(url) = headers.iter().find_map(|(k, v)| {
+                (k == "url")
+                    .then_some(v)
+                    .and_then(|v| std::str::from_utf8(v).ok())
+                    .and_then(|v| Url::parse(v).ok())
+            }) {
+                match double_echo(request, &url).await {
+                    Ok((request_copy, response)) => {
+                        let mut stream = executor::incoming_body(
+                            response.consume().expect("response should be consumable"),
+                        );
+
+                        let response = OutgoingResponse::new(
+                            200,
+                            Fields::from_list(
+                                &headers
+                                    .into_iter()
+                                    .filter_map(|(k, v)| (k == "content-type").then_some((k, v)))
+                                    .collect::<Vec<_>>(),
+                            )
+                            .unwrap(),
+                        );
+
+                        let mut body = executor::outgoing_body(
+                            response.body().expect("response should be writable"),
+                        );
+
+                        ResponseOutparam::set(response_out, Ok(response));
+
+                        let response_copy = async move {
+                            while let Some(chunk) = stream.next().await {
+                                body.send(chunk?).await?;
+                            }
+                            Ok::<_, anyhow::Error>(())
+                        };
 
-            ResponseOutparam::set(response_out, Ok(response));
+                        let (request_copy, response_copy) =
+                            future::join(request_copy, response_copy).await;
+                        if let Err(e) = request_copy.and(response_copy) {
+                            eprintln!("error piping to and from {url}: {e}");
+                        }
+                    }
+
+                    Err(e) => {
+                        eprintln!("Error sending outgoing request to {url}: {e}");
+                        server_error(response_out);
+                    }
+                }
+            } else {
+                bad_request(response_out);
+            }
+        }
+
+        _ => method_not_allowed(response_out),
+    }
+}
+
+async fn double_echo(
+    incoming_request: IncomingRequest,
+    url: &Url,
+) -> Result<(impl Future<Output = Result<()>>, IncomingResponse)> {
+    let outgoing_request = OutgoingRequest::new(
+        &Method::Post,
+        Some(url.path()),
+        Some(&match url.scheme() {
+            "http" => Scheme::Http,
+            "https" => Scheme::Https,
+            scheme => Scheme::Other(scheme.into()),
+        }),
+        Some(&format!(
+            "{}{}",
+            url.host_str().unwrap_or(""),
+            if let Some(port) = url.port() {
+                format!(":{port}")
+            } else {
+                String::new()
+            }
+        )),
+        Fields::new(),
+    );
+
+    let mut body = executor::outgoing_body(
+        outgoing_request
+            .body()
+            .expect("request body should be writable"),
+    );
+
+    let response = executor::outgoing_request_send(outgoing_request);
+
+    let mut stream = executor::incoming_body(
+        incoming_request
+            .consume()
+            .expect("request should be consumable"),
+    );
 
-            OutgoingBody::finish(body, None);
+    let copy = async move {
+        while let Some(chunk) = stream.next().await {
+            body.send(chunk?).await?;
         }
+        Ok::<_, anyhow::Error>(())
+    };
+
+    let response = response.await?;
+
+    let status = response.status();
+
+    if !(200..300).contains(&status) {
+        bail!("unexpected status: {status}");
     }
+
+    Ok((copy, response))
+}
+
+fn server_error(response_out: ResponseOutparam) {
+    respond(500, response_out)
+}
+
+fn bad_request(response_out: ResponseOutparam) {
+    respond(400, response_out)
+}
+
+fn method_not_allowed(response_out: ResponseOutparam) {
+    respond(405, response_out)
+}
+
+fn respond(status: u16, response_out: ResponseOutparam) {
+    let response = OutgoingResponse::new(status, Fields::new());
+
+    let body = response.body().expect("response should be writable");
+
+    ResponseOutparam::set(response_out, Ok(response));
+
+    OutgoingBody::finish(body, None);
 }
 
 async fn hash(url: &Url) -> Result<String> {
diff --git a/crates/wasi-http/tests/all/main.rs b/crates/wasi-http/tests/all/main.rs
index 84e3ef12d..9f8386b5e 100644
--- a/crates/wasi-http/tests/all/main.rs
+++ b/crates/wasi-http/tests/all/main.rs
@@ -2,7 +2,7 @@ use crate::http_server::Server;
 use anyhow::{anyhow, Context, Result};
 use futures::{channel::oneshot, future, stream, FutureExt};
 use http_body::Frame;
-use http_body_util::{combinators::BoxBody, Collected, StreamBody};
+use http_body_util::{combinators::BoxBody, Collected, Empty, StreamBody};
 use hyper::{body::Bytes, server::conn::http1, service::service_fn, Method, StatusCode};
 use sha2::{Digest, Sha256};
 use std::{collections::HashMap, iter, net::Ipv4Addr, str, sync::Arc};
@@ -268,7 +268,7 @@ async fn do_wasi_http_hash_all(override_send_request: bool) -> Result<()> {
                 Ok(view.table().push(HostFutureIncomingResponse::Ready(
                     handle(request.into_parts().0).map(|resp| IncomingResponseInternal {
                         resp,
-                        worker: preview2::spawn(future::ready(Ok(()))),
+                        worker: Arc::new(preview2::spawn(future::ready(Ok(())))),
                         between_bytes_timeout,
                     }),
                 ))?)
@@ -360,6 +360,73 @@ async fn do_wasi_http_hash_all(override_send_request: bool) -> Result<()> {
 
 #[test_log::test(tokio::test)]
 async fn wasi_http_echo() -> Result<()> {
+    do_wasi_http_echo("echo", None).await
+}
+
+#[test_log::test(tokio::test)]
+async fn wasi_http_double_echo() -> Result<()> {
+    let listener = tokio::net::TcpListener::bind((Ipv4Addr::new(127, 0, 0, 1), 0)).await?;
+
+    let prefix = format!("http://{}", listener.local_addr()?);
+
+    let (_tx, rx) = oneshot::channel::<()>();
+
+    let server = async move {
+        loop {
+            let (stream, _) = listener.accept().await?;
+            task::spawn(async move {
+                if let Err(e) = http1::Builder::new()
+                    .keep_alive(true)
+                    .serve_connection(
+                        stream,
+                        service_fn(
+                            move |request: hyper::Request<hyper::body::Incoming>| async move {
+                                use http_body_util::BodyExt;
+
+                                if let (&Method::POST, "/echo") =
+                                    (request.method(), request.uri().path())
+                                {
+                                    Ok::<_, anyhow::Error>(hyper::Response::new(
+                                        request.into_body().boxed(),
+                                    ))
+                                } else {
+                                    Ok(hyper::Response::builder()
+                                        .status(StatusCode::METHOD_NOT_ALLOWED)
+                                        .body(BoxBody::new(
+                                            Empty::new().map_err(|_| unreachable!()),
+                                        ))?)
+                                }
+                            },
+                        ),
+                    )
+                    .await
+                {
+                    eprintln!("error serving connection: {e:?}");
+                }
+            });
+
+            // Help rustc with type inference:
+            if false {
+                return Ok::<_, anyhow::Error>(());
+            }
+        }
+    }
+    .then(|result| {
+        if let Err(e) = result {
+            eprintln!("error listening for connections: {e:?}");
+        }
+        future::ready(())
+    })
+    .boxed();
+
+    task::spawn(async move {
+        drop(future::select(server, rx).await);
+    });
+
+    do_wasi_http_echo("double-echo", Some(&format!("{prefix}/echo"))).await
+}
+
+async fn do_wasi_http_echo(uri: &str, url_header: Option<&str>) -> Result<()> {
     let body = {
         // A sorta-random-ish megabyte
         let mut n = 0_u8;
@@ -371,13 +438,18 @@ async fn wasi_http_echo() -> Result<()> {
         .collect::<Vec<_>>()
     };
 
-    let request = hyper::Request::post("/echo")
-        .header("content-type", "application/octet-stream")
-        .body(BoxBody::new(StreamBody::new(stream::iter(
-            body.chunks(16 * 1024)
-                .map(|chunk| Ok::<_, anyhow::Error>(Frame::data(Bytes::copy_from_slice(chunk))))
-                .collect::<Vec<_>>(),
-        ))))?;
+    let mut request =
+        hyper::Request::post(&format!("/{uri}")).header("content-type", "application/octet-stream");
+
+    if let Some(url_header) = url_header {
+        request = request.header("url", url_header);
+    }
+
+    let request = request.body(BoxBody::new(StreamBody::new(stream::iter(
+        body.chunks(16 * 1024)
+            .map(|chunk| Ok::<_, anyhow::Error>(Frame::data(Bytes::copy_from_slice(chunk))))
+            .collect::<Vec<_>>(),
+    ))))?;
 
     let response = run_wasi_http(
         test_programs_artifacts::API_PROXY_STREAMING_COMPONENT,