test(decompression-plz): added tests for extra body

Author: darkseid-is

decompression-plz/src/decompression/state.rs

@@ -184,6 +184,10 @@ impl<'a> DecompressionState<'a> {
             || matches!(self, Self::EndExtraMainSeparate(..))
             || matches!(self, Self::EndExtraRawMainDone(..))
     }
+
+    pub fn is_extra_raw(&self) -> bool {
+        matches!(self, Self::EndExtraRawMainDone(..))
+    }
 }
 
 impl<'a> From<DecompressionState<'a>> for (BytesMut, Option<BytesMut>) {

decompression-plz/src/state.rs

@@ -118,7 +118,7 @@ fn apply_encoding<T>(
     encoding_info: &mut [EncodingInfo],
 ) -> Result<(), MultiDecompressErrorReason>
 where
-    T: DecompressTrait,
+    T: DecompressTrait + std::fmt::Debug,
 {
     match decompression_runner(
         &decode_struct.body,
@@ -127,7 +127,12 @@ where
         decode_struct.buf,
     ) {
         Ok(state) => {
-            (decode_struct.body, decode_struct.extra_body) = state.into();
+            let is_extra_raw = state.is_extra_raw();
+            let (body, extra_body) = state.into();
+            decode_struct.body = body;
+            if !is_extra_raw {
+                decode_struct.extra_body = extra_body;
+            }
             let iter = encoding_info.iter().map(|einfo| einfo.header_index);
             // remove applied headers
             decode_struct

decompression-plz/tests/test_cases/single_compression.rs

@@ -0,0 +1,56 @@
+use super::*;
+
+// CE only
+#[test]
+fn assert_decode_state_single_ce_brotli() {
+    assert_case_single_compression(
+        CONTENT_ENCODING,
+        ContentEncoding::Brotli,
+        None,
+    );
+}
+
+#[test]
+fn assert_decode_state_single_ce_compress() {
+    assert_case_single_compression(
+        CONTENT_ENCODING,
+        ContentEncoding::Compress,
+        None,
+    );
+}
+
+#[test]
+fn assert_decode_state_single_ce_deflate() {
+    assert_case_single_compression(
+        CONTENT_ENCODING,
+        ContentEncoding::Deflate,
+        None,
+    );
+}
+
+#[test]
+fn assert_decode_state_single_ce_gzip() {
+    assert_case_single_compression(
+        CONTENT_ENCODING,
+        ContentEncoding::Gzip,
+        None,
+    );
+}
+
+#[test]
+fn assert_decode_state_single_ce_identity() {
+    assert_case_single_compression(
+        CONTENT_ENCODING,
+        ContentEncoding::Identity,
+        None,
+    );
+}
+
+#[test]
+fn assert_decode_state_single_ce_zstd() {
+    assert_case_single_compression(
+        CONTENT_ENCODING,
+        ContentEncoding::Zstd,
+        None,
+    );
+}

decompression-plz/tests/test_cases/single_compression/body_only/ce.rs

@@ -0,0 +1,3 @@
+use super::*;
+mod ce;
+mod te;

decompression-plz/tests/test_cases/single_compression/body_only/mod.rs

@@ -0,0 +1,56 @@
+use super::*;
+
+// Transfer-Encoding
+#[test]
+fn assert_decode_state_single_te_brotli() {
+    assert_case_single_compression(
+        TRANSFER_ENCODING,
+        ContentEncoding::Brotli,
+        None,
+    );
+}
+
+#[test]
+fn assert_decode_state_single_te_compress() {
+    assert_case_single_compression(
+        TRANSFER_ENCODING,
+        ContentEncoding::Compress,
+        None,
+    );
+}
+
+#[test]
+fn assert_decode_state_single_te_deflate() {
+    assert_case_single_compression(
+        TRANSFER_ENCODING,
+        ContentEncoding::Deflate,
+        None,
+    );
+}
+
+#[test]
+fn assert_decode_state_single_te_gzip() {
+    assert_case_single_compression(
+        TRANSFER_ENCODING,
+        ContentEncoding::Gzip,
+        None,
+    );
+}
+
+#[test]
+fn assert_decode_state_single_te_identity() {
+    assert_case_single_compression(
+        TRANSFER_ENCODING,
+        ContentEncoding::Identity,
+        None,
+    );
+}
+
+#[test]
+fn assert_decode_state_single_te_zstd() {
+    assert_case_single_compression(
+        TRANSFER_ENCODING,
+        ContentEncoding::Zstd,
+        None,
+    );
+}

decompression-plz/tests/test_cases/single_compression/body_only/te.rs

@@ -0,0 +1,56 @@
+use header_plz::{
+    body_headers::content_encoding::ContentEncoding,
+    const_headers::{CONTENT_ENCODING, TRANSFER_ENCODING},
+};
+use tests_utils::single_compression;
+
+use super::*;
+mod body_only;
+mod with_extra;
+
+fn assert_case_single_compression(
+    header_name: &str,
+    content_encoding: ContentEncoding,
+    extra: Option<BytesMut>,
+) {
+    let body: Vec<u8> = single_compression(&content_encoding);
+    let headers = format!(
+        "Host: example.com\r\n\
+        Content-Type: text/html; charset=utf-8\r\n\
+        {}: {}\r\n\
+        Content-Length: {}\r\n\r\n",
+        header_name,
+        content_encoding.as_ref(),
+        body.len()
+    );
+
+    let expected_state = encoding_state(header_name);
+    let with_extra = extra.is_some();
+    let mut tm = TestMessage::build(
+        headers.as_bytes().into(),
+        Body::Raw(body.as_slice().into()),
+        extra,
+    );
+    let mut buf = BytesMut::new();
+    let mut state = DecodeState::init(&mut tm, &mut buf);
+    state = state.try_next().unwrap();
+    assert!((expected_state)(&state));
+    state = state.try_next().unwrap();
+    assert!(matches!(state, DecodeState::UpdateContentLength(_)));
+    state = state.try_next().unwrap();
+    assert!(state.is_ended());
+
+    let verify = if with_extra {
+        "Host: example.com\r\n\
+            Content-Type: text/html; charset=utf-8\r\n\
+            Content-Length: 22\r\n\r\n\
+            hello worldhello world"
+    } else {
+        "Host: example.com\r\n\
+            Content-Type: text/html; charset=utf-8\r\n\
+            Content-Length: 11\r\n\r\n\
+            hello world"
+    };
+    let result = tm.into_bytes();
+    assert_eq!(result, verify);
+}