Repository: bytecodealliance/wasmtime
Pull request #7527 (2023-11-10 22:35:59)

wasi-http: Validate Content-Length when present
Make use of immutable header state to track how much has been written to an output-body stream, ensuring that we write exactly as much as was mentioned in the `Content-Length` header.
<!--
Please make sure you include the following information:

- If this work has been discussed elsewhere, please include a link to that
  conversation. If it was discussed in an issue, just mention "issue #...".

- Explain why this change is needed. If the details are in an issue already,
  this can be brief.

Our development process is documented in the Wasmtime book:
https://docs.wasmtime.dev/contributing-development-process.html

Please ensure all communication follows the code of conduct:
https://github.com/bytecodealliance/wasmtime/blob/main/CODE_OF_CONDUCT.md
-->



Patch:
diff --git a/cranelift/codegen/src/opts.rs b/cranelift/codegen/src/opts.rs
index f7f03a0be..3ec3e1165 100644
--- a/cranelift/codegen/src/opts.rs
+++ b/cranelift/codegen/src/opts.rs
@@ -23,7 +23,16 @@ pub type Range = (usize, usize);
 pub type ValueArray2 = [Value; 2];
 pub type ValueArray3 = [Value; 3];
 
-pub type ConstructorVec<T> = SmallVec<[T; 8]>;
+const MAX_ISLE_RETURNS: usize = 8;
+
+pub type ConstructorVec<T> = SmallVec<[T; MAX_ISLE_RETURNS]>;
+
+impl<T: smallvec::Array> generated_code::Length for SmallVec<T> {
+    #[inline]
+    fn len(&self) -> usize {
+        SmallVec::len(self)
+    }
+}
 
 pub(crate) mod generated_code;
 use generated_code::{ContextIter, IntoContextIter};
diff --git a/cranelift/isle/isle/isle_examples/link/multi_constructor_main.rs b/cranelift/isle/isle/isle_examples/link/multi_constructor_main.rs
index 11396a93b..3c05532ce 100644
--- a/cranelift/isle/isle/isle_examples/link/multi_constructor_main.rs
+++ b/cranelift/isle/isle/isle_examples/link/multi_constructor_main.rs
@@ -3,6 +3,8 @@ use multi_constructor::{ContextIter, IntoContextIter};
 
 struct Context;
 
+const MAX_ISLE_RETURNS: usize = 100;
+
 #[derive(Default)]
 struct It {
     i: u32,
diff --git a/cranelift/isle/isle/isle_examples/link/multi_extractor_main.rs b/cranelift/isle/isle/isle_examples/link/multi_extractor_main.rs
index c7c7dd406..ebfa38e30 100644
--- a/cranelift/isle/isle/isle_examples/link/multi_extractor_main.rs
+++ b/cranelift/isle/isle/isle_examples/link/multi_extractor_main.rs
@@ -2,6 +2,8 @@ mod multi_extractor;
 
 use multi_extractor::{ContextIter, IntoContextIter};
 
+const MAX_ISLE_RETURNS: usize = 100;
+
 #[derive(Clone)]
 pub enum A {
     B,
diff --git a/cranelift/isle/isle/src/codegen.rs b/cranelift/isle/isle/src/codegen.rs
index 1c3dac155..d5eae5a09 100644
--- a/cranelift/isle/isle/src/codegen.rs
+++ b/cranelift/isle/isle/src/codegen.rs
@@ -236,6 +236,16 @@ pub trait IntoContextIter {{
     fn into_context_iter(self) -> Self::IntoIter;
 }}
 
+pub trait Length {{
+    fn len(&self) -> usize;
+}}
+
+impl<T> Length for std::vec::Vec<T> {{
+    fn len(&self) -> usize {{
+        std::vec::Vec::len(self)
+    }}
+}}
+
 pub struct ContextIterWrapper<I, C> {{
     iter: I,
     _ctx: std::marker::PhantomData<C>,
@@ -289,6 +299,11 @@ impl<T, E: Extend<T>, C> Extend<T> for ContextIterWrapper<E, C> {{
     fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {{
         self.iter.extend(iter);
     }}
+}}
+impl<L: Length, C> Length for ContextIterWrapper<L, C> {{
+    fn len(&self) -> usize {{
+        self.iter.len()
+    }}
 }}
            "#,
         )
@@ -406,7 +421,7 @@ impl<T, E: Extend<T>, C> Extend<T> for ContextIterWrapper<E, C> {{
             if let ReturnKind::Iterator = sig.ret_kind {
                 writeln!(
                     ctx.out,
-                    "{}    returns: &mut impl Extend<{}>,",
+                    "{}    returns: &mut (impl Extend<{}> + Length),",
                     &ctx.indent, ret
                 )?;
             }
@@ -635,7 +650,14 @@ impl<T, E: Extend<T>, C> Extend<T> for ContextIterWrapper<E, C> {{
                     match ret_kind {
                         ReturnKind::Plain => writeln!(ctx.out, ";")?,
                         ReturnKind::Option => writeln!(ctx.out, ");")?,
-                        ReturnKind::Iterator => writeln!(ctx.out, "));")?,
+                        ReturnKind::Iterator => {
+                            writeln!(ctx.out, "));")?;
+                            writeln!(
+                                ctx.out,
+                                "{}if returns.len() >= MAX_ISLE_RETURNS {{ return; }}",
+                                ctx.indent
+                            )?;
+                        }
                     }
                 }
             }
diff --git a/crates/wasi-http/src/body.rs b/crates/wasi-http/src/body.rs
index c4ada9ef9..41f730286 100644
--- a/crates/wasi-http/src/body.rs
+++ b/crates/wasi-http/src/body.rs
@@ -1,3 +1,4 @@
+use crate::internal_error;
 use crate::{bindings::http::types, types::FieldMap};
 use anyhow::anyhow;
 use bytes::Bytes;
@@ -366,13 +367,47 @@ pub enum FinishMessage {
     Abort,
 }
 
+#[derive(Clone)]
+struct WrittenState {
+    expected: u64,
+    written: Arc<std::sync::atomic::AtomicU64>,
+}
+
+impl WrittenState {
+    fn new(expected_size: u64) -> Self {
+        Self {
+            expected: expected_size,
+            written: Arc::new(std::sync::atomic::AtomicU64::new(0)),
+        }
+    }
+
+    /// Add `len` to the total number of bytes written. Returns `false` if the new total exceeds
+    /// the number of bytes expected to be written.
+    fn update(&self, len: usize) -> bool {
+        let len = len as u64;
+        let old = self
+            .written
+            .fetch_add(len, std::sync::atomic::Ordering::Relaxed);
+        old + len <= self.expected
+    }
+
+    /// Return a comparison of total bytes written to the number of bytes expected to be written.
+    fn finish(self) -> std::cmp::Ordering {
+        let written = self.written.load(std::sync::atomic::Ordering::Relaxed);
+        written.cmp(&self.expected)
+    }
+}
+
 pub struct HostOutgoingBody {
     pub body_output_stream: Option<Box<dyn HostOutputStream>>,
-    pub finish_sender: Option<tokio::sync::oneshot::Sender<FinishMessage>>,
+    written: Option<WrittenState>,
+    finish_sender: Option<tokio::sync::oneshot::Sender<FinishMessage>>,
 }
 
 impl HostOutgoingBody {
-    pub fn new() -> (Self, HyperOutgoingBody) {
+    pub fn new(size: Option<u64>) -> (Self, HyperOutgoingBody) {
+        let written = size.map(WrittenState::new);
+
         use tokio::sync::oneshot::error::RecvError;
         struct BodyImpl {
             body_receiver: mpsc::Receiver<Bytes>,
@@ -423,31 +458,91 @@ impl HostOutgoingBody {
             finish_receiver: Some(finish_receiver),
         }
         .boxed();
+
+        // TODO: this capacity constant is arbitrary, and should be configurable
+        let output_stream = BodyWriteStream::new(1024 * 1024, body_sender, written.clone());
+
         (
             Self {
-                // TODO: this capacity constant is arbitrary, and should be configurable
-                body_output_stream: Some(Box::new(BodyWriteStream::new(1024 * 1024, body_sender))),
+                body_output_stream: Some(Box::new(output_stream)),
+                written,
                 finish_sender: Some(finish_sender),
             },
             body_impl,
         )
     }
+
+    pub fn finish(mut self, ts: Option<FieldMap>) -> Result<(), types::ErrorCode> {
+        // Make sure that the output stream has been dropped, so that the BodyImpl poll function
+        // will immediately pick up the finish sender.
+        drop(self.body_output_stream);
+
+        let sender = self
+            .finish_sender
+            .take()
+            .expect("outgoing-body trailer_sender consumed by a non-owning function");
+
+        if let Some(w) = self.written {
+            use std::cmp::Ordering;
+            let res = w.finish();
+            if res != Ordering::Equal {
+                let msg = match res {
+                    Ordering::Less => "not enough",
+                    Ordering::Greater => "too much",
+                    Ordering::Equal => unreachable!(),
+                };
+
+                let _ = sender.send(FinishMessage::Abort);
+                return Err(internal_error(format!("{msg} written to body stream")));
+            }
+        }
+
+        let message = if let Some(ts) = ts {
+            FinishMessage::Trailers(ts)
+        } else {
+            FinishMessage::Finished
+        };
+
+        // Ignoring failure: receiver died sending body, but we can't report that here.
+        let _ = sender.send(message.into());
+
+        Ok(())
+    }
+
+    pub fn abort(mut self) {
+        // Make sure that the output stream has been dropped, so that the BodyImpl poll function
+        // will immediately pick up the finish sender.
+        drop(self.body_output_stream);
+
+        let sender = self
+            .finish_sender
+            .take()
+            .expect("outgoing-body trailer_sender consumed by a non-owning function");
+
+        let _ = sender.send(FinishMessage::Abort);
+    }
 }
 
 /// Provides a [`HostOutputStream`] impl from a [`tokio::sync::mpsc::Sender`].
-pub struct BodyWriteStream {
+struct BodyWriteStream {
     writer: mpsc::Sender<Bytes>,
     write_budget: usize,
+    written: Option<WrittenState>,
 }
 
 impl BodyWriteStream {
     /// Create a [`BodyWriteStream`].
-    pub fn new(write_budget: usize, writer: mpsc::Sender<Bytes>) -> Self {
+    fn new(
+        write_budget: usize,
+        writer: mpsc::Sender<Bytes>,
+        written: Option<WrittenState>,
+    ) -> Self {
         // at least one capacity is required to send a message
         assert!(writer.max_capacity() >= 1);
         BodyWriteStream {
             writer,
             write_budget,
+            written,
         }
     }
 }
@@ -455,10 +550,21 @@ impl BodyWriteStream {
 #[async_trait::async_trait]
 impl HostOutputStream for BodyWriteStream {
     fn write(&mut self, bytes: Bytes) -> Result<(), StreamError> {
+        let len = bytes.len();
         match self.writer.try_send(bytes) {
             // If the message was sent then it's queued up now in hyper to get
             // received.
-            Ok(()) => Ok(()),
+            Ok(()) => {
+                if let Some(written) = self.written.as_ref() {
+                    if !written.update(len) {
+                        return Err(StreamError::LastOperationFailed(anyhow!(internal_error(
+                            "too much written to output stream".to_owned()
+                        ))));
+                    }
+                }
+
+                Ok(())
+            }
 
             // If this channel is full then that means `check_write` wasn't
             // called. The call to `check_write` always guarantees that there's
diff --git a/crates/wasi-http/src/lib.rs b/crates/wasi-http/src/lib.rs
index 486c6b776..9cd3fe1d3 100644
--- a/crates/wasi-http/src/lib.rs
+++ b/crates/wasi-http/src/lib.rs
@@ -43,3 +43,7 @@ pub(crate) fn dns_error(rcode: String, info_code: u16) -> bindings::http::types:
         info_code: Some(info_code),
     })
 }
+
+pub(crate) fn internal_error(msg: String) -> bindings::http::types::ErrorCode {
+    bindings::http::types::ErrorCode::InternalError(Some(msg))
+}
diff --git a/crates/wasi-http/src/types_impl.rs b/crates/wasi-http/src/types_impl.rs
index cf35b5b3b..18c1b0457 100644
--- a/crates/wasi-http/src/types_impl.rs
+++ b/crates/wasi-http/src/types_impl.rs
@@ -1,6 +1,6 @@
 use crate::{
     bindings::http::types::{self, Headers, Method, Scheme, StatusCode, Trailers},
-    body::{FinishMessage, HostFutureTrailers, HostIncomingBody, HostOutgoingBody},
+    body::{HostFutureTrailers, HostIncomingBody, HostOutgoingBody},
     types::{
         FieldMap, HostFields, HostFutureIncomingResponse, HostIncomingRequest,
         HostIncomingResponse, HostOutgoingRequest, HostOutgoingResponse, HostResponseOutparam,
@@ -26,6 +26,26 @@ impl<T: WasiHttpView> crate::bindings::http::types::Host for T {
     }
 }
 
+/// Extract the `Content-Length` header value from a [`FieldMap`], returning `None` if it's not
+/// present. This function will return `Err` if it's not possible to parse the `Content-Length`
+/// header.
+fn get_content_length(fields: &FieldMap) -> Result<Option<u64>, ()> {
+    let header_val = match fields.get(hyper::header::CONTENT_LENGTH) {
+        Some(val) => val,
+        None => return Ok(None),
+    };
+
+    let header_str = match header_val.to_str() {
+        Ok(val) => val,
+        Err(_) => return Err(()),
+    };
+
+    match header_str.parse() {
+        Ok(len) => Ok(Some(len)),
+        Err(_) => Err(()),
+    }
+}
+
 /// Take ownership of the underlying [`FieldMap`] associated with this fields resource. If the
 /// fields resource references another fields, the returned [`FieldMap`] will be cloned.
 fn move_fields(table: &mut Table, id: Resource<HostFields>) -> wasmtime::Result<FieldMap> {
@@ -374,7 +394,12 @@ impl<T: WasiHttpView> crate::bindings::http::types::HostOutgoingRequest for T {
             return Ok(Err(()));
         }
 
-        let (host_body, hyper_body) = HostOutgoingBody::new();
+        let size = match get_content_length(&req.headers) {
+            Ok(size) => size,
+            Err(e) => return Ok(Err(e)),
+        };
+
+        let (host_body, hyper_body) = HostOutgoingBody::new(size);
 
         req.body = Some(hyper_body);
 
@@ -713,7 +738,12 @@ impl<T: WasiHttpView> crate::bindings::http::types::HostOutgoingResponse for T {
             return Ok(Err(()));
         }
 
-        let (host, body) = HostOutgoingBody::new();
+        let size = match get_content_length(&resp.headers) {
+            Ok(size) => size,
+            Err(e) => return Ok(Err(e)),
+        };
+
+        let (host, body) = HostOutgoingBody::new(size);
 
         resp.body.replace(body);
 
@@ -845,37 +875,20 @@ impl<T: WasiHttpView> crate::bindings::http::types::HostOutgoingBody for T {
         &mut self,
         id: Resource<HostOutgoingBody>,
         ts: Option<Resource<Trailers>>,
-    ) -> wasmtime::Result<()> {
-        let mut body = self.table().delete(id)?;
-
-        let sender = body
-            .finish_sender
-            .take()
-            .expect("outgoing-body trailer_sender consumed by a non-owning function");
+    ) -> wasmtime::Result<Result<(), types::ErrorCode>> {
+        let body = self.table().delete(id)?;
 
-        let message = if let Some(ts) = ts {
-            FinishMessage::Trailers(move_fields(self.table(), ts)?)
+        let ts = if let Some(ts) = ts {
+            Some(move_fields(self.table(), ts)?)
         } else {
-            FinishMessage::Finished
+            None
         };
 
-        // Ignoring failure: receiver died sending body, but we can't report that here.
-        let _ = sender.send(message.into());
-
-        Ok(())
+        Ok(body.finish(ts))
     }
 
     fn drop(&mut self, id: Resource<HostOutgoingBody>) -> wasmtime::Result<()> {
-        let mut body = self.table().delete(id)?;
-
-        let sender = body
-            .finish_sender
-            .take()
-            .expect("outgoing-body trailer_sender consumed by a non-owning function");
-
-        // Ignoring failure: receiver died sending body, but we can't report that here.
-        let _ = sender.send(FinishMessage::Abort);
-
+        self.table().delete(id)?.abort();
         Ok(())
     }
 }
diff --git a/crates/wasi-http/wit/deps/http/types.wit b/crates/wasi-http/wit/deps/http/types.wit
index 41226894e..b9ff389b0 100644
--- a/crates/wasi-http/wit/deps/http/types.wit
+++ b/crates/wasi-http/wit/deps/http/types.wit
@@ -477,7 +477,10 @@ interface types {
     /// called to signal that the response is complete. If the `outgoing-body`
     /// is dropped without calling `outgoing-body.finalize`, the implementation
     /// should treat the body as corrupted.
-    finish: static func(this: outgoing-body, trailers: option<trailers>);
+    finish: static func(
+      this: outgoing-body,
+      trailers: option<trailers>
+    ) -> result<_, error-code>;
   }
 
   /// Represents a future which may eventaully return an incoming HTTP
diff --git a/crates/wasi/wit/deps/http/types.wit b/crates/wasi/wit/deps/http/types.wit
index 41226894e..b9ff389b0 100644
--- a/crates/wasi/wit/deps/http/types.wit
+++ b/crates/wasi/wit/deps/http/types.wit
@@ -477,7 +477,10 @@ interface types {
     /// called to signal that the response is complete. If the `outgoing-body`
     /// is dropped without calling `outgoing-body.finalize`, the implementation
     /// should treat the body as corrupted.
-    finish: static func(this: outgoing-body, trailers: option<trailers>);
+    finish: static func(
+      this: outgoing-body,
+      trailers: option<trailers>
+    ) -> result<_, error-code>;
   }
 
   /// Represents a future which may eventaully return an incoming HTTP


Test patch:
diff --git a/crates/test-programs/src/bin/api_proxy.rs b/crates/test-programs/src/bin/api_proxy.rs
index ff36e2519..c9f292330 100644
--- a/crates/test-programs/src/bin/api_proxy.rs
+++ b/crates/test-programs/src/bin/api_proxy.rs
@@ -42,7 +42,8 @@ impl bindings::exports::wasi::http::incoming_handler::Guest for T {
             .expect("writing response");
 
         drop(out);
-        bindings::wasi::http::types::OutgoingBody::finish(body, None);
+        bindings::wasi::http::types::OutgoingBody::finish(body, None)
+            .expect("outgoing-body.finish");
     }
 }
 
diff --git a/crates/test-programs/src/bin/api_proxy_streaming.rs b/crates/test-programs/src/bin/api_proxy_streaming.rs
index 9c54eb474..b2a354650 100644
--- a/crates/test-programs/src/bin/api_proxy_streaming.rs
+++ b/crates/test-programs/src/bin/api_proxy_streaming.rs
@@ -258,7 +258,7 @@ fn respond(status: u16, response_out: ResponseOutparam) {
 
     ResponseOutparam::set(response_out, Ok(response));
 
-    OutgoingBody::finish(body, None);
+    OutgoingBody::finish(body, None).expect("outgoing-body.finish");
 }
 
 async fn hash(url: &Url) -> Result<String> {
@@ -390,7 +390,7 @@ mod executor {
             fn drop(&mut self) {
                 if let Some((stream, body)) = self.0.take() {
                     drop(stream);
-                    OutgoingBody::finish(body, None);
+                    OutgoingBody::finish(body, None).expect("outgoing-body.finish");
                 }
             }
         }
diff --git a/crates/test-programs/src/bin/http_outbound_request_content_length.rs b/crates/test-programs/src/bin/http_outbound_request_content_length.rs
new file mode 100644
index 000000000..c4c4e4a4a
--- /dev/null
+++ b/crates/test-programs/src/bin/http_outbound_request_content_length.rs
@@ -0,0 +1,92 @@
+use test_programs::wasi::http::types as http_types;
+
+fn make_request() -> http_types::OutgoingRequest {
+    let request = http_types::OutgoingRequest::new(
+        http_types::Headers::from_list(&[("Content-Length".to_string(), b"11".to_vec())]).unwrap(),
+    );
+
+    request
+        .set_method(&http_types::Method::Post)
+        .expect("setting method");
+    request
+        .set_scheme(Some(&http_types::Scheme::Http))
+        .expect("setting scheme");
+    let addr = std::env::var("HTTP_SERVER").unwrap();
+    request
+        .set_authority(Some(&addr))
+        .expect("setting authority");
+    request
+        .set_path_with_query(Some("/"))
+        .expect("setting path with query");
+
+    request
+}
+
+fn main() {
+    {
+        println!("writing enough");
+        let request = make_request();
+        let outgoing_body = request.body().unwrap();
+
+        {
+            let request_body = outgoing_body.write().unwrap();
+            request_body
+                .blocking_write_and_flush("long enough".as_bytes())
+                .unwrap();
+        }
+
+        http_types::OutgoingBody::finish(outgoing_body, None).expect("enough written")
+    }
+
+    {
+        println!("writing too little");
+        let request = make_request();
+        let outgoing_body = request.body().unwrap();
+
+        {
+            let request_body = outgoing_body.write().unwrap();
+            request_body
+                .blocking_write_and_flush("msg".as_bytes())
+                .unwrap();
+        }
+
+        let e =
+            http_types::OutgoingBody::finish(outgoing_body, None).expect_err("finish should fail");
+
+        assert!(
+            matches!(
+                &e,
+                http_types::ErrorCode::InternalError(Some(s))
+                  if s == "not enough written to body stream",
+            ),
+            "unexpected error: {e:#?}"
+        );
+    }
+
+    {
+        println!("writing too much");
+        let request = make_request();
+        let outgoing_body = request.body().unwrap();
+
+        {
+            let request_body = outgoing_body.write().unwrap();
+            request_body
+                .blocking_write_and_flush("more than 11 bytes".as_bytes())
+                .expect_err("write should fail");
+
+            // TODO: show how to use http-error-code to unwrap this error
+        }
+
+        let e =
+            http_types::OutgoingBody::finish(outgoing_body, None).expect_err("finish should fail");
+
+        assert!(
+            matches!(
+                &e,
+                http_types::ErrorCode::InternalError(Some(s))
+                  if s == "too much written to body stream",
+            ),
+            "unexpected error: {e:#?}"
+        );
+    }
+}
diff --git a/crates/test-programs/src/http.rs b/crates/test-programs/src/http.rs
index 64a182ac8..93ccb8243 100644
--- a/crates/test-programs/src/http.rs
+++ b/crates/test-programs/src/http.rs
@@ -110,7 +110,7 @@ pub fn request(
 
     let future_response = outgoing_handler::handle(request, None)?;
 
-    http_types::OutgoingBody::finish(outgoing_body, None);
+    http_types::OutgoingBody::finish(outgoing_body, None)?;
 
     let incoming_response = match future_response.get() {
         Some(result) => result.map_err(|()| anyhow!("response already taken"))?,
diff --git a/crates/wasi-http/tests/all/async_.rs b/crates/wasi-http/tests/all/async_.rs
index 3b74482e5..e9d735279 100644
--- a/crates/wasi-http/tests/all/async_.rs
+++ b/crates/wasi-http/tests/all/async_.rs
@@ -84,3 +84,9 @@ async fn http_outbound_request_response_build() -> Result<()> {
     let server = Server::http1()?;
     run(HTTP_OUTBOUND_REQUEST_RESPONSE_BUILD_COMPONENT, &server).await
 }
+
+#[test_log::test(tokio::test(flavor = "multi_thread"))]
+async fn http_outbound_request_content_length() -> Result<()> {
+    let server = Server::http1()?;
+    run(HTTP_OUTBOUND_REQUEST_CONTENT_LENGTH_COMPONENT, &server).await
+}
diff --git a/crates/wasi-http/tests/all/sync.rs b/crates/wasi-http/tests/all/sync.rs
index ea9328aff..78a4c77cd 100644
--- a/crates/wasi-http/tests/all/sync.rs
+++ b/crates/wasi-http/tests/all/sync.rs
@@ -83,3 +83,9 @@ fn http_outbound_request_response_build() -> Result<()> {
     let server = Server::http1()?;
     run(HTTP_OUTBOUND_REQUEST_RESPONSE_BUILD_COMPONENT, &server)
 }
+
+#[test_log::test]
+fn http_outbound_request_content_length() -> Result<()> {
+    let server = Server::http1()?;
+    run(HTTP_OUTBOUND_REQUEST_CONTENT_LENGTH_COMPONENT, &server)
+}