Initial commit for Rust powered modules

Author: Schamper

.gitignore

@@ -9,3 +9,6 @@ __pycache__/
 tests/docs/api
 tests/docs/build
 .tox/
+*/*/_native.src/target
+*/*/_native.*.so
+*/*/_native.*.dll

dissect/util/_native.src/Cargo.lock

@@ -0,0 +1,13 @@
+[package]
+name = "_native"
+version = "0.0.0-git"
+edition = "2021"
+
+[lib]
+name = "_native"
+crate-type = ["cdylib"]
+
+[dependencies]
+pyo3 = "0.22"
+lz4_flex = "0.11"
+lzokay-native = "0.1"

dissect/util/_native.src/Cargo.toml

@@ -0,0 +1,11 @@
+use pyo3::prelude::*;
+
+mod lz4;
+mod lzo;
+
+pub fn create_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    let submodule = PyModule::new_bound(m.py(), "compression")?;
+    lz4::create_submodule(&submodule)?;
+    lzo::create_submodule(&submodule)?;
+    m.add_submodule(&submodule)
+}

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

@@ -0,0 +1,58 @@
+use pyo3::exceptions::PyValueError;
+use pyo3::prelude::*;
+use pyo3::types::{PyByteArray, PyBytes};
+
+const MAX_DISCOVER_OUTPUT_SIZE: usize = 1024 * 1024 * 1024;
+
+#[pyfunction]
+#[pyo3(signature = (src, uncompressed_size=-1, return_bytearray=false))]
+fn decompress(
+    py: Python<'_>,
+    src: Vec<u8>,
+    uncompressed_size: isize,
+    return_bytearray: bool,
+) -> 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()));
+                }
+            }
+        }
+    } else {
+        lz4_flex::block::decompress(&src, uncompressed_size as usize)
+            .map_err(|e| PyErr::new::<PyValueError, _>(e.to_string()))?
+    };
+
+    let pyresult = PyBytes::new_bound(py, &result);
+    if return_bytearray {
+        Ok(PyByteArray::from_bound(&pyresult)?.into())
+    } else {
+        Ok(pyresult.into())
+    }
+}
+
+pub fn create_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    let submodule = PyModule::new_bound(m.py(), "lz4")?;
+    submodule.add_function(wrap_pyfunction!(decompress, m)?)?;
+    m.add_submodule(&submodule)
+}

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

@@ -0,0 +1,39 @@
+use pyo3::exceptions::PyValueError;
+use pyo3::prelude::*;
+use pyo3::types::PyBytes;
+
+#[pyfunction]
+#[pyo3(signature = (src, header=true, buflen=-1))]
+fn decompress(
+    py: Python<'_>,
+    src: Vec<u8>,
+    header: bool,
+    buflen: isize,
+) -> PyResult<Bound<'_, PyBytes>> {
+    let (src, out_len) = if header {
+        if src.len() < 8 || src[0] < 0xf0 || src[0] > 0xf1 {
+            return Err(PyErr::new::<PyValueError, _>(
+                "Invalid header value".to_string(),
+            ));
+        }
+        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(),
+        ));
+    } else {
+        (src, buflen as usize)
+    };
+
+    let mut cursor = std::io::Cursor::new(src);
+    lzokay_native::decompress(&mut cursor, Some(out_len))
+        .map_err(|e| PyErr::new::<PyValueError, _>(e.to_string()))
+        .map(|result| PyBytes::new_bound(py, &result))
+}
+
+pub fn create_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    let submodule = PyModule::new_bound(m.py(), "lzo")?;
+    submodule.add_function(wrap_pyfunction!(decompress, m)?)?;
+    m.add_submodule(&submodule)
+}

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

@@ -0,0 +1,10 @@
+use pyo3::prelude::*;
+
+mod compression;
+
+
+#[pymodule]
+fn _native(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    compression::create_submodule(m)?;
+    Ok(())
+}