Process review

Author: Schamper

dissect/util/_native.src/src/compression/lz4.rs

@@ -4,6 +4,33 @@ use pyo3::types::{PyByteArray, PyBytes};
 
 const MAX_DISCOVER_OUTPUT_SIZE: usize = 1024 * 1024 * 1024;
 
+fn decompress_to_unknown_size(src: &[u8]) -> Result<Vec<u8>, PyErr> {
+    let mut output_size = lz4_flex::block::get_maximum_output_size(src.len());
+    loop {
+        // If the output size is too large, we should not attempt to decompress further
+        if output_size > MAX_DISCOVER_OUTPUT_SIZE {
+            return Err(PyErr::new::<PyValueError, _>(
+                "output size is too large".to_string(),
+            ));
+        }
+
+        match lz4_flex::block::decompress(&src, output_size) {
+            Ok(result) => {
+                break Ok(result);
+            }
+            Err(lz4_flex::block::DecompressError::OutputTooSmall {
+                expected,
+                actual: _,
+            }) => {
+                output_size = expected;
+            }
+            Err(e) => {
+                return Err(PyErr::new::<PyValueError, _>(e.to_string()));
+            }
+        }
+    }
+}
+
 /// LZ4 decompress bytes up to a certain length. Assumes no header.
 ///
 /// Unlike the Python implementation, this function does not support streaming decompression
@@ -15,7 +42,7 @@ const MAX_DISCOVER_OUTPUT_SIZE: usize = 1024 * 1024 * 1024;
 ///     return_bytearray: Whether to return ``bytearray`` or ``bytes``.
 ///
 /// Returns:
-///     The decompressed data or a tuple of the decompressed data and the amount of bytes read.
+///     The decompressed data.
 ///
 #[pyfunction]
 #[pyo3(signature = (src, uncompressed_size=-1, return_bytearray=false))]
@@ -27,30 +54,7 @@ fn decompress(
 ) -> PyResult<PyObject> {
     let result = if uncompressed_size < 0 {
         // If the uncompressed size is not provided, we need to discover it first
-        let mut output_size = lz4_flex::block::get_maximum_output_size(src.len());
-        loop {
-            // If the output size is too large, we should not attempt to decompress further
-            if output_size > MAX_DISCOVER_OUTPUT_SIZE {
-                return Err(PyErr::new::<PyValueError, _>(
-                    "output size is too large".to_string(),
-                ));
-            }
-
-            match lz4_flex::block::decompress(&src, output_size) {
-                Ok(result) => {
-                    break result;
-                }
-                Err(lz4_flex::block::DecompressError::OutputTooSmall {
-                    expected,
-                    actual: _,
-                }) => {
-                    output_size = expected;
-                }
-                Err(e) => {
-                    return Err(PyErr::new::<PyValueError, _>(e.to_string()));
-                }
-            }
-        }
+        decompress_to_unknown_size(&src)?
     } else {
         lz4_flex::block::decompress(&src, uncompressed_size as usize)
             .map_err(|e| PyErr::new::<PyValueError, _>(e.to_string()))?

dissect/util/_native.src/src/compression/lzo.rs

@@ -24,24 +24,20 @@ fn decompress(
     header: bool,
     buflen: isize,
 ) -> PyResult<Bound<'_, PyBytes>> {
-    let (src, out_len) = if header {
+    let (body, out_len) = if header {
+        // Compatibility with python-lzo, which can include a header
+        // https://github.com/jd-boyd/python-lzo/blob/80ca60416c6657d373c5308a1eb511903a3ff9b1/lzomodule.c#L238-L269
         if src.len() < 8 || src[0] < 0xf0 || src[0] > 0xf1 {
-            return Err(PyErr::new::<PyValueError, _>(
-                "Invalid header value".to_string(),
-            ));
+            return Err(PyValueError::new_err("Invalid header value"));
         }
         let len = u32::from_le_bytes([src[1], src[2], src[3], src[4]]) as usize;
-        (src[5..].to_vec(), len)
-    } else if buflen < 0 {
-        return Err(PyErr::new::<PyValueError, _>(
-            "Buffer length must be provided".to_string(),
-        ));
+        (src[5..].to_vec(), Some(len))
     } else {
-        (src, buflen as usize)
+        (src, (buflen >= 0).then_some(buflen as usize))
     };
 
-    let mut cursor = std::io::Cursor::new(src);
-    lzokay_native::decompress(&mut cursor, Some(out_len))
+    let mut cursor = std::io::Cursor::new(body);
+    lzokay_native::decompress(&mut cursor, out_len)
         .map_err(|e| PyErr::new::<PyValueError, _>(e.to_string()))
         .map(|result| PyBytes::new(py, &result))
 }

dissect/util/compression/lz4.py

@@ -37,7 +37,7 @@ def decompress(
         return_bytearray: Whether to return ``bytearray`` or ``bytes``.
 
     Returns:
-        The decompressed data or a tuple of the decompressed data and the amount of bytes read.
+        The decompressed data.
     """
     if not hasattr(src, "read"):
         src = io.BytesIO(src)