Some improvements to the wasi-http client implementation of write. (#6161)

* Improve write implementation for streams

* Add trailers implementation for responses.

* Improve tests.

* Update tests.

crates/wasi-http/src/http_impl.rs

@@ -8,6 +8,7 @@ use bytes::{BufMut, Bytes, BytesMut};
 use http_body_util::{BodyExt, Full};
 use hyper::Method;
 use hyper::Request;
+use std::collections::HashMap;
 #[cfg(not(any(target_arch = "riscv64", target_arch = "s390x")))]
 use std::sync::Arc;
 use std::time::Duration;
@@ -201,6 +202,23 @@ impl WasiHttp {
             if let Some(chunk) = frame.data_ref() {
                 buf.put(chunk.clone());
             }
+            if let Some(trailers) = frame.trailers_ref() {
+                response.trailers = self.fields_id_base;
+                self.fields_id_base += 1;
+                let mut map: HashMap<String, Vec<String>> = HashMap::new();
+                for (name, value) in trailers.iter() {
+                    let key = name.to_string();
+                    match map.get_mut(&key) {
+                        Some(vec) => vec.push(value.to_str()?.to_string()),
+                        None => {
+                            let mut vec = Vec::new();
+                            vec.push(value.to_str()?.to_string());
+                            map.insert(key, vec);
+                        }
+                    };
+                }
+                self.fields.insert(response.trailers, map);
+            }
         }
         response.body = self.streams_id_base;
         self.streams_id_base = self.streams_id_base + 1;

crates/wasi-http/src/streams_impl.rs

@@ -2,6 +2,7 @@ use crate::poll::Pollable;
 use crate::streams::{InputStream, OutputStream, StreamError};
 use crate::WasiHttp;
 use anyhow::{anyhow, bail};
+use bytes::BufMut;
 use std::vec::Vec;
 
 impl crate::streams::Host for WasiHttp {
@@ -27,10 +28,23 @@ impl crate::streams::Host for WasiHttp {
 
     fn skip(
         &mut self,
-        _this: InputStream,
-        _len: u64,
+        stream: InputStream,
+        len: u64,
     ) -> wasmtime::Result<Result<(u64, bool), StreamError>> {
-        bail!("unimplemented: skip");
+        let s = self
+            .streams
+            .get_mut(&stream)
+            .ok_or_else(|| anyhow!("stream not found: {stream}"))?;
+        if len == 0 {
+            Ok(Ok((0, s.len() > 0)))
+        } else if s.len() > len.try_into()? {
+            s.truncate(len.try_into()?);
+            Ok(Ok((len, false)))
+        } else {
+            let bytes = s.len();
+            s.truncate(s.len());
+            Ok(Ok((bytes.try_into()?, true)))
+        }
     }
 
     fn subscribe_to_input_stream(&mut self, _this: InputStream) -> wasmtime::Result<Pollable> {
@@ -38,11 +52,7 @@ impl crate::streams::Host for WasiHttp {
     }
 
     fn drop_input_stream(&mut self, stream: InputStream) -> wasmtime::Result<()> {
-        let r = self
-            .streams
-            .get_mut(&stream)
-            .ok_or_else(|| anyhow!("no such input-stream {stream}"))?;
-        r.truncate(0);
+        self.streams.remove(&stream);
         Ok(())
     }
 
@@ -51,17 +61,32 @@ impl crate::streams::Host for WasiHttp {
         this: OutputStream,
         buf: Vec<u8>,
     ) -> wasmtime::Result<Result<u64, StreamError>> {
-        // TODO: Make this a real write not a replace.
-        self.streams.insert(this, bytes::Bytes::from(buf.clone()));
+        match self.streams.get(&this) {
+            Some(data) => {
+                let mut new = bytes::BytesMut::with_capacity(data.len() + buf.len());
+                new.put(data.clone());
+                new.put(bytes::Bytes::from(buf.clone()));
+                self.streams.insert(this, new.freeze());
+            }
+            None => {
+                self.streams.insert(this, bytes::Bytes::from(buf.clone()));
+            }
+        }
         Ok(Ok(buf.len().try_into()?))
     }
 
     fn write_zeroes(
         &mut self,
-        _this: OutputStream,
-        _len: u64,
+        this: OutputStream,
+        len: u64,
     ) -> wasmtime::Result<Result<u64, StreamError>> {
-        bail!("unimplemented: write_zeroes");
+        let mut data = Vec::with_capacity(len.try_into()?);
+        let mut i = 0;
+        while i < len {
+            data.push(0);
+            i = i + 1;
+        }
+        self.write(this, data)
     }
 
     fn splice(
@@ -85,9 +110,8 @@ impl crate::streams::Host for WasiHttp {
         bail!("unimplemented: subscribe_to_output_stream");
     }
 
-    fn drop_output_stream(&mut self, _this: OutputStream) -> wasmtime::Result<()> {
-        //bail!("unimplemented: drop_output_stream");
-        //FIXME: intentionally ignoring
+    fn drop_output_stream(&mut self, stream: OutputStream) -> wasmtime::Result<()> {
+        self.streams.remove(&stream);
         Ok(())
     }
 }

crates/wasi-http/src/struct.rs

@@ -34,6 +34,7 @@ pub struct ActiveResponse {
     pub status: u16,
     pub body: u32,
     pub response_headers: HashMap<String, Vec<String>>,
+    pub trailers: u32,
 }
 
 impl ActiveRequest {
@@ -60,6 +61,7 @@ impl ActiveResponse {
             status: 0,
             body: 0,
             response_headers: HashMap::new(),
+            trailers: 0,
         }
     }
 }

crates/wasi-http/src/types_impl.rs

@@ -101,8 +101,16 @@ impl crate::types::Host for WasiHttp {
         self.fields.insert(id, m.clone());
         Ok(id)
     }
-    fn finish_incoming_stream(&mut self, _s: IncomingStream) -> wasmtime::Result<Option<Trailers>> {
-        bail!("unimplemented: finish_incoming_stream")
+    fn finish_incoming_stream(&mut self, s: IncomingStream) -> wasmtime::Result<Option<Trailers>> {
+        for (_, value) in self.responses.iter() {
+            if value.body == s {
+                return match value.trailers {
+                    0 => Ok(None),
+                    _ => Ok(Some(value.trailers)),
+                };
+            }
+        }
+        bail!("unknown stream!")
     }
     fn finish_outgoing_stream(
         &mut self,
@@ -181,8 +189,10 @@ impl crate::types::Host for WasiHttp {
             .requests
             .get_mut(&request)
             .ok_or_else(|| anyhow!("unknown request: {request}"))?;
-        req.body = self.streams_id_base;
-        self.streams_id_base = self.streams_id_base + 1;
+        if req.body == 0 {
+            req.body = self.streams_id_base;
+            self.streams_id_base = self.streams_id_base + 1;
+        }
         Ok(Ok(req.body))
     }
     fn drop_response_outparam(&mut self, _response: ResponseOutparam) -> wasmtime::Result<()> {