diff --git a/vm/src/stdlib/zlib.rs b/vm/src/stdlib/zlib.rs index 5f17978d6..27707a09f 100644 --- a/vm/src/stdlib/zlib.rs +++ b/vm/src/stdlib/zlib.rs @@ -96,19 +96,25 @@ fn zlib_compress(data: PyBytesRef, level: OptionalArg, vm: &VirtualMachine) Ok(vm.ctx.new_bytes(encoded_bytes)) } +// TODO: validate wbits value here +fn header_from_wbits(wbits: OptionalArg) -> (bool, u8) { + let wbits = wbits.unwrap_or(MAX_WBITS as i8); + (wbits > 0, wbits.abs() as u8) +} + /// Returns a bytes object containing the uncompressed data. fn zlib_decompress( data: PyBytesRef, - wbits: OptionalArg, + wbits: OptionalArg, bufsize: OptionalArg, vm: &VirtualMachine, ) -> PyResult { let encoded_bytes = data.get_value(); - let wbits = wbits.unwrap_or(MAX_WBITS); + let (header, wbits) = header_from_wbits(wbits); let bufsize = bufsize.unwrap_or(DEF_BUF_SIZE); - let mut decompressor = Decompress::new_with_window_bits(true, wbits); + let mut decompressor = Decompress::new_with_window_bits(header, wbits); let mut decoded_bytes = Vec::with_capacity(bufsize); match decompressor.decompress_vec(&encoded_bytes, &mut decoded_bytes, FlushDecompress::Finish) { @@ -119,12 +125,12 @@ fn zlib_decompress( } fn zlib_decompressobj( - wbits: OptionalArg, + wbits: OptionalArg, zdict: OptionalArg, vm: &VirtualMachine, ) -> PyDecompress { - let wbits = wbits.unwrap_or(MAX_WBITS); - let mut decompress = Decompress::new_with_window_bits(true, wbits); + let (header, wbits) = header_from_wbits(wbits); + let mut decompress = Decompress::new_with_window_bits(header, wbits); if let OptionalArg::Present(dict) = zdict { dict.with_ref(|d| decompress.set_dictionary(d).unwrap()); }