Merge pull request #1243 from mpajkowski/zlib_module

Implement zlib module

vm/src/stdlib/mod.rs

@@ -43,6 +43,8 @@ mod pwd;
 pub mod signal;
 #[cfg(not(target_arch = "wasm32"))]
 mod subprocess;
+#[cfg(not(target_arch = "wasm32"))]
+mod zlib;
 
 use crate::pyobject::PyObjectRef;
 
@@ -99,6 +101,7 @@ pub fn get_module_inits() -> HashMap<String, StdlibInitFunc> {
         modules.insert("socket".to_string(), Box::new(socket::make_module));
         modules.insert("signal".to_string(), Box::new(signal::make_module));
         modules.insert("subprocess".to_string(), Box::new(subprocess::make_module));
+        modules.insert("zlib".to_string(), Box::new(zlib::make_module));
     }
 
     // Unix-only

vm/src/stdlib/zlib.rs

@@ -0,0 +1,153 @@
+use crate::function::OptionalArg;
+use crate::obj::{objbytes::PyBytesRef, objint::PyIntRef};
+use crate::pyobject::{create_type, ItemProtocol, PyObjectRef, PyResult};
+use crate::vm::VirtualMachine;
+
+use adler32::RollingAdler32 as Adler32;
+use crc32fast::Hasher as Crc32;
+use flate2::{write::ZlibEncoder, Compression, Decompress, FlushDecompress, Status};
+use libz_sys as libz;
+use num_traits::cast::ToPrimitive;
+
+use std::io::Write;
+
+// copied from zlibmodule.c (commit 530f506ac91338)
+const MAX_WBITS: u8 = 15;
+const DEF_BUF_SIZE: usize = 16 * 1024;
+
+pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
+    let ctx = &vm.ctx;
+
+    let zlib_error = create_type("error", &ctx.type_type, &ctx.exceptions.exception_type);
+
+    py_module!(vm, "zlib", {
+        "crc32" => ctx.new_rustfunc(zlib_crc32),
+        "adler32" => ctx.new_rustfunc(zlib_adler32),
+        "compress" => ctx.new_rustfunc(zlib_compress),
+        "decompress" => ctx.new_rustfunc(zlib_decompress),
+        "error" => zlib_error,
+        "Z_DEFAULT_COMPRESSION" => ctx.new_int(libz::Z_DEFAULT_COMPRESSION),
+        "Z_NO_COMPRESSION" => ctx.new_int(libz::Z_NO_COMPRESSION),
+        "Z_BEST_SPEED" => ctx.new_int(libz::Z_BEST_SPEED),
+        "Z_BEST_COMPRESSION" => ctx.new_int(libz::Z_BEST_COMPRESSION),
+        "DEF_BUF_SIZE" => ctx.new_int(DEF_BUF_SIZE),
+        "MAX_WBITS" => ctx.new_int(MAX_WBITS),
+    })
+}
+
+/// Compute an Adler-32 checksum of data.
+fn zlib_adler32(
+    data: PyBytesRef,
+    begin_state: OptionalArg<PyIntRef>,
+    vm: &VirtualMachine,
+) -> PyResult<PyObjectRef> {
+    let data = data.get_value();
+
+    let begin_state = begin_state
+        .into_option()
+        .as_ref()
+        .map(|v| v.as_bigint().to_i32().unwrap())
+        .unwrap_or(1);
+
+    let mut hasher = Adler32::from_value(begin_state as u32);
+    hasher.update_buffer(data);
+
+    let checksum: u32 = hasher.hash();
+
+    Ok(vm.new_int(checksum))
+}
+
+/// Compute a CRC-32 checksum of data.
+fn zlib_crc32(
+    data: PyBytesRef,
+    begin_state: OptionalArg<PyIntRef>,
+    vm: &VirtualMachine,
+) -> PyResult<PyObjectRef> {
+    let data = data.get_value();
+
+    let begin_state = begin_state
+        .into_option()
+        .as_ref()
+        .map(|v| v.as_bigint().to_i32().unwrap())
+        .unwrap_or(0);
+
+    let mut hasher = Crc32::new_with_initial(begin_state as u32);
+    hasher.update(data);
+
+    let checksum: u32 = hasher.finalize();
+
+    Ok(vm.new_int(checksum))
+}
+
+/// Returns a bytes object containing compressed data.
+fn zlib_compress(
+    data: PyBytesRef,
+    level: OptionalArg<PyIntRef>,
+    vm: &VirtualMachine,
+) -> PyResult<PyObjectRef> {
+    let input_bytes = data.get_value();
+
+    let level = level
+        .into_option()
+        .as_ref()
+        .map(|v| v.as_bigint().to_i32().unwrap())
+        .unwrap_or(libz::Z_DEFAULT_COMPRESSION);
+
+    let compression = match level {
+        valid_level @ libz::Z_NO_COMPRESSION...libz::Z_BEST_COMPRESSION => {
+            Compression::new(valid_level as u32)
+        }
+        libz::Z_DEFAULT_COMPRESSION => Compression::default(),
+        _ => return Err(zlib_error("Bad compression level", vm)),
+    };
+
+    let mut encoder = ZlibEncoder::new(Vec::new(), compression);
+    encoder.write_all(input_bytes).unwrap();
+    let encoded_bytes = encoder.finish().unwrap();
+
+    Ok(vm.ctx.new_bytes(encoded_bytes))
+}
+
+/// Returns a bytes object containing the uncompressed data.
+fn zlib_decompress(
+    data: PyBytesRef,
+    wbits: OptionalArg<PyIntRef>,
+    bufsize: OptionalArg<PyIntRef>,
+    vm: &VirtualMachine,
+) -> PyResult<PyObjectRef> {
+    let encoded_bytes = data.get_value();
+
+    let wbits = wbits
+        .into_option()
+        .as_ref()
+        .map(|wbits| wbits.as_bigint().to_u8().unwrap())
+        .unwrap_or(MAX_WBITS);
+
+    let bufsize = bufsize
+        .into_option()
+        .as_ref()
+        .map(|bufsize| bufsize.as_bigint().to_usize().unwrap())
+        .unwrap_or(DEF_BUF_SIZE);
+
+    let mut decompressor = Decompress::new_with_window_bits(true, wbits);
+    let mut decoded_bytes = Vec::with_capacity(bufsize);
+
+    match decompressor.decompress_vec(&encoded_bytes, &mut decoded_bytes, FlushDecompress::Finish) {
+        Ok(Status::BufError) => Err(zlib_error("inconsistent or truncated state", vm)),
+        Err(_) => Err(zlib_error("invalid input data", vm)),
+        _ => Ok(vm.ctx.new_bytes(decoded_bytes)),
+    }
+}
+
+fn zlib_error(message: &str, vm: &VirtualMachine) -> PyObjectRef {
+    let module = vm
+        .get_attribute(vm.sys_module.clone(), "modules")
+        .unwrap()
+        .get_item("zlib", vm)
+        .unwrap();
+
+    let zlib_error = vm.get_attribute(module, "error").unwrap();
+    let zlib_error = zlib_error.downcast().unwrap();
+
+    vm.new_exception(zlib_error, message.to_string())
+}