Repository: bearcove/loona (language: rust)

Issue:
Error out if we write too few / too many bytes compared to the announced content-length



Discussion / hints:

Fix patch:
diff --git a/crates/fluke/src/h1/encode.rs b/crates/fluke/src/h1/encode.rs
index fa550351..c0ff6c23 100644
--- a/crates/fluke/src/h1/encode.rs
+++ b/crates/fluke/src/h1/encode.rs
@@ -1,11 +1,11 @@
 use std::io::Write;
 
 use eyre::Context;
-use http::{StatusCode, Version};
+use http::{header, StatusCode, Version};
 
 use crate::{
     types::{Headers, Request, Response},
-    Encoder,
+    Encoder, HeadersExt,
 };
 use fluke_buffet::{Piece, PieceList, RollMut, WriteOwned};
 
@@ -148,13 +148,44 @@ where
     T: WriteOwned,
 {
     pub(crate) transport_w: T,
+    mode: BodyWriteMode,
+}
+
+impl<T> H1Encoder<T>
+where
+    T: WriteOwned,
+{
+    pub fn new(transport_w: T) -> Self {
+        Self {
+            transport_w,
+            mode: BodyWriteMode::Empty,
+        }
+    }
 }
 
 impl<T> Encoder for H1Encoder<T>
 where
     T: WriteOwned,
 {
-    async fn write_response(&mut self, res: Response) -> eyre::Result<()> {
+    async fn write_response(&mut self, mut res: Response) -> eyre::Result<()> {
+        // TODO: set BodyWriteMode here (and take it out of the Encoder trait,
+        // h2 doesn't need it)
+
+        if !res.status.is_informational() {
+            // after this, we expect a body — time to determine the body write mode
+            if !res.means_empty_body() {
+                self.mode = match res.headers.content_length() {
+                    Some(0) => BodyWriteMode::Empty,
+                    Some(_) => BodyWriteMode::ContentLength,
+                    None => {
+                        res.headers
+                            .insert(header::TRANSFER_ENCODING, "chunked".into());
+                        BodyWriteMode::Chunked
+                    }
+                };
+            };
+        }
+
         let mut list = PieceList::default();
         encode_response(res, &mut list)?;
 
@@ -167,14 +198,12 @@ where
     }
 
     // TODO: move `mode` into `H1Encoder`? we don't need it for h2
-    async fn write_body_chunk(&mut self, chunk: Piece, mode: BodyWriteMode) -> eyre::Result<()> {
-        // TODO: inline
-        write_h1_body_chunk(&mut self.transport_w, chunk, mode).await
+    async fn write_body_chunk(&mut self, chunk: Piece) -> eyre::Result<()> {
+        write_h1_body_chunk(&mut self.transport_w, chunk, self.mode).await
     }
 
-    async fn write_body_end(&mut self, mode: BodyWriteMode) -> eyre::Result<()> {
-        // TODO: inline
-        write_h1_body_end(&mut self.transport_w, mode).await
+    async fn write_body_end(&mut self) -> eyre::Result<()> {
+        write_h1_body_end(&mut self.transport_w, self.mode).await
     }
 
     async fn write_trailers(&mut self, trailers: Box<Headers>) -> eyre::Result<()> {
diff --git a/crates/fluke/src/h1/server.rs b/crates/fluke/src/h1/server.rs
index 4ed521ae..354fa57f 100644
--- a/crates/fluke/src/h1/server.rs
+++ b/crates/fluke/src/h1/server.rs
@@ -93,7 +93,7 @@ pub async fn serve(
             },
         );
 
-        let responder = Responder::new(H1Encoder { transport_w });
+        let responder = Responder::new(H1Encoder::new(transport_w));
 
         let resp = driver
             .handle(req, &mut req_body, responder)
diff --git a/crates/fluke/src/h2/encode.rs b/crates/fluke/src/h2/encode.rs
index fb8ce56f..ac407f76 100644
--- a/crates/fluke/src/h2/encode.rs
+++ b/crates/fluke/src/h2/encode.rs
@@ -4,7 +4,7 @@ use tokio::sync::mpsc;
 use tracing::debug;
 
 use super::types::{H2Event, H2EventPayload};
-use crate::{h1::body::BodyWriteMode, Encoder, Response};
+use crate::{Encoder, Response};
 use fluke_h2_parse::StreamId;
 
 #[derive(Debug, PartialEq, Eq)]
@@ -64,7 +64,7 @@ impl Encoder for H2Encoder {
     }
 
     // TODO: BodyWriteMode is not relevant for h2
-    async fn write_body_chunk(&mut self, chunk: Piece, _mode: BodyWriteMode) -> eyre::Result<()> {
+    async fn write_body_chunk(&mut self, chunk: Piece) -> eyre::Result<()> {
         assert!(matches!(self.state, EncoderState::ExpectResponseBody));
 
         self.send(H2EventPayload::BodyChunk(chunk)).await?;
@@ -72,7 +72,7 @@ impl Encoder for H2Encoder {
     }
 
     // TODO: BodyWriteMode is not relevant for h2
-    async fn write_body_end(&mut self, _mode: BodyWriteMode) -> eyre::Result<()> {
+    async fn write_body_end(&mut self) -> eyre::Result<()> {
         assert!(matches!(self.state, EncoderState::ExpectResponseBody));
 
         self.send(H2EventPayload::BodyEnd).await?;
diff --git a/crates/fluke/src/responder.rs b/crates/fluke/src/responder.rs
index bef622d2..557a1484 100644
--- a/crates/fluke/src/responder.rs
+++ b/crates/fluke/src/responder.rs
@@ -1,7 +1,7 @@
 use fluke_buffet::Piece;
 use http::header;
 
-use crate::{h1::body::BodyWriteMode, Body, BodyChunk, Headers, HeadersExt, Response};
+use crate::{Body, BodyChunk, Headers, HeadersExt, Response};
 
 pub trait ResponseState {}
 
@@ -9,7 +9,8 @@ pub struct ExpectResponseHeaders;
 impl ResponseState for ExpectResponseHeaders {}
 
 pub struct ExpectResponseBody {
-    mode: BodyWriteMode,
+    pub announced_content_length: Option<u64>,
+    pub bytes_written: u64,
 }
 impl ResponseState for ExpectResponseBody {}
 
@@ -47,44 +48,36 @@ where
         Ok(())
     }
 
-    /// Send the final response headers
-    /// Errors out if the response status is < 200.
-    /// Errors out if the client sent `expect: 100-continue`
-    pub async fn write_final_response(
+    async fn write_final_response_internal(
         mut self,
-        mut res: Response,
+        res: Response,
+        announced_content_length: Option<u64>,
     ) -> eyre::Result<Responder<E, ExpectResponseBody>> {
         if res.status.is_informational() {
             return Err(eyre::eyre!("final response must have status code >= 200"));
         }
-
-        let mode = if res.means_empty_body() {
-            // do nothing
-            BodyWriteMode::Empty
-        } else {
-            match res.headers.content_length() {
-                Some(0) => BodyWriteMode::Empty,
-                Some(len) => {
-                    // TODO: can probably save that heap allocation
-                    res.headers
-                        .insert(header::CONTENT_LENGTH, format!("{len}").into_bytes().into());
-                    BodyWriteMode::ContentLength
-                }
-                None => {
-                    res.headers
-                        .insert(header::TRANSFER_ENCODING, "chunked".into());
-                    BodyWriteMode::Chunked
-                }
-            }
-        };
         self.encoder.write_response(res).await?;
-
         Ok(Responder {
-            state: ExpectResponseBody { mode },
+            state: ExpectResponseBody {
+                announced_content_length,
+                bytes_written: 0,
+            },
             encoder: self.encoder,
         })
     }
 
+    /// Send the final response headers
+    /// Errors out if the response status is < 200.
+    /// Errors out if the client sent `expect: 100-continue`
+    pub async fn write_final_response(
+        self,
+        res: Response,
+    ) -> eyre::Result<Responder<E, ExpectResponseBody>> {
+        let announced_content_length = res.headers.content_length();
+        self.write_final_response_internal(res, announced_content_length)
+            .await
+    }
+
     /// Writes a response with the given body. Sets `content-length` or
     /// `transfer-encoding` as needed.
     pub async fn write_final_response_with_body(
@@ -97,7 +90,13 @@ where
                 .entry(header::CONTENT_LENGTH)
                 .or_insert_with(|| {
                     // TODO: can probably get rid of this heap allocation, also
-                    // use `itoa`
+                    // use `itoa`, cf. https://docs.rs/itoa/latest/itoa/
+                    // TODO: also, for `write_final_response`, we could avoid
+                    // parsing the content-length header, since we have the
+                    // content-length already.
+                    // TODO: also, instead of doing a heap allocation with a `Vec<u8>`, we
+                    // could probably take a tiny bit of `RollMut` which is cheaper to
+                    // allocate (or at least was, last I measured).
                     format!("{clen}").into_bytes().into()
                 });
         }
@@ -110,8 +109,6 @@ where
                     this.write_chunk(chunk).await?;
                 }
                 BodyChunk::Done { trailers } => {
-                    // TODO: should we do something here in case of
-                    // content-length mismatches?
                     return this.finish_body(trailers).await;
                 }
             }
@@ -125,8 +122,10 @@ where
 {
     /// Send a response body chunk. Errors out if sending more than the
     /// announced content-length.
+    #[inline]
     pub async fn write_chunk(&mut self, chunk: Piece) -> eyre::Result<()> {
-        self.encoder.write_body_chunk(chunk, self.state.mode).await
+        self.state.bytes_written += chunk.len() as u64;
+        self.encoder.write_body_chunk(chunk).await
     }
 
     /// Finish the body, with optional trailers, cf. <https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/TE>
@@ -139,14 +138,20 @@ where
         mut self,
         trailers: Option<Box<Headers>>,
     ) -> eyre::Result<Responder<E, ResponseDone>> {
-        self.encoder.write_body_end(self.state.mode).await?;
+        if let Some(announced_content_length) = self.state.announced_content_length {
+            if self.state.bytes_written != announced_content_length {
+                eyre::bail!(
+                    "content-length mismatch: announced {announced_content_length}, wrote {}",
+                    self.state.bytes_written
+                );
+            }
+        }
+        self.encoder.write_body_end().await?;
 
         if let Some(trailers) = trailers {
             self.encoder.write_trailers(trailers).await?;
         }
 
-        // TODO: check announced content-length size vs actual, etc.
-
         Ok(Responder {
             state: ResponseDone,
             encoder: self.encoder,
@@ -166,7 +171,63 @@ where
 #[allow(async_fn_in_trait)] // we never require Send
 pub trait Encoder {
     async fn write_response(&mut self, res: Response) -> eyre::Result<()>;
-    async fn write_body_chunk(&mut self, chunk: Piece, mode: BodyWriteMode) -> eyre::Result<()>;
-    async fn write_body_end(&mut self, mode: BodyWriteMode) -> eyre::Result<()>;
+    async fn write_body_chunk(&mut self, chunk: Piece) -> eyre::Result<()>;
+    async fn write_body_end(&mut self) -> eyre::Result<()>;
     async fn write_trailers(&mut self, trailers: Box<Headers>) -> eyre::Result<()>;
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use fluke_buffet::Piece;
+    use http::{StatusCode, Version};
+
+    struct MockEncoder;
+
+    impl Encoder for MockEncoder {
+        async fn write_response(&mut self, _: Response) -> eyre::Result<()> {
+            Ok(())
+        }
+        async fn write_body_chunk(&mut self, _: Piece) -> eyre::Result<()> {
+            Ok(())
+        }
+        async fn write_body_end(&mut self) -> eyre::Result<()> {
+            Ok(())
+        }
+        async fn write_trailers(&mut self, _: Box<Headers>) -> eyre::Result<()> {
+            Ok(())
+        }
+    }
+
+    #[tokio::test]
+    async fn test_content_length_mismatch() {
+        let encoder = MockEncoder;
+        let mut res = Response {
+            status: StatusCode::OK,
+            version: Version::HTTP_11,
+            headers: Default::default(),
+        };
+        res.headers.insert(header::CONTENT_LENGTH, "10".into());
+
+        // Test writing fewer bytes than announced
+        let mut responder = Responder::new(encoder)
+            .write_final_response(res.clone())
+            .await
+            .unwrap();
+        responder.write_chunk(b"12345".into()).await.unwrap();
+        let result = responder.finish_body(None).await;
+        assert!(result.is_err());
+        assert!(matches!(result, Err(e) if e.to_string().contains("content-length mismatch")));
+
+        // Test writing more bytes than announced
+        let encoder = MockEncoder;
+        let mut responder = Responder::new(encoder)
+            .write_final_response(res)
+            .await
+            .unwrap();
+        responder.write_chunk(b"12345678901".into()).await.unwrap();
+        let result = responder.finish_body(None).await;
+        assert!(result.is_err());
+        assert!(matches!(result, Err(e) if e.to_string().contains("content-length mismatch")));
+    }
+}


Test patch: