spansy: support Transfer-Encoding

spansy/Cargo.toml

@@ -20,3 +20,4 @@ thiserror.workspace = true
 httparse = "1.8"
 pest = { version = "2.7" }
 pest_derive = { version = "2.7" }
+flate2 = "1.0.30"

spansy/src/http/mod.rs

@@ -2,6 +2,7 @@
 
 mod span;
 mod types;
+mod parse;
 
 use bytes::Bytes;
 
@@ -10,6 +11,7 @@ pub use types::{
     Body, BodyContent, Code, Header, HeaderName, HeaderValue, Method, Reason, Request, RequestLine,
     Response, Status, Target,
 };
+pub use parse::{parse_chunked_body, parse_deflate_body, parse_gzip_body, parse_identity_body};
 
 use crate::ParseError;
 

spansy/src/http/parse.rs

@@ -0,0 +1,63 @@
+use crate::ParseError;
+use bytes::{Bytes, BytesMut};
+use flate2::read::{GzDecoder, DeflateDecoder};
+use std::io::Read;
+
+// Parsing functions for Transfer-Encoding header types
+/// Parse Transfer-Encoding: chunked body
+pub fn parse_chunked_body(src: &Bytes, offset: usize) -> Result<(Bytes, usize), ParseError> {
+    let mut body = BytesMut::new();
+    let mut pos = offset;
+
+    loop {
+        let chunk_size_end = src[pos..]
+            .windows(2)
+            .position(|w| w == b"\r\n")
+            .ok_or_else(|| ParseError("Invalid chunked encoding: missing chunk size CRLF".to_string()))?
+            + pos;
+
+        let chunk_size_str = std::str::from_utf8(&src[pos..chunk_size_end])
+            .map_err(|_| ParseError("Invalid chunk size encoding".to_string()))?;
+        let chunk_size = usize::from_str_radix(chunk_size_str.trim(), 16)
+            .map_err(|_| ParseError("Invalid chunk size value".to_string()))?;
+
+        pos = chunk_size_end + 2;
+
+        if chunk_size == 0 {
+            break;
+        }
+
+        let chunk_data_end = pos + chunk_size;
+        if chunk_data_end > src.len() {
+            return Err(ParseError("Chunk data exceeds source length".to_string()));
+        }
+        body.extend_from_slice(&src[pos..chunk_data_end]);
+
+        pos = chunk_data_end + 2;
+    }
+
+    pos += 2;
+
+    Ok((body.freeze(), pos))
+}
+
+/// Parse Transfer-Encoding: gzip body
+pub fn parse_gzip_body(src: &Bytes) -> Result<Bytes, ParseError> {
+    let mut decoder = GzDecoder::new(&src[..]);
+    let mut decompressed = Vec::new();
+    decoder.read_to_end(&mut decompressed).map_err(|e| ParseError(format!("Failed to decompress gzip body: {}", e)))?;
+    Ok(Bytes::from(decompressed))
+}
+
+/// Parse Transfer-Encoding: deflate body
+pub fn parse_deflate_body(src: &Bytes) -> Result<Bytes, ParseError> {
+    let mut decoder = DeflateDecoder::new(&src[..]);
+    let mut decompressed = Vec::new();
+    decoder.read_to_end(&mut decompressed).map_err(|e| ParseError(format!("Failed to decompress deflate body: {}", e)))?;
+    Ok(Bytes::from(decompressed))
+}
+
+/// Parse Transfer-Encoding: identity body
+pub fn parse_identity_body(src: &Bytes) -> Result<Bytes, ParseError> {
+    Ok(src.clone())
+}

spansy/src/http/span.rs

@@ -6,7 +6,7 @@ use crate::{
     helpers::get_span_range,
     http::{
         Body, BodyContent, Code, Header, HeaderName, HeaderValue, Method, Reason, Request,
-        RequestLine, Response, Status, Target,
+        RequestLine, Response, Status, Target, parse_chunked_body, parse_deflate_body, parse_gzip_body, parse_identity_body
     },
     json, ParseError, Span,
 };
@@ -168,7 +168,23 @@ pub(crate) fn parse_response_from_bytes(
 
     let body_len = response_body_len(&response)?;
 
-    if body_len > 0 {
+    if body_len == usize::MAX {
+        // Handle different transfer encodings
+        let transfer_encoding = response.headers_with_name("Transfer-Encoding").next().unwrap().value.as_bytes();
+        let (body_bytes, end_pos) = match transfer_encoding {
+            b"chunked" => parse_chunked_body(src, head_end)?,
+            b"gzip" => (parse_gzip_body(&src.slice(head_end..))?, src.len()),
+            b"deflate" => (parse_deflate_body(&src.slice(head_end..))?, src.len()),
+            b"identity" => (parse_identity_body(&src.slice(head_end..))?, src.len()),
+            _ => return Err(ParseError("Unsupported Transfer-Encoding".to_string())),
+        };
+        let body_span = Span::new_bytes(body_bytes.clone(), 0..body_bytes.len());
+        response.body = Some(Body {
+            span: Span::new_bytes(src.clone(), head_end..end_pos),
+            content: BodyContent::Unknown(body_span),
+        });
+        response.span = Span::new_bytes(src.clone(), offset..end_pos);
+    } else if body_len > 0 {
         let range = head_end..head_end + body_len;
 
         if range.end > src.len() {
@@ -256,14 +272,12 @@ fn response_body_len(response: &Response) -> Result<usize, ParseError> {
         _ => {}
     }
 
-    if response
-        .headers_with_name("Transfer-Encoding")
-        .next()
-        .is_some()
-    {
-        Err(ParseError(
-            "Transfer-Encoding not supported yet".to_string(),
-        ))
+    if let Some(transfer_encoding) = response.headers_with_name("Transfer-Encoding").next() {
+        match transfer_encoding.value.as_bytes() {
+            b"chunked" => Ok(usize::MAX),
+            b"gzip" | b"deflate" | b"identity" => Ok(usize::MAX),
+            _ => Err(ParseError("Unsupported Transfer-Encoding".to_string())),
+        }
     } else if let Some(h) = response.headers_with_name("Content-Length").next() {
         // If a valid Content-Length header field is present without Transfer-Encoding, its decimal value
         // defines the expected message body length in octets.