@@ -96,19 +96,25 @@ fn zlib_compress(data: PyBytesRef, level: OptionalArg<i32>, vm: &VirtualMachine)
9696 Ok ( vm. ctx . new_bytes ( encoded_bytes) )
9797}
9898
99+ // TODO: validate wbits value here
100+ fn header_from_wbits ( wbits : OptionalArg < i8 > ) -> ( bool , u8 ) {
101+ let wbits = wbits. unwrap_or ( MAX_WBITS as i8 ) ;
102+ ( wbits > 0 , wbits. abs ( ) as u8 )
103+ }
104+
99105/// Returns a bytes object containing the uncompressed data.
100106fn zlib_decompress (
101107 data : PyBytesRef ,
102- wbits : OptionalArg < u8 > ,
108+ wbits : OptionalArg < i8 > ,
103109 bufsize : OptionalArg < usize > ,
104110 vm : & VirtualMachine ,
105111) -> PyResult {
106112 let encoded_bytes = data. get_value ( ) ;
107113
108- let wbits = wbits . unwrap_or ( MAX_WBITS ) ;
114+ let ( header , wbits) = header_from_wbits ( wbits ) ;
109115 let bufsize = bufsize. unwrap_or ( DEF_BUF_SIZE ) ;
110116
111- let mut decompressor = Decompress :: new_with_window_bits ( true , wbits) ;
117+ let mut decompressor = Decompress :: new_with_window_bits ( header , wbits) ;
112118 let mut decoded_bytes = Vec :: with_capacity ( bufsize) ;
113119
114120 match decompressor. decompress_vec ( & encoded_bytes, & mut decoded_bytes, FlushDecompress :: Finish ) {
@@ -119,12 +125,12 @@ fn zlib_decompress(
119125}
120126
121127fn zlib_decompressobj (
122- wbits : OptionalArg < u8 > ,
128+ wbits : OptionalArg < i8 > ,
123129 zdict : OptionalArg < PyBytesLike > ,
124130 vm : & VirtualMachine ,
125131) -> PyDecompress {
126- let wbits = wbits . unwrap_or ( MAX_WBITS ) ;
127- let mut decompress = Decompress :: new_with_window_bits ( true , wbits) ;
132+ let ( header , wbits) = header_from_wbits ( wbits ) ;
133+ let mut decompress = Decompress :: new_with_window_bits ( header , wbits) ;
128134 if let OptionalArg :: Present ( dict) = zdict {
129135 dict. with_ref ( |d| decompress. set_dictionary ( d) . unwrap ( ) ) ;
130136 }
0 commit comments