diff --git a/Cargo.lock b/Cargo.lock index 239f19bf824..e4b6a747703 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,12 +2,6 @@ # It is not intended for manual editing. version = 4 -[[package]] -name = "adler2" -version = "2.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" - [[package]] name = "adler32" version = "1.2.0" @@ -1314,16 +1308,6 @@ dependencies = [ "serde_json", ] -[[package]] -name = "flate2" -version = "1.1.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c" -dependencies = [ - "miniz_oxide", - "zlib-rs", -] - [[package]] name = "fnv" version = "1.0.7" @@ -2269,16 +2253,6 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" -[[package]] -name = "miniz_oxide" -version = "0.8.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316" -dependencies = [ - "adler2", - "simd-adler32", -] - [[package]] name = "mt19937" version = "3.3.0" @@ -3366,6 +3340,7 @@ dependencies = [ "getrandom 0.4.3", "itertools 0.15.0", "libc", + "libz-rs-sys", "lock_api", "malachite-base", "malachite-bigint", @@ -3379,6 +3354,9 @@ dependencies = [ "rustpython-unicode", "rustpython-wtf8", "siphasher", + "xz", + "xz-sys", + "zlib-rs", ] [[package]] @@ -3611,7 +3589,6 @@ dependencies = [ "digest", "dyn-clone", "flame", - "flate2", "foreign-types-shared", "hex", "hmac", @@ -3621,7 +3598,6 @@ dependencies = [ "jiff", "libc", "libsqlite3-sys", - "libz-rs-sys", "malachite-bigint", "md-5", "memchr", @@ -3669,8 +3645,6 @@ dependencies = [ "x509-cert", "x509-parser", "xml", - "xz", - "xz-sys", ] [[package]] @@ -4029,12 +4003,6 @@ dependencies = [ "rand_core 0.10.1", ] -[[package]] -name = "simd-adler32" -version = "0.3.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "703d5c7ef118737c72f1af64ad2f6f8c5e1921f818cdcb97b8fe6fc69bf66214" - [[package]] name = "simd_cesu8" version = "1.1.1" diff --git a/Cargo.toml b/Cargo.toml index f2d1579f589..678bec69601 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -215,7 +215,6 @@ dyn-clone = "1.0.10" exitcode = "1.1.2" flame = "0.2.2" flamer = "0.5" -flate2 = { version = "1.1.9", default-features = false } # Bump only when the openssl crate bumps it foreign-types-shared = "0.1" gethostname = "1.0.2" @@ -241,6 +240,7 @@ xz = "0.4.6-rc.0" xz-sys = "0.4.6-rc.0" libsqlite3-sys = "0.38" libz-rs-sys = "0.6" +zlib-rs = { version = "=0.6.7", default-features = false, features = ["rust-allocator", "__internal-api"] } lock_api = "0.4" log = "0.4.30" lz4_flex = "0.13" diff --git a/Lib/test/test_zlib.py b/Lib/test/test_zlib.py index bb1366cb21c..4d97fe56f3a 100644 --- a/Lib/test/test_zlib.py +++ b/Lib/test/test_zlib.py @@ -489,7 +489,6 @@ def test_clear_unconsumed_tail(self): ddata += dco.decompress(dco.unconsumed_tail) self.assertEqual(dco.unconsumed_tail, b"") - @unittest.expectedFailure # TODO: RUSTPYTHON; Z_BLOCK support in flate2 def test_flushes(self): # Test flush() with the various options, using all the # different levels in order to provide more variations. @@ -809,7 +808,6 @@ def test_large_unconsumed_tail(self, size): finally: comp = uncomp = data = None - @unittest.expectedFailure # TODO: RUSTPYTHON; wbits=0 support in flate2 def test_wbits(self): # wbits=0 only supported since zlib v1.2.3.5 supports_wbits_0 = ZLIB_RUNTIME_VERSION_TUPLE >= (1, 2, 3, 5) diff --git a/crates/common/Cargo.toml b/crates/common/Cargo.toml index 83f646e518a..21b700a5526 100644 --- a/crates/common/Cargo.toml +++ b/crates/common/Cargo.toml @@ -14,9 +14,11 @@ binascii = ["dep:base64", "dep:crc32fast"] cjk-codecs = [] inet = ["std"] json = ["dep:memchr", "std"] +lzma = ["std", "dep:xz", "dep:xz-sys"] std = [] threading = ["parking_lot", "std"] wasm_js = ["getrandom/wasm_js"] +zlib = ["std", "dep:libz-rs-sys", "dep:zlib-rs"] [dependencies] rustpython-literal = { workspace = true } @@ -30,6 +32,7 @@ crc32fast = { workspace = true, optional = true } getrandom = { workspace = true } itertools = { workspace = true } libc = { workspace = true } +libz-rs-sys = { workspace = true, optional = true } malachite-bigint = { workspace = true } malachite-q = { workspace = true } malachite-base = { workspace = true } @@ -37,10 +40,16 @@ memchr = { workspace = true, optional = true } num-traits = { workspace = true } parking_lot = { workspace = true, optional = true } radium = { workspace = true } - +zlib-rs = { workspace = true, optional = true } lock_api = { workspace = true } siphasher = { workspace = true } num-complex = { workspace = true } +[target.'cfg(not(any(target_os = "android", target_arch = "wasm32")))'.dependencies] +# RustPython does not expose `_lzma` on these targets; keep the native engine's +# dependency boundary with its common-module owner rather than in stdlib. +xz = { workspace = true, optional = true } +xz-sys = { workspace = true, optional = true } + [lints] workspace = true diff --git a/crates/common/src/compression/lzma.rs b/crates/common/src/compression/lzma.rs new file mode 100644 index 00000000000..1b7943d07fb --- /dev/null +++ b/crates/common/src/compression/lzma.rs @@ -0,0 +1,685 @@ +// spell-checker:ignore ARMTHUMB chunker memlimit + +//! VM-independent liblzma stream engine. +//! +//! `_lzma` is not built on Android or WebAssembly in RustPython, so the xz +//! dependency and this module have the same target boundary. Python object +//! conversion and exception construction remain in `rustpython-stdlib`. + +use xz::stream::{ + Action, Check, Error as XzError, Filters, LzmaOptions, MatchFinder, Mode, Status, Stream, + TELL_ANY_CHECK, TELL_NO_CHECK, +}; + +use super::{CHUNKSIZE, Chunker}; + +pub const BUFSIZ: usize = 8192; +const DEF_BUF_SIZE: usize = 16 * 1024; +const USE_AFTER_FINISH_ERR: &str = "Error -2: inconsistent stream state"; +const LZMA_FILTERS_MAX: usize = 4; + +pub const CHECK_NONE: i32 = xz_sys::LZMA_CHECK_NONE as _; +pub const CHECK_CRC32: i32 = xz_sys::LZMA_CHECK_CRC32 as _; +pub const CHECK_CRC64: i32 = xz_sys::LZMA_CHECK_CRC64 as _; +pub const CHECK_SHA256: i32 = xz_sys::LZMA_CHECK_SHA256 as _; +pub const CHECK_ID_MAX: i32 = 15; +pub const CHECK_UNKNOWN: i32 = CHECK_ID_MAX + 1; + +pub const MF_HC3: i32 = xz_sys::LZMA_MF_HC3 as _; +pub const MF_HC4: i32 = xz_sys::LZMA_MF_HC4 as _; +pub const MF_BT2: i32 = xz_sys::LZMA_MF_BT2 as _; +pub const MF_BT3: i32 = xz_sys::LZMA_MF_BT3 as _; +pub const MF_BT4: i32 = xz_sys::LZMA_MF_BT4 as _; + +pub const MODE_FAST: i32 = xz_sys::LZMA_MODE_FAST as _; +pub const MODE_NORMAL: i32 = xz_sys::LZMA_MODE_NORMAL as _; + +pub const FORMAT_AUTO: i32 = 0; +pub const FORMAT_XZ: i32 = 1; +pub const FORMAT_ALONE: i32 = 2; +pub const FORMAT_RAW: i32 = 3; + +pub const FILTER_LZMA1: u64 = xz_sys::LZMA_FILTER_LZMA1; +pub const FILTER_LZMA2: u64 = xz_sys::LZMA_FILTER_LZMA2; +pub const FILTER_DELTA: u64 = xz_sys::LZMA_FILTER_DELTA; +pub const FILTER_X86: u64 = xz_sys::LZMA_FILTER_X86; +pub const FILTER_POWERPC: u64 = xz_sys::LZMA_FILTER_POWERPC; +pub const FILTER_IA64: u64 = xz_sys::LZMA_FILTER_IA64; +pub const FILTER_ARM: u64 = xz_sys::LZMA_FILTER_ARM; +pub const FILTER_ARMTHUMB: u64 = xz_sys::LZMA_FILTER_ARMTHUMB; +pub const FILTER_SPARC: u64 = xz_sys::LZMA_FILTER_SPARC; + +pub const PRESET_DEFAULT: u32 = xz_sys::LZMA_PRESET_DEFAULT; +pub const PRESET_EXTREME: u32 = xz_sys::LZMA_PRESET_EXTREME; + +const DEFAULT_LC: u32 = xz_sys::LZMA_LC_DEFAULT; +const DEFAULT_LP: u32 = xz_sys::LZMA_LP_DEFAULT; +const DEFAULT_PB: u32 = xz_sys::LZMA_PB_DEFAULT; +const DICT_POW2: [u8; 10] = [18, 20, 21, 22, 22, 23, 23, 24, 25, 26]; + +#[derive(Debug)] +pub enum Error { + Memory, + Value(String), + Lzma(String), + Eof, +} + +impl From for Error { + fn from(err: XzError) -> Self { + match err { + XzError::UnsupportedCheck => Self::Lzma("Unsupported integrity check".to_owned()), + XzError::Mem => Self::Memory, + XzError::MemLimit => Self::Lzma("Memory usage limit exceeded".to_owned()), + XzError::Format => Self::Lzma("Input format not supported by decoder".to_owned()), + XzError::Options => Self::Lzma("Invalid or unsupported options".to_owned()), + XzError::Data | XzError::NoCheck => Self::Lzma("Corrupt input data".to_owned()), + XzError::Program => Self::Lzma("Internal error".to_owned()), + } + } +} + +#[derive(Clone, Debug, Default)] +pub struct FilterSpec { + pub id: u64, + pub preset: Option, + pub dict_size: Option, + pub lc: Option, + pub lp: Option, + pub pb: Option, + pub mode: Option, + pub nice_len: Option, + pub mf: Option, + pub depth: Option, + pub dist: Option, + pub start_offset: Option, +} + +fn int_to_check(check: i32) -> Option { + if check == -1 { + return Some(Check::Crc64); + } + match check { + CHECK_NONE => Some(Check::None), + CHECK_CRC32 => Some(Check::Crc32), + CHECK_CRC64 => Some(Check::Crc64), + CHECK_SHA256 => Some(Check::Sha256), + _ => None, + } +} + +fn u32_to_mode(value: u32) -> Option { + match value as i32 { + MODE_FAST => Some(Mode::Fast), + MODE_NORMAL => Some(Mode::Normal), + _ => None, + } +} + +fn u32_to_mf(value: u32) -> Option { + match value as i32 { + MF_HC3 => Some(MatchFinder::HashChain3), + MF_HC4 => Some(MatchFinder::HashChain4), + MF_BT2 => Some(MatchFinder::BinaryTree2), + MF_BT3 => Some(MatchFinder::BinaryTree3), + MF_BT4 => Some(MatchFinder::BinaryTree4), + _ => None, + } +} + +fn lzma_options(spec: &FilterSpec) -> Result { + let preset = spec.preset.unwrap_or(PRESET_DEFAULT); + let mut options = LzmaOptions::new_preset(preset) + .map_err(|_| Error::Lzma(format!("Invalid compression preset: {preset}")))?; + if let Some(value) = spec.dict_size { + options.dict_size(value); + } + if let Some(value) = spec.lc { + options.literal_context_bits(value); + } + if let Some(value) = spec.lp { + options.literal_position_bits(value); + } + if let Some(value) = spec.pb { + options.position_bits(value); + } + if let Some(value) = spec.mode { + let mode = u32_to_mode(value) + .ok_or_else(|| Error::Value("Invalid filter specifier for LZMA filter".to_owned()))?; + options.mode(mode); + } + if let Some(value) = spec.nice_len { + options.nice_len(value); + } + if let Some(value) = spec.mf { + let mf = u32_to_mf(value) + .ok_or_else(|| Error::Value("Invalid filter specifier for LZMA filter".to_owned()))?; + options.match_finder(mf); + } + if let Some(value) = spec.depth { + options.depth(value); + } + Ok(options) +} + +fn add_bcj_filter(filters: &mut Filters, id: u64, start_offset: u32) -> Result<(), Error> { + if start_offset == 0 { + match id { + FILTER_X86 => filters.x86(), + FILTER_POWERPC => filters.powerpc(), + FILTER_IA64 => filters.ia64(), + FILTER_ARM => filters.arm(), + FILTER_ARMTHUMB => filters.arm_thumb(), + FILTER_SPARC => filters.sparc(), + _ => unreachable!(), + }; + } else { + let properties = start_offset.to_le_bytes(); + match id { + FILTER_X86 => filters.x86_properties(&properties)?, + FILTER_POWERPC => filters.powerpc_properties(&properties)?, + FILTER_IA64 => filters.ia64_properties(&properties)?, + FILTER_ARM => filters.arm_properties(&properties)?, + FILTER_ARMTHUMB => filters.arm_thumb_properties(&properties)?, + FILTER_SPARC => filters.sparc_properties(&properties)?, + _ => unreachable!(), + }; + } + Ok(()) +} + +fn build_filters(specs: &[FilterSpec]) -> Result { + if specs.len() > LZMA_FILTERS_MAX { + return Err(Error::Lzma(format!( + "Too many filters - liblzma supports a maximum of {LZMA_FILTERS_MAX}" + ))); + } + let mut filters = Filters::new(); + for spec in specs { + match spec.id { + FILTER_LZMA1 => { + filters.lzma1(&lzma_options(spec)?); + } + FILTER_LZMA2 => { + filters.lzma2(&lzma_options(spec)?); + } + FILTER_DELTA => { + let dist = spec.dist.unwrap_or(1); + if !(1..=256).contains(&dist) { + return Err(Error::Value( + "Invalid filter specifier for delta filter".to_owned(), + )); + } + filters.delta_properties(&[(dist - 1) as u8])?; + } + FILTER_X86 | FILTER_POWERPC | FILTER_IA64 | FILTER_ARM | FILTER_ARMTHUMB + | FILTER_SPARC => { + add_bcj_filter(&mut filters, spec.id, spec.start_offset.unwrap_or(0))?; + } + id => return Err(Error::Value(format!("Invalid filter ID: {id}"))), + } + } + Ok(filters) +} + +fn preset_dict_size(preset: u32) -> u32 { + let level = (preset & xz_sys::LZMA_PRESET_LEVEL_MASK) as usize; + DICT_POW2.get(level).map_or(0, |power| 1u32 << power) +} + +fn lzma2_dict_size_from_prop(prop: u8) -> u32 { + if prop >= 40 { + return u32::MAX; + } + let prop = u32::from(prop); + (2 | (prop & 1)) << (prop / 2 + 11) +} + +fn lzma2_prop_from_dict_size(dict_size: u32) -> u8 { + if dict_size == u32::MAX { + return 40; + } + (0u8..40) + .find(|&property| lzma2_dict_size_from_prop(property) >= dict_size) + .unwrap_or(40) +} + +pub fn encode_filter_properties(spec: &FilterSpec) -> Result, Error> { + match spec.id { + FILTER_LZMA1 => { + let preset = spec.preset.unwrap_or(PRESET_DEFAULT); + let lc = spec.lc.unwrap_or(DEFAULT_LC); + let lp = spec.lp.unwrap_or(DEFAULT_LP); + let pb = spec.pb.unwrap_or(DEFAULT_PB); + if lc > 4 || lp > 4 || lc + lp > 4 || pb > 4 { + return Err(Error::Lzma("Invalid or unsupported options".to_owned())); + } + let dict_size = spec.dict_size.unwrap_or_else(|| preset_dict_size(preset)); + let mut result = vec![0u8; 5]; + result[0] = ((pb * 5 + lp) * 9 + lc) as u8; + result[1..].copy_from_slice(&dict_size.to_le_bytes()); + Ok(result) + } + FILTER_LZMA2 => { + let preset = spec.preset.unwrap_or(PRESET_DEFAULT); + let dict_size = spec.dict_size.unwrap_or_else(|| preset_dict_size(preset)); + Ok(vec![lzma2_prop_from_dict_size(dict_size)]) + } + FILTER_DELTA => { + let dist = spec.dist.unwrap_or(1); + if !(1..=256).contains(&dist) { + return Err(Error::Value( + "Invalid filter specifier for delta filter".to_owned(), + )); + } + Ok(vec![(dist - 1) as u8]) + } + FILTER_X86 | FILTER_POWERPC | FILTER_IA64 | FILTER_ARM | FILTER_ARMTHUMB | FILTER_SPARC => { + let start_offset = spec.start_offset.unwrap_or(0); + Ok(if start_offset == 0 { + vec![] + } else { + start_offset.to_le_bytes().to_vec() + }) + } + id => Err(Error::Value(format!("Invalid filter ID: {id}"))), + } +} + +pub fn decode_filter_properties(id: u64, properties: &[u8]) -> Result { + let mut spec = FilterSpec { + id, + ..FilterSpec::default() + }; + match id { + FILTER_LZMA1 => { + let [property, a, b, c, d, ..] = properties else { + return Err(Error::Lzma("Invalid or unsupported options".to_owned())); + }; + let mut value = u32::from(*property); + spec.lc = Some(value % 9); + value /= 9; + spec.lp = Some(value % 5); + spec.pb = Some(value / 5); + spec.dict_size = Some(u32::from_le_bytes([*a, *b, *c, *d])); + } + FILTER_LZMA2 => { + let [property] = properties else { + return Err(Error::Lzma("Invalid or unsupported options".to_owned())); + }; + spec.dict_size = Some(lzma2_dict_size_from_prop(*property)); + } + FILTER_DELTA => { + let [property] = properties else { + return Err(Error::Lzma("Invalid or unsupported options".to_owned())); + }; + spec.dist = Some(u32::from(*property) + 1); + } + FILTER_X86 | FILTER_POWERPC | FILTER_IA64 | FILTER_ARM | FILTER_ARMTHUMB | FILTER_SPARC => { + match properties { + [] => {} + [a, b, c, d] => spec.start_offset = Some(u32::from_le_bytes([*a, *b, *c, *d])), + _ => return Err(Error::Lzma("Invalid or unsupported options".to_owned())), + } + } + _ => return Err(Error::Value(format!("Invalid filter ID: {id}"))), + } + Ok(spec) +} + +#[must_use] +pub fn is_check_supported(check_id: i32) -> bool { + unsafe { xz_sys::lzma_check_is_supported(check_id as _) != 0 } +} + +struct LzmaStream { + stream: Stream, + check: i32, + header_buf: [u8; 8], + header_collected: u8, + track_header: bool, +} + +impl LzmaStream { + fn new(stream: Stream, check: i32, track_header: bool) -> Self { + Self { + stream, + check, + header_buf: [0; 8], + header_collected: 0, + track_header, + } + } + + fn process(&mut self, input: &[u8], output: &mut Vec) -> Result { + if self.track_header && self.header_collected < 8 { + let count = (8 - usize::from(self.header_collected)).min(input.len()); + let start = usize::from(self.header_collected); + self.header_buf[start..start + count].copy_from_slice(&input[..count]); + self.header_collected += count as u8; + } + match self.stream.process_vec(input, output, Action::Run) { + Ok(Status::GetCheck) => { + if self.header_collected >= 8 { + self.check = i32::from(self.header_buf[7] & 0x0f); + } + Ok(Status::Ok) + } + Err(XzError::NoCheck) => { + self.check = CHECK_NONE; + Ok(Status::Ok) + } + other => other, + } + } +} + +fn decompress_chunks( + chunks: &mut Chunker<'_>, + stream: &mut LzmaStream, + max_length: Option, +) -> Result<(Vec, bool), XzError> { + if chunks.is_empty() { + return Ok((Vec::new(), true)); + } + let max_length = max_length.unwrap_or(usize::MAX); + let mut output = Vec::new(); + 'outer: loop { + let chunk = chunks.chunk(); + loop { + let additional = BUFSIZ.min(max_length - output.capacity()); + if additional == 0 { + return Ok((output, false)); + } + output.reserve_exact(additional); + let previous_in = stream.stream.total_in(); + let result = stream.process(chunk, &mut output); + let consumed = (stream.stream.total_in() - previous_in) as usize; + chunks.advance(consumed); + let status = result?; + if status == Status::StreamEnd || chunks.is_empty() { + output.shrink_to_fit(); + return Ok((output, status == Status::StreamEnd)); + } + if !chunk.is_empty() && consumed == 0 { + continue; + } + continue 'outer; + } + } +} + +pub struct Decompressor { + stream: LzmaStream, + unused_data: Vec, + input_buffer: Vec, + eof: bool, + needs_input: bool, +} + +impl Decompressor { + pub fn new( + format: i32, + memlimit: Option, + filters: Option>, + ) -> Result { + if format == FORMAT_RAW && memlimit.is_some() { + return Err(Error::Value( + "Cannot specify memory limit with FORMAT_RAW".to_owned(), + )); + } + if format == FORMAT_RAW && filters.is_none() { + return Err(Error::Value( + "Must specify filters for FORMAT_RAW".to_owned(), + )); + } + if format != FORMAT_RAW && filters.is_some() { + return Err(Error::Value( + "Cannot specify filters except with FORMAT_RAW".to_owned(), + )); + } + let memlimit = memlimit.unwrap_or(u64::MAX); + let flags = TELL_ANY_CHECK | TELL_NO_CHECK; + let stream = match format { + FORMAT_AUTO => LzmaStream::new( + Stream::new_auto_decoder(memlimit, flags)?, + CHECK_UNKNOWN, + true, + ), + FORMAT_XZ => LzmaStream::new( + Stream::new_stream_decoder(memlimit, flags)?, + CHECK_UNKNOWN, + true, + ), + FORMAT_ALONE => LzmaStream::new(Stream::new_lzma_decoder(memlimit)?, CHECK_NONE, false), + FORMAT_RAW => { + let filters = build_filters(filters.as_deref().expect("validated raw filters"))?; + LzmaStream::new(Stream::new_raw_decoder(&filters)?, CHECK_NONE, false) + } + _ => return Err(Error::Value(format!("Invalid container format: {format}"))), + }; + Ok(Self { + stream, + unused_data: Vec::new(), + input_buffer: Vec::new(), + eof: false, + needs_input: true, + }) + } + + pub fn decompress(&mut self, data: &[u8], max_length: Option) -> Result, Error> { + if self.eof { + return Err(Error::Eof); + } + let input_buffer = &mut self.input_buffer; + let stream = &mut self.stream; + let mut chunks = Chunker::chain(input_buffer, data); + let previous_len = chunks.len(); + let result = decompress_chunks(&mut chunks, stream, max_length); + let stream_end = match &result { + Ok((_, stream_end)) => *stream_end, + Err(_) => false, + }; + let consumed = previous_len - chunks.len(); + self.eof |= stream_end; + if self.eof { + self.needs_input = false; + if !chunks.is_empty() { + self.unused_data = chunks.to_vec(); + } + } else if chunks.is_empty() { + input_buffer.clear(); + self.needs_input = true; + } else { + self.needs_input = false; + if let Some(consumed_from_data) = consumed.checked_sub(input_buffer.len()) { + input_buffer.clear(); + input_buffer.extend_from_slice(&data[consumed_from_data..]); + } else { + input_buffer.drain(..consumed); + input_buffer.extend_from_slice(data); + } + } + result.map(|(output, _)| output).map_err(Error::from) + } + + #[must_use] + pub fn check(&self) -> i32 { + self.stream.check + } + #[must_use] + pub fn eof(&self) -> bool { + self.eof + } + #[must_use] + pub fn unused_data(&self) -> &[u8] { + &self.unused_data + } + #[must_use] + pub fn needs_input(&self) -> bool { + self.needs_input + } +} + +pub struct Compressor { + stream: Option, +} + +impl Compressor { + pub fn new( + format: i32, + check: i32, + preset: u32, + filters: Option>, + ) -> Result { + if format != FORMAT_XZ && check != -1 && check != CHECK_NONE { + return Err(Error::Lzma( + "Integrity checks are only supported by FORMAT_XZ".to_owned(), + )); + } + let stream = match format { + FORMAT_XZ => { + let check = int_to_check(check) + .ok_or_else(|| Error::Value("Invalid check value".to_owned()))?; + if let Some(specs) = filters { + Stream::new_stream_encoder(&build_filters(&specs)?, check)? + } else { + Stream::new_easy_encoder(preset, check)? + } + } + FORMAT_ALONE => { + let options = match filters { + None => LzmaOptions::new_preset(preset).map_err(|_| { + Error::Lzma(format!("Invalid compression preset: {preset}")) + })?, + Some(specs) => match specs.as_slice() { + [spec] if spec.id == FILTER_LZMA1 => lzma_options(spec)?, + _ => { + return Err(Error::Value( + "Invalid filter chain for FORMAT_ALONE - must be a single LZMA1 filter" + .to_owned(), + )); + } + }, + }; + Stream::new_lzma_encoder(&options)? + } + FORMAT_RAW => { + let specs = filters.ok_or_else(|| { + Error::Value("Must specify filters for FORMAT_RAW".to_owned()) + })?; + Stream::new_raw_encoder(&build_filters(&specs)?)? + } + _ => return Err(Error::Value(format!("Invalid container format: {format}"))), + }; + Ok(Self { + stream: Some(stream), + }) + } + + pub fn compress(&mut self, data: &[u8]) -> Result, Error> { + let stream = self + .stream + .as_mut() + .ok_or_else(|| Error::Lzma(USE_AFTER_FINISH_ERR.to_owned()))?; + let mut output = Vec::new(); + for mut chunk in data.chunks(CHUNKSIZE) { + while !chunk.is_empty() { + output.reserve(DEF_BUF_SIZE); + let previous_in = stream.total_in(); + stream.process_vec(chunk, &mut output, Action::Run)?; + let consumed = (stream.total_in() - previous_in) as usize; + chunk = &chunk[consumed..]; + } + } + output.shrink_to_fit(); + Ok(output) + } + + pub fn flush(&mut self) -> Result, Error> { + let stream = self + .stream + .as_mut() + .ok_or_else(|| Error::Lzma(USE_AFTER_FINISH_ERR.to_owned()))?; + let mut output = Vec::new(); + let status = loop { + if output.len() == output.capacity() { + output.reserve(DEF_BUF_SIZE); + } + let status = stream.process_vec(&[], &mut output, Action::Finish)?; + if output.len() != output.capacity() { + break status; + } + }; + if status == Status::StreamEnd { + self.stream = None; + } + output.shrink_to_fit(); + Ok(output) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn encode_lzma1_properties_rejects_invalid_lclppb() { + for (lc, lp, pb) in [(5, 0, 0), (0, 5, 0), (4, 1, 0), (0, 0, 5)] { + let spec = FilterSpec { + id: FILTER_LZMA1, + lc: Some(lc), + lp: Some(lp), + pb: Some(pb), + ..FilterSpec::default() + }; + assert!(matches!( + encode_filter_properties(&spec), + Err(Error::Lzma(message)) if message == "Invalid or unsupported options" + )); + } + } + + #[test] + fn format_alone_uses_the_single_lzma1_filter() { + let dict_size = 1 << 20; + let spec = FilterSpec { + id: FILTER_LZMA1, + dict_size: Some(dict_size), + ..FilterSpec::default() + }; + let mut compressor = + Compressor::new(FORMAT_ALONE, CHECK_NONE, PRESET_DEFAULT, Some(vec![spec])).unwrap(); + let mut encoded = compressor.compress(b"hello").unwrap(); + encoded.extend(compressor.flush().unwrap()); + assert_eq!(&encoded[1..5], &dict_size.to_le_bytes()); + } + + #[test] + fn format_alone_rejects_other_filter_chains() { + for filters in [ + vec![], + vec![FilterSpec { + id: FILTER_LZMA2, + ..FilterSpec::default() + }], + vec![ + FilterSpec { + id: FILTER_LZMA1, + ..FilterSpec::default() + }, + FilterSpec { + id: FILTER_LZMA1, + ..FilterSpec::default() + }, + ], + ] { + assert!(matches!( + Compressor::new(FORMAT_ALONE, CHECK_NONE, PRESET_DEFAULT, Some(filters)), + Err(Error::Value(message)) + if message + == "Invalid filter chain for FORMAT_ALONE - must be a single LZMA1 filter" + )); + } + } +} diff --git a/crates/common/src/compression/mod.rs b/crates/common/src/compression/mod.rs new file mode 100644 index 00000000000..f0d2b7b5f19 --- /dev/null +++ b/crates/common/src/compression/mod.rs @@ -0,0 +1,78 @@ +// spell-checker:ignore chunker + +//! VM-independent compression engines shared by RustPython components. + +const CHUNKSIZE: usize = u32::MAX as usize; + +/// A two-slice cursor used by streaming decompressors to consume buffered and +/// newly supplied input without joining them first. +#[doc(hidden)] +#[derive(Clone)] +pub struct Chunker<'a> { + data1: &'a [u8], + data2: &'a [u8], +} + +impl<'a> Chunker<'a> { + /// Start a cursor over one input slice. + #[must_use] + pub const fn new(data: &'a [u8]) -> Self { + Self { + data1: data, + data2: &[], + } + } + + /// Chain previously buffered input in front of newly supplied input. + #[must_use] + pub const fn chain(data1: &'a [u8], data2: &'a [u8]) -> Self { + if data1.is_empty() { + Self { + data1: data2, + data2: &[], + } + } else { + Self { data1, data2 } + } + } + + /// Return the number of bytes that have not been consumed. + #[must_use] + pub const fn len(&self) -> usize { + self.data1.len() + self.data2.len() + } + + /// Return whether all input has been consumed. + #[must_use] + pub const fn is_empty(&self) -> bool { + self.data1.is_empty() + } + + /// Copy the remaining input into a contiguous vector. + #[must_use] + pub fn to_vec(&self) -> Vec { + [self.data1, self.data2].concat() + } + + /// Return the next native-codec-sized input chunk. + #[must_use] + pub fn chunk(&self) -> &'a [u8] { + self.data1.get(..CHUNKSIZE).unwrap_or(self.data1) + } + + /// Advance the cursor by `consumed` bytes. + pub fn advance(&mut self, consumed: usize) { + self.data1 = &self.data1[consumed..]; + if self.data1.is_empty() { + self.data1 = core::mem::take(&mut self.data2); + } + } +} + +#[cfg(all( + feature = "lzma", + not(any(target_os = "android", target_arch = "wasm32")) +))] +pub mod lzma; +#[cfg(feature = "zlib")] +pub mod zlib; diff --git a/crates/common/src/compression/zlib.rs b/crates/common/src/compression/zlib.rs new file mode 100644 index 00000000000..224d04a4ac1 --- /dev/null +++ b/crates/common/src/compression/zlib.rs @@ -0,0 +1,1007 @@ +// spell-checker:ignore chunker zdict amet consectetur adipiscing elit + +//! VM-independent zlib DEFLATE engine. +//! +//! The engine owns its native stream state and reports plain Rust errors so +//! interpreter and embedding layers can provide their own object and exception +//! adapters. + +use core::ffi::CStr; +use core::mem::MaybeUninit; + +use zlib_rs::{ + DeflateError, DeflateFlush, InflateError, InflateFlush, ReturnCode, Status, + c_api::z_stream, + deflate::{self, DeflateConfig, DeflateStream, Method, Strategy}, + inflate::{self, InflateConfig, InflateStream}, +}; + +use super::{CHUNKSIZE, Chunker}; + +pub use zlib_rs::c_api::{ + Z_BEST_COMPRESSION, Z_BEST_SPEED, Z_BLOCK, Z_DEFAULT_COMPRESSION, Z_DEFAULT_STRATEGY, + Z_DEFLATED, Z_FILTERED, Z_FINISH, Z_FIXED, Z_FULL_FLUSH, Z_HUFFMAN_ONLY, Z_NO_COMPRESSION, + Z_NO_FLUSH, Z_PARTIAL_FLUSH, Z_RLE, Z_SYNC_FLUSH, Z_TREES, +}; + +pub const MAX_WBITS: i32 = 15; +pub const DEF_BUF_SIZE: usize = 16 * 1024; + +const USE_AFTER_FINISH_ERR: &str = "Error -2: inconsistent stream state"; + +fn valid_level(level: i32) -> bool { + matches!( + level, + Z_DEFAULT_COMPRESSION | Z_NO_COMPRESSION..=Z_BEST_COMPRESSION + ) +} + +#[must_use] +pub const fn version() -> &'static str { + unsafe { + match CStr::from_ptr(libz_rs_sys::zlibVersion()).to_str() { + Ok(version) => version, + Err(_) => unreachable!(), + } + } +} + +const fn z_size_to_u64(value: core::ffi::c_ulong) -> u64 { + #[cfg(all(target_pointer_width = "64", not(windows)))] + { + value + } + #[cfg(any(target_pointer_width = "32", windows))] + { + value as u64 + } +} + +const fn z_checksum_to_u32(value: core::ffi::c_ulong) -> u32 { + #[cfg(all(target_pointer_width = "64", not(windows)))] + { + value as u32 + } + #[cfg(any(target_pointer_width = "32", windows))] + { + value + } +} + +enum InitOptions { + Standard { header: bool, wbits: i32 }, + Gzip { wbits: i32 }, + Auto { wbits: i32 }, + HeaderWindow, +} + +/// Errors raised while constructing a zlib stream. Invalid Python-level +/// options are kept distinct from failures reported by the zlib engine. +#[derive(Debug)] +pub enum InitError { + InvalidOption, + Zlib(String), +} + +impl InitOptions { + fn new(wbits: i32) -> Result { + let header = wbits >= 0; + let wbits = wbits.checked_abs().ok_or(InitError::InvalidOption)?; + match wbits { + 0 if header => Ok(Self::HeaderWindow), + 8..=15 => Ok(Self::Standard { header, wbits }), + 24..=31 if header => Ok(Self::Gzip { wbits: wbits - 16 }), + 40..=47 if header => Ok(Self::Auto { wbits: wbits - 32 }), + _ => Err(InitError::InvalidOption), + } + } + + fn inflate_window_bits(self) -> i32 { + match self { + Self::Standard { header, wbits } => { + if header { + wbits + } else { + -wbits + } + } + Self::Gzip { wbits } => wbits + 16, + Self::Auto { wbits } => wbits + 32, + Self::HeaderWindow => 0, + } + } + + fn deflate_window_bits(self) -> Result { + match self { + Self::Standard { + header: false, + wbits: 8, + } + | Self::Gzip { wbits: 8 } => Err(InitError::InvalidOption), + Self::Standard { header, wbits } => Ok(if header { wbits } else { -wbits }), + Self::Gzip { wbits } => Ok(wbits + 16), + Self::Auto { .. } | Self::HeaderWindow => Err(InitError::InvalidOption), + } + } +} + +fn return_code_message(code: ReturnCode) -> &'static str { + match code { + ReturnCode::Ok => "", + ReturnCode::StreamEnd => "stream end", + ReturnCode::NeedDict => "need dictionary", + ReturnCode::ErrNo => "file error", + ReturnCode::StreamError => "stream error", + ReturnCode::DataError => "data error", + ReturnCode::MemError => "insufficient memory", + ReturnCode::BufError => "buffer error", + ReturnCode::VersionError => "incompatible version", + } +} + +fn stream_message(stream: &z_stream) -> Option<&str> { + if stream.msg.is_null() { + None + } else { + unsafe { CStr::from_ptr(stream.msg).to_str().ok() } + } +} + +/// The public `zlib_rs::Deflate` wrapper exposes only the common constructor. +/// Python's streaming API needs the full zlib configuration and exact native +/// state copies, so the engine owns the crate's low-level `z_stream` instead. +struct RawDeflate { + stream: z_stream, +} + +// The raw pointers belong to this stream and are never accessed without the +// mutable owner. Moving the owner between threads does not share the stream. +unsafe impl Send for RawDeflate {} + +impl RawDeflate { + fn new(config: DeflateConfig) -> Result { + let mut stream = z_stream::default(); + let code = deflate::init(&mut stream, config); + if code == ReturnCode::Ok { + Ok(Self { stream }) + } else { + Err(return_code_message(code).to_owned()) + } + } + + fn total_in(&self) -> u64 { + z_size_to_u64(self.stream.total_in) + } + + fn total_out(&self) -> u64 { + z_size_to_u64(self.stream.total_out) + } + + fn compress( + &mut self, + input: &[u8], + output: &mut [u8], + flush: DeflateFlush, + ) -> Result { + self.stream.avail_in = input.len().min(u32::MAX as usize) as u32; + self.stream.avail_out = output.len().min(u32::MAX as usize) as u32; + self.stream.next_in = input.as_ptr(); + self.stream.next_out = output.as_mut_ptr(); + let stream = unsafe { DeflateStream::from_stream_mut(&mut self.stream) } + .expect("initialized deflate stream"); + match deflate::deflate(stream, flush) { + ReturnCode::Ok => Ok(Status::Ok), + ReturnCode::StreamEnd => Ok(Status::StreamEnd), + ReturnCode::BufError => Ok(Status::BufError), + ReturnCode::StreamError => Err(DeflateError::StreamError), + ReturnCode::DataError => Err(DeflateError::DataError), + ReturnCode::MemError => Err(DeflateError::MemError), + ReturnCode::NeedDict => unreachable!("compression does not use dictionaries here"), + ReturnCode::ErrNo | ReturnCode::VersionError => { + unreachable!("pure-Rust deflate returned an unsupported status") + } + } + } + + fn set_dictionary(&mut self, dictionary: &[u8]) -> Result<(), DeflateError> { + let stream = unsafe { DeflateStream::from_stream_mut(&mut self.stream) } + .expect("initialized deflate stream"); + match deflate::set_dictionary(stream, dictionary) { + ReturnCode::Ok => Ok(()), + ReturnCode::StreamError => Err(DeflateError::StreamError), + ReturnCode::DataError => Err(DeflateError::DataError), + ReturnCode::MemError => Err(DeflateError::MemError), + code => unreachable!("deflate set_dictionary returned {code:?}"), + } + } + + fn copy(&mut self) -> Result { + let mut destination = MaybeUninit::>::uninit(); + let code = { + let source = unsafe { DeflateStream::from_stream_mut(&mut self.stream) } + .expect("initialized deflate stream"); + deflate::copy(&mut destination, source) + }; + if code != ReturnCode::Ok { + return Err(return_code_message(code).to_owned()); + } + let copied = unsafe { destination.assume_init() }; + let stream = unsafe { core::mem::transmute::, z_stream>(copied) }; + Ok(Self { stream }) + } +} + +impl Drop for RawDeflate { + fn drop(&mut self) { + if let Some(stream) = unsafe { DeflateStream::from_stream_mut(&mut self.stream) } { + let _ = deflate::end(stream); + } + } +} + +/// Low-level inflate owner paired with [`RawDeflate`]. `inflate::copy` fixes +/// every internal allocation in the clone, exactly like zlib's `inflateCopy`; +/// no replay log or side table is involved. +struct RawInflate { + stream: z_stream, +} + +// See `RawDeflate`: this is an owned, movable stream, not a shared handle. +unsafe impl Send for RawInflate {} + +impl RawInflate { + fn new(window_bits: i32) -> Result { + let mut stream = z_stream::default(); + let code = inflate::init(&mut stream, InflateConfig { window_bits }); + if code == ReturnCode::Ok { + // zlib-rs' low-level `inflate::copy` requires a non-null next_out + // even when avail_out is zero. Its stable wrapper gets exactly + // such a dangling empty-slice pointer from `Writer::new(&mut [])`; + // reproduce that initialized-stream shape so a pristine + // `Decompress.copy()` succeeds before the first input byte. + stream.next_out = core::ptr::NonNull::::dangling().as_ptr(); + Ok(Self { stream }) + } else { + Err(return_code_message(code).to_owned()) + } + } + + fn total_in(&self) -> u64 { + z_size_to_u64(self.stream.total_in) + } + + fn total_out(&self) -> u64 { + z_size_to_u64(self.stream.total_out) + } + + fn error_message(&self) -> Option<&str> { + stream_message(&self.stream) + } + + fn decompress( + &mut self, + input: &[u8], + output: &mut [u8], + flush: InflateFlush, + ) -> Result { + self.stream.avail_in = input.len().min(u32::MAX as usize) as u32; + self.stream.avail_out = output.len().min(u32::MAX as usize) as u32; + self.stream.next_in = input.as_ptr(); + self.stream.next_out = output.as_mut_ptr(); + let stream = unsafe { InflateStream::from_stream_mut(&mut self.stream) } + .expect("initialized inflate stream"); + match unsafe { inflate::inflate(stream, flush) } { + ReturnCode::Ok => Ok(Status::Ok), + ReturnCode::StreamEnd => Ok(Status::StreamEnd), + ReturnCode::BufError => Ok(Status::BufError), + ReturnCode::NeedDict => Err(InflateError::NeedDict { + dict_id: z_checksum_to_u32(self.stream.adler), + }), + ReturnCode::StreamError => Err(InflateError::StreamError), + ReturnCode::DataError => Err(InflateError::DataError), + ReturnCode::MemError => Err(InflateError::MemError), + ReturnCode::ErrNo | ReturnCode::VersionError => { + unreachable!("pure-Rust inflate returned an unsupported status") + } + } + } + + fn set_dictionary(&mut self, dictionary: &[u8]) -> Result<(), InflateError> { + let stream = unsafe { InflateStream::from_stream_mut(&mut self.stream) } + .expect("initialized inflate stream"); + match inflate::set_dictionary(stream, dictionary) { + ReturnCode::Ok => Ok(()), + ReturnCode::StreamError => Err(InflateError::StreamError), + ReturnCode::DataError => Err(InflateError::DataError), + code => unreachable!("inflate set_dictionary returned {code:?}"), + } + } + + fn copy(&self) -> Result { + let source = unsafe { InflateStream::from_stream_ref(&self.stream) } + .expect("initialized inflate stream"); + let mut destination = MaybeUninit::>::uninit(); + let code = unsafe { inflate::copy(&mut destination, source) }; + if code != ReturnCode::Ok { + return Err(return_code_message(code).to_owned()); + } + let copied = unsafe { destination.assume_init() }; + let stream = unsafe { core::mem::transmute::, z_stream>(copied) }; + Ok(Self { stream }) + } +} + +impl Drop for RawInflate { + fn drop(&mut self) { + if let Some(stream) = unsafe { InflateStream::from_stream_mut(&mut self.stream) } { + inflate::end(stream); + } + } +} + +/// Drive an inflate stream over chained input, retrying with an optional +/// preset dictionary when the stream requests one. +fn decompress_chunks( + data: &mut Chunker<'_>, + d: &mut RawInflate, + zdict: Option<&[u8]>, + bufsize: usize, + max_length: Option, + calc_flush: impl Fn(bool) -> InflateFlush, +) -> Result<(Vec, bool), String> { + if data.is_empty() { + return Ok((Vec::new(), false)); + } + let max_length = max_length.unwrap_or(usize::MAX); + let mut buf = Vec::new(); + + 'outer: loop { + let chunk = data.chunk(); + let flush = calc_flush(chunk.len() == data.len()); + loop { + let additional = core::cmp::min(bufsize, max_length - buf.len()); + if additional == 0 { + return Ok((buf, false)); + } + let mut output = vec![0; additional]; + + let prev_in = d.total_in(); + let prev_out = d.total_out(); + let res = d.decompress(chunk, &mut output, flush); + let consumed = d.total_in() - prev_in; + let produced = (d.total_out() - prev_out) as usize; + buf.extend_from_slice(&output[..produced]); + + data.advance(consumed as usize); + + match res { + Ok(status) => { + let stream_end = status == Status::StreamEnd; + if stream_end || data.is_empty() { + buf.shrink_to_fit(); + return Ok((buf, stream_end)); + } else if !chunk.is_empty() && consumed == 0 { + continue; + } + continue 'outer; + } + Err(e) => { + // maybe_set_dict: retry once with the stored dictionary. + match zdict.filter(|_| matches!(e, InflateError::NeedDict { .. })) { + Some(zd) => { + d.set_dictionary(zd).map_err(|e| e.as_str().to_owned())?; + continue 'outer; + } + None => { + return Err(d.error_message().unwrap_or(e.as_str()).to_owned()); + } + } + } + } + } + } +} + +fn decompress_all( + data: &[u8], + d: &mut RawInflate, + zdict: Option<&[u8]>, + bufsize: usize, + max_length: Option, + calc_flush: impl Fn(bool) -> InflateFlush, +) -> Result<(Vec, bool), String> { + let mut chunker = Chunker::new(data); + decompress_chunks(&mut chunker, d, zdict, bufsize, max_length, calc_flush) +} + +/// `zlib.compress(data, level, wbits)`. +pub fn compress(data: &[u8], level: i32, wbits: i32) -> Result, InitError> { + if !valid_level(level) { + return Err(InitError::Zlib("Bad compression level".to_owned())); + } + let mut compressor = Compressor::new(level, 8, wbits, 8, 0, None)?; + let mut output = compressor.compress(data).map_err(InitError::Zlib)?; + output.extend(compressor.flush(Z_FINISH).map_err(InitError::Zlib)?); + Ok(output) +} + +/// `zlib.decompress(data, wbits, bufsize)`. +pub fn decompress(data: &[u8], wbits: i32, bufsize: usize) -> Result, InitError> { + let mut d = + RawInflate::new(InitOptions::new(wbits)?.inflate_window_bits()).map_err(InitError::Zlib)?; + let (buf, stream_end) = decompress_all(data, &mut d, None, bufsize, None, |_| { + InflateFlush::SyncFlush + }) + .map_err(InitError::Zlib)?; + if !stream_end { + return Err(InitError::Zlib( + "Error -5 while decompressing data: incomplete or truncated stream".to_owned(), + )); + } + Ok(buf) +} + +pub struct Compressor { + compress: Option, +} + +impl Compressor { + pub fn new( + level: i32, + method: i32, + wbits: i32, + mem_level: i32, + strategy: i32, + zdict: Option<&[u8]>, + ) -> Result { + if !valid_level(level) { + return Err(InitError::InvalidOption); + } + if method != Z_DEFLATED { + return Err(InitError::InvalidOption); + } + let method = Method::try_from(method).map_err(|_| InitError::InvalidOption)?; + let strategy = Strategy::try_from(strategy).map_err(|_| InitError::InvalidOption)?; + if !(1..=9).contains(&mem_level) { + return Err(InitError::InvalidOption); + } + let window_bits = InitOptions::new(wbits)?.deflate_window_bits()?; + let mut compress = RawDeflate::new(DeflateConfig { + level, + method, + window_bits, + mem_level, + strategy, + }) + .map_err(InitError::Zlib)?; + if let Some(zdict) = zdict { + compress + .set_dictionary(zdict) + .map_err(|e| InitError::Zlib(e.as_str().to_owned()))?; + } + Ok(Self { + compress: Some(compress), + }) + } + + pub fn compress(&mut self, data: &[u8]) -> Result, String> { + let compressor = self + .compress + .as_mut() + .ok_or_else(|| USE_AFTER_FINISH_ERR.to_owned())?; + let mut buf = Vec::new(); + for mut chunk in data.chunks(CHUNKSIZE) { + while !chunk.is_empty() { + let mut output = [0; DEF_BUF_SIZE]; + let prev_in = compressor.total_in(); + let prev_out = compressor.total_out(); + compressor + .compress(chunk, &mut output, DeflateFlush::NoFlush) + .map_err(|e| e.as_str().to_owned())?; + let consumed = (compressor.total_in() - prev_in) as usize; + let produced = (compressor.total_out() - prev_out) as usize; + buf.extend_from_slice(&output[..produced]); + chunk = &chunk[consumed..]; + } + } + buf.shrink_to_fit(); + Ok(buf) + } + + /// Returns the flushed bytes; `finished` is true once a `Z_FINISH` flush + /// has consumed the stream (the object may no longer be used). + pub fn flush(&mut self, mode: i32) -> Result, String> { + let flush = match mode { + Z_NO_FLUSH => return Ok(vec![]), + Z_PARTIAL_FLUSH => DeflateFlush::PartialFlush, + Z_SYNC_FLUSH => DeflateFlush::SyncFlush, + Z_FULL_FLUSH => DeflateFlush::FullFlush, + Z_FINISH => DeflateFlush::Finish, + 5 => DeflateFlush::Block, + _ => return Err("invalid mode".to_owned()), + }; + let compressor = self + .compress + .as_mut() + .ok_or_else(|| USE_AFTER_FINISH_ERR.to_owned())?; + let mut buf = Vec::new(); + let status = loop { + let mut output = [0; DEF_BUF_SIZE]; + let prev_out = compressor.total_out(); + let status = compressor + .compress(&[], &mut output, flush) + .map_err(|e| e.as_str().to_owned())?; + let produced = (compressor.total_out() - prev_out) as usize; + buf.extend_from_slice(&output[..produced]); + if produced != output.len() || status == Status::StreamEnd { + break status; + } + }; + if status == Status::StreamEnd { + if mode == Z_FINISH { + self.compress = None; + } else { + return Err("unexpected eof".to_owned()); + } + } + buf.shrink_to_fit(); + Ok(buf) + } + + #[must_use] + pub fn is_finished(&self) -> bool { + self.compress.is_none() + } + + /// Clone the live native deflate stream. The caller must serialize access + /// to this object while copying. + pub fn copy(&mut self) -> Result { + let compress = self + .compress + .as_mut() + .ok_or_else(|| "Compressor was already flushed".to_owned())? + .copy()?; + Ok(Self { + compress: Some(compress), + }) + } +} + +pub struct Decompressor { + decompress: Option, + zdict: Option>, + eof: bool, + unused_data: Vec, + unconsumed_tail: Vec, +} + +impl Decompressor { + pub fn new(wbits: i32, zdict: Option>) -> Result { + let mut decompress = RawInflate::new(InitOptions::new(wbits)?.inflate_window_bits()) + .map_err(InitError::Zlib)?; + if let Some(d) = &zdict + && wbits < 0 + { + decompress + .set_dictionary(d) + .map_err(|e| InitError::Zlib(e.as_str().to_owned()))?; + } + Ok(Self { + decompress: Some(decompress), + zdict, + eof: false, + unused_data: Vec::new(), + unconsumed_tail: Vec::new(), + }) + } + + #[must_use] + pub fn eof(&self) -> bool { + self.eof + } + #[must_use] + pub fn unused_data(&self) -> &[u8] { + &self.unused_data + } + #[must_use] + pub fn unconsumed_tail(&self) -> &[u8] { + &self.unconsumed_tail + } + #[must_use] + pub fn is_finished(&self) -> bool { + self.decompress.is_none() + } + + /// Clone the inflate stream and all externally visible buffered state. + pub fn copy(&self) -> Result { + let decompress = self + .decompress + .as_ref() + .ok_or_else(|| "Decompressor was already flushed".to_owned())? + .copy()?; + Ok(Self { + decompress: Some(decompress), + zdict: self.zdict.clone(), + eof: self.eof, + unused_data: self.unused_data.clone(), + unconsumed_tail: self.unconsumed_tail.clone(), + }) + } + + /// Run the streaming decompressor and update its unused-input state. + fn decompress_inner( + &mut self, + data: &[u8], + bufsize: usize, + max_length: Option, + is_flush: bool, + ) -> Result<(Result, String>, bool), String> { + let Self { + decompress, + zdict, + unused_data, + unconsumed_tail, + .. + } = self; + let Some(d) = decompress.as_mut() else { + return Err(USE_AFTER_FINISH_ERR.to_owned()); + }; + + let prev_in = d.total_in(); + let res = if is_flush { + // ignore zdict on a flush, finish on the final chunk + let calc_flush = |final_chunk| { + if final_chunk { + InflateFlush::Finish + } else { + InflateFlush::NoFlush + } + }; + decompress_all(data, d, None, bufsize, max_length, calc_flush) + } else { + decompress_all(data, d, zdict.as_deref(), bufsize, max_length, |_| { + InflateFlush::SyncFlush + }) + }; + let (ret, stream_end) = match res { + Ok((buf, stream_end)) => (Ok(buf), stream_end), + Err(err) => (Err(err), false), + }; + let consumed = (d.total_in() - prev_in) as usize; + + // save unused input + let unconsumed = &data[consumed..]; + if !unconsumed.is_empty() { + if stream_end { + unused_data.extend_from_slice(unconsumed); + } else { + *unconsumed_tail = unconsumed.to_vec(); + } + } else if !unconsumed_tail.is_empty() { + unconsumed_tail.clear(); + } + + Ok((ret, stream_end)) + } + + /// `Decompress.decompress(data, max_length)`; `max_length` of `None` is + /// unlimited. + pub fn decompress( + &mut self, + data: &[u8], + max_length: Option, + ) -> Result, String> { + let (ret, stream_end) = self.decompress_inner(data, DEF_BUF_SIZE, max_length, false)?; + self.eof |= stream_end; + ret + } + + /// `Decompress.flush(length)`. + pub fn flush(&mut self, length: usize) -> Result, String> { + let data = core::mem::take(&mut self.unconsumed_tail); + let (ret, stream_end) = self.decompress_inner(&data, length, None, true)?; + self.eof |= stream_end; + if self.eof { + self.decompress = None; + } + ret + } +} + +/// Error surface of [`ZlibDecompressor::decompress`]: `Zlib` maps to +/// `zlib.error`, `Eof` to `EOFError` ("End of stream already reached"). +#[derive(Debug)] +pub enum DecompressError { + Zlib(String), + Eof, +} + +pub struct ZlibDecompressor { + decompress: Option, + zdict: Option>, + unused_data: Vec, + input_buffer: Vec, + eof: bool, + needs_input: bool, +} + +impl ZlibDecompressor { + pub fn new(wbits: i32, zdict: Option>) -> Result { + let mut decompress = RawInflate::new(InitOptions::new(wbits)?.inflate_window_bits()) + .map_err(InitError::Zlib)?; + if let Some(d) = &zdict + && wbits < 0 + { + decompress + .set_dictionary(d) + .map_err(|e| InitError::Zlib(e.as_str().to_owned()))?; + } + Ok(Self { + decompress: Some(decompress), + zdict, + unused_data: Vec::new(), + input_buffer: Vec::new(), + eof: false, + needs_input: true, + }) + } + + #[must_use] + pub fn eof(&self) -> bool { + self.eof + } + #[must_use] + pub fn unused_data(&self) -> &[u8] { + &self.unused_data + } + #[must_use] + pub fn needs_input(&self) -> bool { + self.needs_input + } + + /// Decompress from new input and any bytes buffered by an earlier call. + pub fn decompress( + &mut self, + data: &[u8], + max_length: Option, + ) -> Result, DecompressError> { + if self.eof { + return Err(DecompressError::Eof); + } + + let Self { + decompress: decompress_slot, + zdict, + unused_data, + input_buffer, + eof, + needs_input, + } = self; + let decompress = decompress_slot.as_mut().ok_or(DecompressError::Eof)?; + + let mut chunks = Chunker::chain(input_buffer.as_slice(), data); + + let prev_len = chunks.len(); + let (ret, stream_end) = match decompress_chunks( + &mut chunks, + decompress, + zdict.as_deref(), + DEF_BUF_SIZE, + max_length, + |_| InflateFlush::SyncFlush, + ) { + Ok((buf, stream_end)) => (Ok(buf), stream_end), + Err(err) => (Err(err), false), + }; + let consumed = prev_len - chunks.len(); + + *eof |= stream_end; + + if *eof { + *needs_input = false; + if !chunks.is_empty() { + *unused_data = chunks.to_vec(); + } + // Release the native stream immediately at EOF instead of keeping + // a full inflate window alive until the wrapper is dropped. + *decompress_slot = None; + } else if chunks.is_empty() { + input_buffer.clear(); + *needs_input = true; + } else { + *needs_input = false; + if let Some(n_consumed_from_data) = consumed.checked_sub(input_buffer.len()) { + input_buffer.clear(); + input_buffer.extend_from_slice(&data[n_consumed_from_data..]); + } else { + input_buffer.drain(..consumed); + input_buffer.extend_from_slice(data); + } + } + + ret.map_err(DecompressError::Zlib) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn roundtrip_default() { + let data = b"Lorem ipsum dolor sit amet, consectetur adipiscing elit"; + let c = compress(data, -1, MAX_WBITS).unwrap(); + let d = decompress(&c, MAX_WBITS, DEF_BUF_SIZE).unwrap(); + assert_eq!(d, data); + } + + #[test] + fn roundtrip_all_levels_and_wbits() { + let data = b"the quick brown fox jumps over the lazy dog".repeat(20); + for level in 0..=9 { + for &wbits in &[15i32, 31, -15] { + let c = compress(&data, level, wbits).unwrap(); + let d = decompress(&c, wbits, DEF_BUF_SIZE).unwrap(); + assert_eq!(d, data, "level={level} wbits={wbits}"); + } + } + } + + #[test] + fn bad_level_rejected() { + assert!(compress(b"x", 10, MAX_WBITS).is_err()); + assert!(compress(b"x", -40, MAX_WBITS).is_err()); + } + + #[test] + fn streaming_roundtrip() { + let mut co = Compressor::new(-1, 8, MAX_WBITS, 8, 0, None).unwrap(); + let mut out = co.compress(b"hello ").unwrap(); + out.extend(co.compress(b"world").unwrap()); + out.extend(co.flush(Z_FINISH).unwrap()); + assert!(co.is_finished()); + + let mut do_ = Decompressor::new(MAX_WBITS, None).unwrap(); + let got = do_.decompress(&out, None).unwrap(); + assert_eq!(got, b"hello world"); + assert!(do_.eof()); + } + + #[test] + fn streaming_block_flush_roundtrip() { + let data = b"block-flush payload ".repeat(200); + let mut co = Compressor::new(6, 8, MAX_WBITS, 8, 0, None).unwrap(); + let mut encoded = co.compress(&data[..1000]).unwrap(); + encoded.extend(co.flush(5).unwrap()); + encoded.extend(co.compress(&data[1000..]).unwrap()); + encoded.extend(co.flush(Z_FINISH).unwrap()); + assert_eq!(decompress(&encoded, MAX_WBITS, DEF_BUF_SIZE).unwrap(), data); + } + + #[test] + fn compressor_options_are_validated_by_deflate_init() { + for result in [ + Compressor::new(6, 7, 15, 8, 0, None), + Compressor::new(6, 8, 15, 0, 0, None), + Compressor::new(6, 8, 15, 8, 99, None), + Compressor::new(6, 8, -8, 8, 0, None), + Compressor::new(6, 8, 24, 8, 0, None), + ] { + assert!(matches!(result, Err(InitError::InvalidOption))); + } + assert!(matches!( + compress(b"x", 10, MAX_WBITS), + Err(InitError::Zlib(message)) if message == "Bad compression level" + )); + } + + #[test] + fn negative_gzip_and_auto_wbits_are_rejected() { + for wbits in [-25, -31, -40, -47] { + assert!(matches!( + Decompressor::new(wbits, None), + Err(InitError::InvalidOption) + )); + } + } + + #[test] + fn compressor_copy_forks_native_stream() { + let prefix = b"shared prefix ".repeat(50); + let left = b"left branch".repeat(20); + let right = b"right branch".repeat(20); + let mut original = Compressor::new(6, 8, 15, 8, 0, None).unwrap(); + let shared = original.compress(&prefix).unwrap(); + let mut copied = original.copy().unwrap(); + + let mut left_encoded = shared.clone(); + left_encoded.extend(original.compress(&left).unwrap()); + left_encoded.extend(original.flush(Z_FINISH).unwrap()); + let mut right_encoded = shared; + right_encoded.extend(copied.compress(&right).unwrap()); + right_encoded.extend(copied.flush(Z_FINISH).unwrap()); + + assert_eq!( + decompress(&left_encoded, 15, DEF_BUF_SIZE).unwrap(), + [prefix.as_slice(), left.as_slice()].concat() + ); + assert_eq!( + decompress(&right_encoded, 15, DEF_BUF_SIZE).unwrap(), + [prefix.as_slice(), right.as_slice()].concat() + ); + } + + #[test] + fn decompressor_copy_forks_native_stream() { + let data = b"decompress copy payload ".repeat(100); + let encoded = compress(&data, 6, 15).unwrap(); + let mut original = Decompressor::new(15, None).unwrap(); + let prefix = original.decompress(&encoded[..16], None).unwrap(); + let mut copied = original.copy().unwrap(); + let left = original.decompress(&encoded[16..], None).unwrap(); + let right = copied.decompress(&encoded[16..], None).unwrap(); + assert_eq!([prefix.as_slice(), left.as_slice()].concat(), data); + assert_eq!(left, right); + } + + #[test] + fn pristine_streams_can_be_copied() { + let mut compressor = Compressor::new(6, 8, 15, 8, 0, None).unwrap(); + compressor.copy().unwrap(); + let decompressor = Decompressor::new(15, None).unwrap(); + decompressor.copy().unwrap(); + } + + #[test] + fn header_window_autodetection() { + let data = b"window-sized-from-zlib-header".repeat(20); + let encoded = compress(&data, 1, MAX_WBITS).unwrap(); + assert_eq!(decompress(&encoded, 0, DEF_BUF_SIZE).unwrap(), data); + } + + #[test] + fn buffered_decompressor_gzip_roundtrip() { + // gzip uses _ZlibDecompressor(wbits=-MAX_WBITS) over raw deflate. + let raw = compress(b"gzip payload contents", -1, -15).unwrap(); + let mut d = ZlibDecompressor::new(-15, None).unwrap(); + let got = d.decompress(&raw, None).unwrap(); + assert_eq!(got, b"gzip payload contents"); + assert!(d.eof()); + // decompress after eof raises Eof + assert!(matches!(d.decompress(b"", None), Err(DecompressError::Eof))); + } + + #[test] + fn buffered_decompressor_incremental_needs_input() { + let full = compress(&b"streamed content ".repeat(100), -1, MAX_WBITS).unwrap(); + let mut d = ZlibDecompressor::new(MAX_WBITS, None).unwrap(); + let mut out = Vec::new(); + for byte in full.chunks(1) { + out.extend(d.decompress(byte, None).unwrap()); + } + assert_eq!(out, b"streamed content ".repeat(100)); + assert!(d.eof()); + } + + #[test] + fn streaming_max_length_unconsumed_tail() { + let full = compress(&b"abcdefghij".repeat(50), -1, MAX_WBITS).unwrap(); + let mut d = Decompressor::new(MAX_WBITS, None).unwrap(); + let first = d.decompress(&full, Some(5)).unwrap(); + assert_eq!(first.len(), 5); + assert!(!d.unconsumed_tail().is_empty()); + let rest = d.flush(DEF_BUF_SIZE).unwrap(); + assert_eq!([first, rest].concat(), b"abcdefghij".repeat(50)); + assert!(d.eof()); + } + + #[test] + fn empty_input_does_not_finish_a_stream() { + let encoded = compress(b"later input", -1, MAX_WBITS).unwrap(); + let mut d = Decompressor::new(MAX_WBITS, None).unwrap(); + assert_eq!(d.decompress(b"", None).unwrap(), b""); + assert!(!d.eof()); + assert_eq!(d.decompress(&encoded, None).unwrap(), b"later input"); + assert!(d.eof()); + } +} diff --git a/crates/common/src/lib.rs b/crates/common/src/lib.rs index 8bfde803090..575bf64e365 100644 --- a/crates/common/src/lib.rs +++ b/crates/common/src/lib.rs @@ -11,6 +11,8 @@ pub mod binascii; pub mod borrow; pub mod boxvec; pub mod cformat; +#[cfg(any(feature = "lzma", feature = "zlib"))] +pub mod compression; pub mod encodings; pub mod float_ops; pub mod format; diff --git a/crates/stdlib/Cargo.toml b/crates/stdlib/Cargo.toml index ec1c6ce0167..cd955597d94 100644 --- a/crates/stdlib/Cargo.toml +++ b/crates/stdlib/Cargo.toml @@ -30,7 +30,7 @@ __ssl-rustls = ["ssl", "rustls", "rustls-native-certs", "rustls-pemfile", "rustl # rustpython crates rustpython-derive = { workspace = true } rustpython-vm = { workspace = true, default-features = false, features = ["compiler"]} -rustpython-common = { workspace = true, features = ["binascii", "cjk-codecs", "inet", "json"] } +rustpython-common = { workspace = true, features = ["binascii", "cjk-codecs", "inet", "json", "lzma", "zlib"] } rustpython-host_env = { workspace = true } rustpython-unicode = { workspace = true } @@ -79,8 +79,6 @@ constant_time_eq = { workspace = true } # compression adler32 = { workspace = true } -flate2 = { workspace = true, features = ["zlib-rs"] } -libz-rs-sys = { workspace = true } bzip2 = { workspace = true } # tkinter @@ -117,8 +115,6 @@ pkcs8 = { workspace = true, features = ["encryption", "pkcs5", "pem"], optional [target.'cfg(not(any(target_os = "android", target_arch = "wasm32")))'.dependencies] libsqlite3-sys = { workspace = true, features = ["bundled"], optional = true } -xz = { workspace = true } -xz-sys = { workspace = true } [target.'cfg(windows)'.dependencies] paste = { workspace = true } diff --git a/crates/stdlib/src/compression.rs b/crates/stdlib/src/compression.rs index 3c1160e1f0d..563a32f7b35 100644 --- a/crates/stdlib/src/compression.rs +++ b/crates/stdlib/src/compression.rs @@ -4,14 +4,11 @@ use crate::vm::function::{ArgBytesLike, ArgSize, OptionalArg}; use crate::vm::{ - PyResult, VirtualMachine, + VirtualMachine, builtins::{PyBaseExceptionRef, PyBytesRef}, convert::ToPyException, }; - -pub(crate) const USE_AFTER_FINISH_ERR: &str = "Error -2: inconsistent stream state"; -// TODO: don't hardcode -const CHUNKSIZE: usize = u32::MAX as usize; +use rustpython_common::compression::Chunker; #[derive(FromArgs)] pub(crate) struct DecompressArgs { @@ -70,59 +67,6 @@ pub(crate) const fn flush_sync(_final_chunk: bool) -> T T::SYNC } -#[derive(Clone)] -pub(crate) struct Chunker<'a> { - data1: &'a [u8], - data2: &'a [u8], -} -impl<'a> Chunker<'a> { - pub(crate) const fn new(data: &'a [u8]) -> Self { - Self { - data1: data, - data2: &[], - } - } - pub(crate) const fn chain(data1: &'a [u8], data2: &'a [u8]) -> Self { - if data1.is_empty() { - Self { - data1: data2, - data2: &[], - } - } else { - Self { data1, data2 } - } - } - pub(crate) const fn len(&self) -> usize { - self.data1.len() + self.data2.len() - } - pub(crate) const fn is_empty(&self) -> bool { - self.data1.is_empty() - } - pub(crate) fn to_vec(&self) -> Vec { - [self.data1, self.data2].concat() - } - pub(crate) fn chunk(&self) -> &'a [u8] { - self.data1.get(..CHUNKSIZE).unwrap_or(self.data1) - } - pub(crate) fn advance(&mut self, consumed: usize) { - self.data1 = &self.data1[consumed..]; - if self.data1.is_empty() { - self.data1 = core::mem::take(&mut self.data2); - } - } -} - -pub(crate) fn _decompress( - data: &[u8], - d: &mut D, - bufsize: usize, - max_length: Option, - calc_flush: impl Fn(bool) -> D::Flush, -) -> Result<(Vec, bool), D::Error> { - let mut data = Chunker::new(data); - _decompress_chunks(&mut data, d, bufsize, max_length, calc_flush) -} - pub(crate) fn _decompress_chunks( data: &mut Chunker<'_>, d: &mut D, @@ -176,101 +120,6 @@ pub(crate) fn _decompress_chunks( } } -pub(crate) trait Compressor { - type Status: CompressStatusKind; - type Flush: CompressFlushKind; - const CHUNKSIZE: usize; - const DEF_BUF_SIZE: usize; - - fn compress_vec( - &mut self, - input: &[u8], - output: &mut Vec, - flush: Self::Flush, - vm: &VirtualMachine, - ) -> PyResult; - - fn total_in(&mut self) -> usize; - - fn new_error(message: impl Into, vm: &VirtualMachine) -> PyBaseExceptionRef; -} - -pub(crate) trait CompressFlushKind: Copy { - const NONE: Self; - const FINISH: Self; - - fn to_usize(self) -> usize; -} - -pub(crate) trait CompressStatusKind: Copy { - const EOF: Self; - - fn to_usize(self) -> usize; -} - -#[derive(Debug)] -pub(crate) struct CompressState { - compressor: Option, -} - -impl CompressState { - pub(crate) const fn new(compressor: C) -> Self { - Self { - compressor: Some(compressor), - } - } - - fn get_compressor(&mut self, vm: &VirtualMachine) -> PyResult<&mut C> { - self.compressor - .as_mut() - .ok_or_else(|| C::new_error(USE_AFTER_FINISH_ERR, vm)) - } - - pub(crate) fn compress(&mut self, data: &[u8], vm: &VirtualMachine) -> PyResult> { - let mut buf = Vec::new(); - let compressor = self.get_compressor(vm)?; - - for mut chunk in data.chunks(C::CHUNKSIZE) { - while !chunk.is_empty() { - buf.reserve(C::DEF_BUF_SIZE); - let prev_in = compressor.total_in(); - compressor.compress_vec(chunk, &mut buf, C::Flush::NONE, vm)?; - let consumed = compressor.total_in() - prev_in; - chunk = &chunk[consumed..]; - } - } - - buf.shrink_to_fit(); - Ok(buf) - } - - pub(crate) fn flush(&mut self, mode: C::Flush, vm: &VirtualMachine) -> PyResult> { - let mut buf = Vec::new(); - let compressor = self.get_compressor(vm)?; - - let status = loop { - if buf.len() == buf.capacity() { - buf.reserve(C::DEF_BUF_SIZE); - } - let status = compressor.compress_vec(&[], &mut buf, mode, vm)?; - if buf.len() != buf.capacity() { - break status; - } - }; - - if status.to_usize() == C::Status::EOF.to_usize() { - if mode.to_usize() == C::Flush::FINISH.to_usize() { - self.compressor = None; - } else { - return Err(C::new_error("unexpected eof", vm)); - } - } - - buf.shrink_to_fit(); - Ok(buf) - } -} - #[derive(Debug)] pub(crate) struct DecompressState { decompress: D, @@ -295,11 +144,6 @@ impl DecompressState { self.eof } - #[cfg_attr(target_os = "android", allow(dead_code))] - pub(crate) const fn decompressor(&self) -> &D { - &self.decompress - } - pub(crate) fn unused_data(&self) -> PyBytesRef { self.unused_data.clone() } diff --git a/crates/stdlib/src/lzma.rs b/crates/stdlib/src/lzma.rs index 6e8a913abaa..23f3a4393a7 100644 --- a/crates/stdlib/src/lzma.rs +++ b/crates/stdlib/src/lzma.rs @@ -4,81 +4,75 @@ pub(crate) use _lzma::module_def; #[pymodule] mod _lzma { - use crate::compression::{ - CompressFlushKind, CompressState, CompressStatusKind, Compressor, DecompressArgs, - DecompressError, DecompressState, DecompressStatus, Decompressor, - }; + use crate::compression::DecompressArgs; use alloc::fmt; - use xz::stream::{ - Action, Check, Error, Filters, LzmaOptions, MatchFinder, Mode, Status, Stream, - TELL_ANY_CHECK, TELL_NO_CHECK, - }; - // lzma_check, lzma_mode, lzma_match_finder have platform-dependent signedness - // (i32 on Windows, u32 elsewhere). Define as fixed-type const to avoid mismatch. - use rustpython_common::lock::PyMutex; + use rustpython_common::{compression::lzma as backend, lock::PyMutex}; use rustpython_vm::builtins::{PyBaseExceptionRef, PyBytesRef, PyDict, PyType, PyTypeRef}; - use rustpython_vm::convert::ToPyException; use rustpython_vm::function::ArgBytesLike; use rustpython_vm::types::Constructor; use rustpython_vm::{Py, PyObjectRef, PyPayload, PyResult, VirtualMachine}; - #[pyattr] - use xz_sys::{ - LZMA_FILTER_ARM as FILTER_ARM, LZMA_FILTER_ARMTHUMB as FILTER_ARMTHUMB, - LZMA_FILTER_DELTA as FILTER_DELTA, LZMA_FILTER_IA64 as FILTER_IA64, - LZMA_FILTER_LZMA1 as FILTER_LZMA1, LZMA_FILTER_LZMA2 as FILTER_LZMA2, - LZMA_FILTER_POWERPC as FILTER_POWERPC, LZMA_FILTER_SPARC as FILTER_SPARC, - LZMA_FILTER_X86 as FILTER_X86, - }; - #[pyattr] - use xz_sys::{LZMA_PRESET_DEFAULT as PRESET_DEFAULT, LZMA_PRESET_EXTREME as PRESET_EXTREME}; - - const BUFSIZ: usize = 8192; - // xz_sys enum types have platform-dependent signedness; `as _` normalizes to i32 #[pyattr] - const CHECK_NONE: i32 = xz_sys::LZMA_CHECK_NONE as _; + const CHECK_NONE: i32 = backend::CHECK_NONE; #[pyattr] - const CHECK_CRC32: i32 = xz_sys::LZMA_CHECK_CRC32 as _; + const CHECK_CRC32: i32 = backend::CHECK_CRC32; #[pyattr] - const CHECK_CRC64: i32 = xz_sys::LZMA_CHECK_CRC64 as _; + const CHECK_CRC64: i32 = backend::CHECK_CRC64; #[pyattr] - const CHECK_SHA256: i32 = xz_sys::LZMA_CHECK_SHA256 as _; + const CHECK_SHA256: i32 = backend::CHECK_SHA256; #[pyattr] - const CHECK_ID_MAX: i32 = 15; + const CHECK_ID_MAX: i32 = backend::CHECK_ID_MAX; #[pyattr] - const CHECK_UNKNOWN: i32 = CHECK_ID_MAX + 1; + const CHECK_UNKNOWN: i32 = backend::CHECK_UNKNOWN; #[pyattr] - const MF_HC3: i32 = xz_sys::LZMA_MF_HC3 as _; + const MF_HC3: i32 = backend::MF_HC3; #[pyattr] - const MF_HC4: i32 = xz_sys::LZMA_MF_HC4 as _; + const MF_HC4: i32 = backend::MF_HC4; #[pyattr] - const MF_BT2: i32 = xz_sys::LZMA_MF_BT2 as _; + const MF_BT2: i32 = backend::MF_BT2; #[pyattr] - const MF_BT3: i32 = xz_sys::LZMA_MF_BT3 as _; + const MF_BT3: i32 = backend::MF_BT3; #[pyattr] - const MF_BT4: i32 = xz_sys::LZMA_MF_BT4 as _; + const MF_BT4: i32 = backend::MF_BT4; #[pyattr] - const MODE_FAST: i32 = xz_sys::LZMA_MODE_FAST as _; + const MODE_FAST: i32 = backend::MODE_FAST; #[pyattr] - const MODE_NORMAL: i32 = xz_sys::LZMA_MODE_NORMAL as _; + const MODE_NORMAL: i32 = backend::MODE_NORMAL; - enum Format { - Auto = 0, - Xz = 1, - Alone = 2, - Raw = 3, - } + #[pyattr] + const FORMAT_AUTO: i32 = backend::FORMAT_AUTO; + #[pyattr] + const FORMAT_XZ: i32 = backend::FORMAT_XZ; + #[pyattr] + const FORMAT_ALONE: i32 = backend::FORMAT_ALONE; + #[pyattr] + const FORMAT_RAW: i32 = backend::FORMAT_RAW; #[pyattr] - const FORMAT_AUTO: i32 = Format::Auto as i32; + const FILTER_LZMA1: u64 = backend::FILTER_LZMA1; + #[pyattr] + const FILTER_LZMA2: u64 = backend::FILTER_LZMA2; + #[pyattr] + const FILTER_DELTA: u64 = backend::FILTER_DELTA; + #[pyattr] + const FILTER_X86: u64 = backend::FILTER_X86; + #[pyattr] + const FILTER_POWERPC: u64 = backend::FILTER_POWERPC; #[pyattr] - const FORMAT_XZ: i32 = Format::Xz as i32; + const FILTER_IA64: u64 = backend::FILTER_IA64; #[pyattr] - const FORMAT_ALONE: i32 = Format::Alone as i32; + const FILTER_ARM: u64 = backend::FILTER_ARM; #[pyattr] - const FORMAT_RAW: i32 = Format::Raw as i32; + const FILTER_ARMTHUMB: u64 = backend::FILTER_ARMTHUMB; + #[pyattr] + const FILTER_SPARC: u64 = backend::FILTER_SPARC; + + #[pyattr] + const PRESET_DEFAULT: u32 = backend::PRESET_DEFAULT; + #[pyattr] + const PRESET_EXTREME: u32 = backend::PRESET_EXTREME; #[pyattr(once, name = "LZMAError")] fn error(vm: &VirtualMachine) -> PyTypeRef { @@ -90,116 +84,16 @@ mod _lzma { } fn new_lzma_error(message: impl Into, vm: &VirtualMachine) -> PyBaseExceptionRef { - let msg: String = message.into(); - vm.new_exception_msg(vm.class("lzma", "LZMAError"), msg.into()) - } - - fn catch_lzma_error(err: Error, vm: &VirtualMachine) -> PyBaseExceptionRef { - match err { - Error::UnsupportedCheck => new_lzma_error("Unsupported integrity check", vm), - Error::Mem => vm.new_memory_error(""), - Error::MemLimit => new_lzma_error("Memory usage limit exceeded", vm), - Error::Format => new_lzma_error("Input format not supported by decoder", vm), - Error::Options => new_lzma_error("Invalid or unsupported options", vm), - Error::Data => new_lzma_error("Corrupt input data", vm), - Error::Program => new_lzma_error("Internal error", vm), - Error::NoCheck => new_lzma_error("Corrupt input data", vm), - } - } - - fn int_to_check(check: i32) -> Option { - if check == -1 { - return Some(Check::Crc64); - } - match check { - CHECK_NONE => Some(Check::None), - CHECK_CRC32 => Some(Check::Crc32), - CHECK_CRC64 => Some(Check::Crc64), - CHECK_SHA256 => Some(Check::Sha256), - _ => None, - } - } - - fn u32_to_mode(val: u32) -> Option { - match val as i32 { - MODE_FAST => Some(Mode::Fast), - MODE_NORMAL => Some(Mode::Normal), - _ => None, - } - } - - fn u32_to_mf(val: u32) -> Option { - match val as i32 { - MF_HC3 => Some(MatchFinder::HashChain3), - MF_HC4 => Some(MatchFinder::HashChain4), - MF_BT2 => Some(MatchFinder::BinaryTree2), - MF_BT3 => Some(MatchFinder::BinaryTree3), - MF_BT4 => Some(MatchFinder::BinaryTree4), - _ => None, - } - } - - struct LzmaStream { - stream: Stream, - check: i32, - header_buf: [u8; 8], - header_collected: u8, - track_header: bool, - } - - impl LzmaStream { - fn new(stream: Stream, check: i32, track_header: bool) -> Self { - Self { - stream, - check, - header_buf: [0u8; 8], - header_collected: 0, - track_header, - } - } + let message: String = message.into(); + vm.new_exception_msg(vm.class("lzma", "LZMAError"), message.into()) } - impl Decompressor for LzmaStream { - type Flush = (); - type Status = Status; - type Error = Error; - - fn total_in(&self) -> u64 { - self.stream.total_in() - } - - fn decompress_vec( - &mut self, - input: &[u8], - output: &mut Vec, - (): Self::Flush, - ) -> Result { - if self.track_header && self.header_collected < 8 { - let need = (8 - self.header_collected) as usize; - let n = need.min(input.len()); - self.header_buf[self.header_collected as usize..][..n].copy_from_slice(&input[..n]); - self.header_collected += n as u8; - } - - match self.stream.process_vec(input, output, Action::Run) { - Ok(Status::GetCheck) => { - if self.header_collected >= 8 { - self.check = (self.header_buf[7] & 0x0F) as i32; - } - Ok(Status::Ok) - } - Err(Error::NoCheck) => { - self.check = CHECK_NONE; - Ok(Status::Ok) - } - other => other, - } - } - } - - impl DecompressStatus for Status { - fn is_stream_end(&self) -> bool { - *self == Self::StreamEnd + fn map_backend_error(error: backend::Error, vm: &VirtualMachine) -> PyBaseExceptionRef { + match error { + backend::Error::Memory => vm.new_memory_error(""), + backend::Error::Value(message) => vm.new_value_error(message), + backend::Error::Lzma(message) => new_lzma_error(message, vm), + backend::Error::Eof => vm.new_eof_error("End of stream already reached"), } } @@ -212,7 +106,7 @@ mod _lzma { vm.new_type_error("Filter specifier must be a dict or dict-like object") })?; match dict.get_item_opt(key, vm)? { - Some(obj) => Ok(Some(obj.try_into_value::(vm)?)), + Some(value) => Ok(Some(value.try_into_value::(vm)?)), None => Ok(None), } } @@ -226,275 +120,140 @@ mod _lzma { vm.new_type_error("Filter specifier must be a dict or dict-like object") })?; match dict.get_item_opt(key, vm)? { - Some(obj) => Ok(Some(obj.try_into_value::(vm)?)), + Some(value) => Ok(Some(value.try_into_value::(vm)?)), None => Ok(None), } } - fn parse_filter_spec_lzma(spec: &PyObjectRef, vm: &VirtualMachine) -> PyResult { - let preset = get_dict_opt_u32(spec, "preset", vm)?.unwrap_or(PRESET_DEFAULT); - - let mut opts = LzmaOptions::new_preset(preset) - .map_err(|_| new_lzma_error(format!("Invalid compression preset: {preset}"), vm))?; - - if let Some(v) = get_dict_opt_u32(spec, "dict_size", vm)? { - opts.dict_size(v); - } - if let Some(v) = get_dict_opt_u32(spec, "lc", vm)? { - opts.literal_context_bits(v); - } - if let Some(v) = get_dict_opt_u32(spec, "lp", vm)? { - opts.literal_position_bits(v); - } - if let Some(v) = get_dict_opt_u32(spec, "pb", vm)? { - opts.position_bits(v); - } - if let Some(v) = get_dict_opt_u32(spec, "mode", vm)? { - let mode = u32_to_mode(v) - .ok_or_else(|| vm.new_value_error("Invalid filter specifier for LZMA filter"))?; - opts.mode(mode); - } - if let Some(v) = get_dict_opt_u32(spec, "nice_len", vm)? { - opts.nice_len(v); - } - if let Some(v) = get_dict_opt_u32(spec, "mf", vm)? { - let mf = u32_to_mf(v) - .ok_or_else(|| vm.new_value_error("Invalid filter specifier for LZMA filter"))?; - opts.match_finder(mf); - } - if let Some(v) = get_dict_opt_u32(spec, "depth", vm)? { - opts.depth(v); - } - - Ok(opts) + fn filter_spec_with_id( + spec: &PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult { + let id = get_dict_opt_u64(spec, "id", vm)? + .ok_or_else(|| vm.new_value_error("Filter specifier must have an \"id\" entry"))?; + Ok(backend::FilterSpec { + id, + ..backend::FilterSpec::default() + }) } - fn parse_filter_spec_delta(spec: &PyObjectRef, vm: &VirtualMachine) -> PyResult { - let dist = get_dict_opt_u32(spec, "dist", vm)?.unwrap_or(1); - if dist == 0 || dist > 256 { - return Err(vm.new_value_error("Invalid filter specifier for delta filter")); + fn parse_filter_chain_item( + spec: &PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult { + let mut parsed = filter_spec_with_id(spec, vm)?; + match parsed.id { + FILTER_LZMA1 | FILTER_LZMA2 => { + parsed.preset = get_dict_opt_u32(spec, "preset", vm)?; + parsed.dict_size = get_dict_opt_u32(spec, "dict_size", vm)?; + parsed.lc = get_dict_opt_u32(spec, "lc", vm)?; + parsed.lp = get_dict_opt_u32(spec, "lp", vm)?; + parsed.pb = get_dict_opt_u32(spec, "pb", vm)?; + parsed.mode = get_dict_opt_u32(spec, "mode", vm)?; + parsed.nice_len = get_dict_opt_u32(spec, "nice_len", vm)?; + parsed.mf = get_dict_opt_u32(spec, "mf", vm)?; + parsed.depth = get_dict_opt_u32(spec, "depth", vm)?; + } + FILTER_DELTA => parsed.dist = get_dict_opt_u32(spec, "dist", vm)?, + FILTER_X86 | FILTER_POWERPC | FILTER_IA64 | FILTER_ARM | FILTER_ARMTHUMB + | FILTER_SPARC => { + parsed.start_offset = get_dict_opt_u32(spec, "start_offset", vm)?; + } + _ => {} } - Ok(dist) - } - - fn parse_filter_spec_bcj(spec: &PyObjectRef, vm: &VirtualMachine) -> PyResult { - Ok(get_dict_opt_u32(spec, "start_offset", vm)?.unwrap_or(0)) + Ok(parsed) } - fn add_bcj_filter( - filters: &mut Filters, - filter_id: u64, - start_offset: u32, - ) -> Result<(), Error> { - if start_offset == 0 { - match filter_id { - FILTER_X86 => { - filters.x86(); - } - FILTER_POWERPC => { - filters.powerpc(); - } - FILTER_IA64 => { - filters.ia64(); - } - FILTER_ARM => { - filters.arm(); - } - FILTER_ARMTHUMB => { - filters.arm_thumb(); - } - FILTER_SPARC => { - filters.sparc(); - } - _ => unreachable!(), + fn parse_filter_properties( + spec: &PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult { + let mut parsed = filter_spec_with_id(spec, vm)?; + match parsed.id { + FILTER_LZMA1 => { + parsed.preset = get_dict_opt_u32(spec, "preset", vm)?; + parsed.lc = get_dict_opt_u32(spec, "lc", vm)?; + parsed.lp = get_dict_opt_u32(spec, "lp", vm)?; + parsed.pb = get_dict_opt_u32(spec, "pb", vm)?; + parsed.dict_size = get_dict_opt_u32(spec, "dict_size", vm)?; } - Ok(()) - } else { - let props = start_offset.to_le_bytes(); - match filter_id { - FILTER_X86 => { - filters.x86_properties(&props)?; - } - FILTER_POWERPC => { - filters.powerpc_properties(&props)?; - } - FILTER_IA64 => { - filters.ia64_properties(&props)?; - } - FILTER_ARM => { - filters.arm_properties(&props)?; - } - FILTER_ARMTHUMB => { - filters.arm_thumb_properties(&props)?; - } - FILTER_SPARC => { - filters.sparc_properties(&props)?; - } - _ => unreachable!(), + FILTER_LZMA2 => { + parsed.preset = get_dict_opt_u32(spec, "preset", vm)?; + parsed.dict_size = get_dict_opt_u32(spec, "dict_size", vm)?; + } + FILTER_DELTA => parsed.dist = get_dict_opt_u32(spec, "dist", vm)?, + FILTER_X86 | FILTER_POWERPC | FILTER_IA64 | FILTER_ARM | FILTER_ARMTHUMB + | FILTER_SPARC => { + parsed.start_offset = get_dict_opt_u32(spec, "start_offset", vm)?; } - Ok(()) + _ => {} } + Ok(parsed) } - fn parse_filter_chain_spec( + fn parse_filter_chain( filter_specs: PyObjectRef, vm: &VirtualMachine, - ) -> PyResult { + ) -> PyResult> { const LZMA_FILTERS_MAX: usize = 4; - let filter_specs_len = filter_specs.length(vm)?; - if filter_specs_len > LZMA_FILTERS_MAX { + let length = filter_specs.length(vm)?; + if length > LZMA_FILTERS_MAX { return Err(new_lzma_error( format!("Too many filters - liblzma supports a maximum of {LZMA_FILTERS_MAX}"), vm, )); } - - let filter_specs = filter_specs.try_sequence(vm)?; - let mut filters = Filters::new(); - for i in 0..filter_specs_len { - let spec = filter_specs.get_item(i as isize, vm)?; - let filter_id = get_dict_opt_u64(&spec, "id", vm)? - .ok_or_else(|| vm.new_value_error("Filter specifier must have an \"id\" entry"))?; - - match filter_id { - FILTER_LZMA1 => { - let opts = parse_filter_spec_lzma(&spec, vm)?; - filters.lzma1(&opts); - } - FILTER_LZMA2 => { - let opts = parse_filter_spec_lzma(&spec, vm)?; - filters.lzma2(&opts); - } - FILTER_DELTA => { - let dist = parse_filter_spec_delta(&spec, vm)?; - filters - .delta_properties(&[(dist - 1) as u8]) - .map_err(|e| catch_lzma_error(e, vm))?; - } - FILTER_X86 | FILTER_POWERPC | FILTER_IA64 | FILTER_ARM | FILTER_ARMTHUMB - | FILTER_SPARC => { - let start_offset = parse_filter_spec_bcj(&spec, vm)?; - add_bcj_filter(&mut filters, filter_id, start_offset) - .map_err(|e| catch_lzma_error(e, vm))?; - } - _ => { - return Err(vm.new_value_error(format!("Invalid filter ID: {filter_id}"))); - } - } - } - - Ok(filters) - } - - const DEFAULT_LC: u32 = xz_sys::LZMA_LC_DEFAULT; - const DEFAULT_LP: u32 = xz_sys::LZMA_LP_DEFAULT; - const DEFAULT_PB: u32 = xz_sys::LZMA_PB_DEFAULT; - const DICT_POW2: [u8; 10] = [18, 20, 21, 22, 22, 23, 23, 24, 25, 26]; - - fn preset_dict_size(preset: u32) -> u32 { - let level = (preset & xz_sys::LZMA_PRESET_LEVEL_MASK) as usize; - if level > 9 { - return 0; - } - 1u32 << DICT_POW2[level] - } - - fn lzma2_dict_size_from_prop(prop: u8) -> u32 { - if prop > 40 { - return u32::MAX; - } - if prop == 40 { - return u32::MAX; - } - let prop = prop as u32; - (2 | (prop & 1)) << (prop / 2 + 11) - } - - fn lzma2_prop_from_dict_size(dict_size: u32) -> u8 { - if dict_size == u32::MAX { - return 40; - } - for i in 0u8..40 { - if lzma2_dict_size_from_prop(i) >= dict_size { - return i; - } - } - 40 - } - - fn encode_lzma1_properties(lc: u32, lp: u32, pb: u32, dict_size: u32) -> Vec { - let mut result = vec![0u8; 5]; - result[0] = ((pb * 5 + lp) * 9 + lc) as u8; - result[1..5].copy_from_slice(&dict_size.to_le_bytes()); - result + let sequence = filter_specs.try_sequence(vm)?; + (0..length) + .map(|index| { + let spec = sequence.get_item(index as isize, vm)?; + parse_filter_chain_item(&spec, vm) + }) + .collect() } - fn decode_lzma1_properties(props: &[u8]) -> Option<(u32, u32, u32, u32)> { - if props.len() < 5 { - return None; - } - let mut d = props[0] as u32; - let lc = d % 9; - d /= 9; - let lp = d % 5; - let pb = d / 5; - let dict_size = u32::from_le_bytes([props[1], props[2], props[3], props[4]]); - Some((lc, lp, pb, dict_size)) + fn filters_to_backend( + filters: Option, + vm: &VirtualMachine, + ) -> PyResult>> { + filters + .map(|filters| parse_filter_chain(filters, vm)) + .transpose() } - fn build_filter_spec( - filter_id: u64, - props: &[u8], + fn filter_spec_to_dict( + spec: backend::FilterSpec, vm: &VirtualMachine, ) -> PyResult { let dict = vm.ctx.new_dict(); - dict.set_item("id", vm.new_pyobj(filter_id), vm)?; - - match filter_id { + dict.set_item("id", vm.new_pyobj(spec.id), vm)?; + match spec.id { FILTER_LZMA1 => { - let (lc, lp, pb, dict_size) = decode_lzma1_properties(props) - .ok_or_else(|| new_lzma_error("Invalid or unsupported options", vm))?; - dict.set_item("lc", vm.new_pyobj(lc), vm)?; - dict.set_item("lp", vm.new_pyobj(lp), vm)?; - dict.set_item("pb", vm.new_pyobj(pb), vm)?; - dict.set_item("dict_size", vm.new_pyobj(dict_size), vm)?; + dict.set_item("lc", vm.new_pyobj(spec.lc.unwrap()), vm)?; + dict.set_item("lp", vm.new_pyobj(spec.lp.unwrap()), vm)?; + dict.set_item("pb", vm.new_pyobj(spec.pb.unwrap()), vm)?; + dict.set_item("dict_size", vm.new_pyobj(spec.dict_size.unwrap()), vm)?; } FILTER_LZMA2 => { - if props.len() != 1 { - return Err(new_lzma_error("Invalid or unsupported options", vm)); - } - let dict_size = lzma2_dict_size_from_prop(props[0]); - dict.set_item("dict_size", vm.new_pyobj(dict_size), vm)?; + dict.set_item("dict_size", vm.new_pyobj(spec.dict_size.unwrap()), vm)?; } FILTER_DELTA => { - if props.len() != 1 { - return Err(new_lzma_error("Invalid or unsupported options", vm)); - } - let dist = props[0] as u32 + 1; - dict.set_item("dist", vm.new_pyobj(dist), vm)?; + dict.set_item("dist", vm.new_pyobj(spec.dist.unwrap()), vm)?; } FILTER_X86 | FILTER_POWERPC | FILTER_IA64 | FILTER_ARM | FILTER_ARMTHUMB | FILTER_SPARC => { - if props.is_empty() { - // default: no start_offset - } else if props.len() == 4 { - let start_offset = u32::from_le_bytes([props[0], props[1], props[2], props[3]]); + if let Some(start_offset) = spec.start_offset { dict.set_item("start_offset", vm.new_pyobj(start_offset), vm)?; - } else { - return Err(new_lzma_error("Invalid or unsupported options", vm)); } } - _ => { - return Err(vm.new_value_error(format!("Invalid filter ID: {filter_id}"))); - } + _ => unreachable!("common backend validated the filter ID"), } - Ok(dict.into()) } #[pyfunction] fn is_check_supported(check_id: i32) -> bool { - unsafe { xz_sys::lzma_check_is_supported(check_id as _) != 0 } + backend::is_check_supported(check_id) } #[pyfunction] @@ -502,45 +261,8 @@ mod _lzma { filter_spec: PyObjectRef, vm: &VirtualMachine, ) -> PyResult> { - let filter_id = get_dict_opt_u64(&filter_spec, "id", vm)? - .ok_or_else(|| vm.new_value_error("Filter specifier must have an \"id\" entry"))?; - - match filter_id { - FILTER_LZMA1 => { - let preset = - get_dict_opt_u32(&filter_spec, "preset", vm)?.unwrap_or(PRESET_DEFAULT); - let lc = get_dict_opt_u32(&filter_spec, "lc", vm)?.unwrap_or(DEFAULT_LC); - let lp = get_dict_opt_u32(&filter_spec, "lp", vm)?.unwrap_or(DEFAULT_LP); - let pb = get_dict_opt_u32(&filter_spec, "pb", vm)?.unwrap_or(DEFAULT_PB); - let dict_size = get_dict_opt_u32(&filter_spec, "dict_size", vm)? - .unwrap_or_else(|| preset_dict_size(preset)); - Ok(encode_lzma1_properties(lc, lp, pb, dict_size)) - } - FILTER_LZMA2 => { - let preset = - get_dict_opt_u32(&filter_spec, "preset", vm)?.unwrap_or(PRESET_DEFAULT); - let dict_size = get_dict_opt_u32(&filter_spec, "dict_size", vm)? - .unwrap_or_else(|| preset_dict_size(preset)); - Ok(vec![lzma2_prop_from_dict_size(dict_size)]) - } - FILTER_DELTA => { - let dist = get_dict_opt_u32(&filter_spec, "dist", vm)?.unwrap_or(1); - if dist == 0 || dist > 256 { - return Err(vm.new_value_error("Invalid filter specifier for delta filter")); - } - Ok(vec![(dist - 1) as u8]) - } - FILTER_X86 | FILTER_POWERPC | FILTER_IA64 | FILTER_ARM | FILTER_ARMTHUMB - | FILTER_SPARC => { - let start_offset = get_dict_opt_u32(&filter_spec, "start_offset", vm)?.unwrap_or(0); - if start_offset == 0 { - Ok(vec![]) - } else { - Ok(start_offset.to_le_bytes().to_vec()) - } - } - _ => Err(vm.new_value_error(format!("Invalid filter ID: {filter_id}"))), - } + let spec = parse_filter_properties(&filter_spec, vm)?; + backend::encode_filter_properties(&spec).map_err(|error| map_backend_error(error, vm)) } #[pyfunction] @@ -549,15 +271,29 @@ mod _lzma { encoded_props: ArgBytesLike, vm: &VirtualMachine, ) -> PyResult { - let props = encoded_props.borrow_buf(); - build_filter_spec(filter_id, &props, vm) + let spec = encoded_props + .with_ref(|properties| backend::decode_filter_properties(filter_id, properties)); + filter_spec_to_dict(spec.map_err(|error| map_backend_error(error, vm))?, vm) + } + + struct DecompressorInner { + backend: backend::Decompressor, + unused_data: PyBytesRef, + } + + impl DecompressorInner { + fn sync_visible_state(&mut self, vm: &VirtualMachine) { + if self.unused_data.as_bytes() != self.backend.unused_data() { + self.unused_data = vm.ctx.new_bytes(self.backend.unused_data().to_vec()); + } + } } #[pyattr] #[pyclass(name = "LZMADecompressor")] #[derive(PyPayload)] struct LZMADecompressor { - state: PyMutex>, + state: PyMutex, } impl fmt::Debug for LZMADecompressor { @@ -583,49 +319,20 @@ mod _lzma { if args.format == FORMAT_RAW && args.memlimit.is_some() { return Err(vm.new_value_error("Cannot specify memory limit with FORMAT_RAW")); } - if args.format == FORMAT_RAW && args.filters.is_none() { return Err(vm.new_value_error("Must specify filters for FORMAT_RAW")); } if args.format != FORMAT_RAW && args.filters.is_some() { return Err(vm.new_value_error("Cannot specify filters except with FORMAT_RAW")); } - - let memlimit = args.memlimit.unwrap_or(u64::MAX); - let decoder_flags = TELL_ANY_CHECK | TELL_NO_CHECK; - - let lzma_stream = match args.format { - FORMAT_AUTO => { - let stream = Stream::new_auto_decoder(memlimit, decoder_flags) - .map_err(|e| catch_lzma_error(e, vm))?; - LzmaStream::new(stream, CHECK_UNKNOWN, true) - } - FORMAT_XZ => { - let stream = Stream::new_stream_decoder(memlimit, decoder_flags) - .map_err(|e| catch_lzma_error(e, vm))?; - LzmaStream::new(stream, CHECK_UNKNOWN, true) - } - FORMAT_ALONE => { - let stream = - Stream::new_lzma_decoder(memlimit).map_err(|e| catch_lzma_error(e, vm))?; - LzmaStream::new(stream, CHECK_NONE, false) - } - FORMAT_RAW => { - let filter_specs = args.filters.unwrap(); // safe: checked above - let filters = parse_filter_chain_spec(filter_specs, vm)?; - let stream = - Stream::new_raw_decoder(&filters).map_err(|e| catch_lzma_error(e, vm))?; - LzmaStream::new(stream, CHECK_NONE, false) - } - _ => { - return Err( - vm.new_value_error(format!("Invalid container format: {}", args.format)) - ); - } - }; - + let filters = filters_to_backend(args.filters, vm)?; + let backend = backend::Decompressor::new(args.format, args.memlimit, filters) + .map_err(|error| map_backend_error(error, vm))?; Ok(Self { - state: PyMutex::new(DecompressState::new(lzma_stream, vm)), + state: PyMutex::new(DecompressorInner { + backend, + unused_data: vm.ctx.empty_bytes.clone(), + }), }) } } @@ -636,88 +343,30 @@ mod _lzma { fn decompress(&self, args: DecompressArgs, vm: &VirtualMachine) -> PyResult> { let max_length = args.max_length(); let data = &*args.data(); - let mut state = self.state.lock(); - state - .decompress(data, max_length, BUFSIZ, vm) - .map_err(|e| match e { - DecompressError::Decompress(err) => catch_lzma_error(err, vm), - DecompressError::Eof(err) => err.to_pyexception(vm), - }) + let result = state.backend.decompress(data, max_length); + state.sync_visible_state(vm); + result.map_err(|error| map_backend_error(error, vm)) } #[pygetset] fn check(&self) -> i32 { - self.state.lock().decompressor().check + self.state.lock().backend.check() } #[pygetset] fn eof(&self) -> bool { - self.state.lock().eof() + self.state.lock().backend.eof() } #[pygetset] fn unused_data(&self) -> PyBytesRef { - self.state.lock().unused_data() + self.state.lock().unused_data.clone() } #[pygetset] fn needs_input(&self) -> bool { - self.state.lock().needs_input() - } - } - - struct CompressorInner { - stream: Stream, - } - - impl CompressStatusKind for Status { - const EOF: Self = Self::StreamEnd; - - fn to_usize(self) -> usize { - self as usize - } - } - - impl CompressFlushKind for Action { - const NONE: Self = Self::Run; - const FINISH: Self = Self::Finish; - - fn to_usize(self) -> usize { - self as usize - } - } - - impl Compressor for CompressorInner { - type Status = Status; - type Flush = Action; - const CHUNKSIZE: usize = u32::MAX as usize; - const DEF_BUF_SIZE: usize = 16 * 1024; - - fn compress_vec( - &mut self, - input: &[u8], - output: &mut Vec, - flush: Self::Flush, - vm: &VirtualMachine, - ) -> PyResult { - self.stream - .process_vec(input, output, flush) - .map_err(|e| catch_lzma_error(e, vm)) - } - - fn total_in(&mut self) -> usize { - self.stream.total_in() as usize - } - - fn new_error(message: impl Into, vm: &VirtualMachine) -> PyBaseExceptionRef { - new_lzma_error(message, vm) - } - } - - impl CompressorInner { - fn new(stream: Stream) -> Self { - Self { stream } + self.state.lock().backend.needs_input() } } @@ -725,7 +374,7 @@ mod _lzma { #[pyclass(name = "LZMACompressor")] #[derive(PyPayload)] struct LZMACompressor { - state: PyMutex>, + state: PyMutex, } impl fmt::Debug for LZMACompressor { @@ -734,52 +383,6 @@ mod _lzma { } } - impl LZMACompressor { - fn init_xz( - check: i32, - preset: u32, - filters: Option, - vm: &VirtualMachine, - ) -> PyResult { - let real_check = - int_to_check(check).ok_or_else(|| vm.new_value_error("Invalid check value"))?; - if let Some(filter_specs) = filters { - let filters = parse_filter_chain_spec(filter_specs, vm)?; - Stream::new_stream_encoder(&filters, real_check) - .map_err(|e| catch_lzma_error(e, vm)) - } else { - Stream::new_easy_encoder(preset, real_check).map_err(|e| catch_lzma_error(e, vm)) - } - } - - fn init_alone( - preset: u32, - filter_specs: Option, - vm: &VirtualMachine, - ) -> PyResult { - if let Some(filter_specs) = filter_specs { - filter_specs.length(vm)?; - // TODO: validate single LZMA1 filter and use its options - let options = LzmaOptions::new_preset(preset).map_err(|_| { - new_lzma_error(format!("Invalid compression preset: {preset}"), vm) - })?; - Stream::new_lzma_encoder(&options).map_err(|e| catch_lzma_error(e, vm)) - } else { - let options = LzmaOptions::new_preset(preset).map_err(|_| { - new_lzma_error(format!("Invalid compression preset: {preset}"), vm) - })?; - Stream::new_lzma_encoder(&options).map_err(|e| catch_lzma_error(e, vm)) - } - } - - fn init_raw(filter_specs: Option, vm: &VirtualMachine) -> PyResult { - let filter_specs = filter_specs - .ok_or_else(|| vm.new_value_error("Must specify filters for FORMAT_RAW"))?; - let filters = parse_filter_chain_spec(filter_specs, vm)?; - Stream::new_raw_encoder(&filters).map_err(|e| catch_lzma_error(e, vm)) - } - } - #[derive(FromArgs)] pub(super) struct LZMACompressorConstructorArgs { #[pyarg(any, default = FORMAT_XZ)] @@ -802,32 +405,24 @@ mod _lzma { vm, )); } - if args.preset.is_some() && args.filters.is_some() { return Err(new_lzma_error( "Cannot specify both preset and filter chain", vm, )); } - - let preset: u32 = match &args.preset { - Some(obj) => obj.clone().try_into_value(vm)?, + let preset = match args.preset { + Some(preset) => preset.try_into_value::(vm)?, None => PRESET_DEFAULT, }; - - let stream = match args.format { - FORMAT_XZ => Self::init_xz(args.check, preset, args.filters, vm)?, - FORMAT_ALONE => Self::init_alone(preset, args.filters, vm)?, - FORMAT_RAW => Self::init_raw(args.filters, vm)?, - _ => { - return Err( - vm.new_value_error(format!("Invalid container format: {}", args.format)) - ); - } + let filters = match args.format { + FORMAT_XZ | FORMAT_ALONE | FORMAT_RAW => filters_to_backend(args.filters, vm)?, + _ => None, }; - + let backend = backend::Compressor::new(args.format, args.check, preset, filters) + .map_err(|error| map_backend_error(error, vm))?; Ok(Self { - state: PyMutex::new(CompressState::new(CompressorInner::new(stream))), + state: PyMutex::new(backend), }) } } @@ -836,14 +431,16 @@ mod _lzma { impl LZMACompressor { #[pymethod] fn compress(&self, data: ArgBytesLike, vm: &VirtualMachine) -> PyResult> { - let mut state = self.state.lock(); - state.compress(&data.borrow_buf(), vm) + data.with_ref(|data| self.state.lock().compress(data)) + .map_err(|error| map_backend_error(error, vm)) } #[pymethod] fn flush(&self, vm: &VirtualMachine) -> PyResult> { - let mut state = self.state.lock(); - state.flush(Action::Finish, vm) + self.state + .lock() + .flush() + .map_err(|error| map_backend_error(error, vm)) } } } diff --git a/crates/stdlib/src/zlib.rs b/crates/stdlib/src/zlib.rs index 08855e43e79..2e1fa8d0562 100644 --- a/crates/stdlib/src/zlib.rs +++ b/crates/stdlib/src/zlib.rs @@ -1,31 +1,24 @@ -// spell-checker:ignore compressobj decompressobj zdict chunksize zlibmodule miniz chunker +// spell-checker:ignore compressobj decompressobj zdict chunksize zlibmodule pub(crate) use zlib::module_def; #[pymodule] mod zlib { - use crate::compression::{ - _decompress, CompressFlushKind, CompressState, CompressStatusKind, Compressor, - DecompressArgs, DecompressError, DecompressFlushKind, DecompressState, DecompressStatus, - Decompressor, USE_AFTER_FINISH_ERR, flush_sync, - }; + use crate::compression::DecompressArgs; use crate::vm::{ - Py, PyObject, PyPayload, PyResult, VirtualMachine, + Py, PyObject, PyObjectRef, PyPayload, PyResult, VirtualMachine, builtins::{PyBaseExceptionRef, PyBytesRef, PyIntRef, PyType, PyTypeRef}, common::lock::PyMutex, - convert::{ToPyException, TryFromBorrowedObject}, + convert::TryFromBorrowedObject, function::{ArgBytesLike, ArgPrimitiveIndex, ArgSize, OptionalArg}, types::Constructor, }; use adler32::RollingAdler32 as Adler32; - use flate2::{ - Compress, Compression, Decompress, FlushCompress, FlushDecompress, Status, - write::ZlibEncoder, - }; - use std::io::Write; + use alloc::fmt; + use rustpython_common::compression::zlib as backend; #[pyattr] - use libz_rs_sys::{ + use rustpython_common::compression::zlib::{ Z_BEST_COMPRESSION, Z_BEST_SPEED, Z_BLOCK, Z_DEFAULT_COMPRESSION, Z_DEFAULT_STRATEGY, Z_DEFLATED as DEFLATED, Z_FILTERED, Z_FINISH, Z_FIXED, Z_FULL_FLUSH, Z_HUFFMAN_ONLY, Z_NO_COMPRESSION, Z_NO_FLUSH, Z_PARTIAL_FLUSH, Z_RLE, Z_SYNC_FLUSH, Z_TREES, @@ -34,22 +27,16 @@ mod zlib { #[pyattr(name = "__version__")] const __VERSION__: &str = "1.0"; - // we're statically linking libz-rs, so the compile-time and runtime - // versions will always be the same + // We statically link libz-rs, so the compile-time and runtime versions + // always match. #[pyattr(name = "ZLIB_RUNTIME_VERSION")] #[pyattr] - const ZLIB_VERSION: &str = unsafe { - match core::ffi::CStr::from_ptr(libz_rs_sys::zlibVersion()).to_str() { - Ok(s) => s, - Err(_) => unreachable!(), - } - }; + const ZLIB_VERSION: &str = backend::version(); - // copied from zlibmodule.c (commit 530f506ac91338) #[pyattr] - const MAX_WBITS: i8 = 15; + const MAX_WBITS: i32 = backend::MAX_WBITS; #[pyattr] - const DEF_BUF_SIZE: usize = 16 * 1024; + const DEF_BUF_SIZE: usize = backend::DEF_BUF_SIZE; #[pyattr] const DEF_MEM_LEVEL: u8 = 8; @@ -66,7 +53,6 @@ mod zlib { fn adler32(data: ArgBytesLike, begin_state: OptionalArg) -> u32 { data.with_ref(|data| { let begin_state = begin_state.map_or(1, |i| i.as_u32_mask()); - let mut hasher = Adler32::from_value(begin_state); hasher.update_buffer(data); hasher.hash() @@ -85,68 +71,19 @@ mod zlib { #[pyarg(any, default = Level::new(Z_DEFAULT_COMPRESSION))] level: Level, #[pyarg(any, default = ArgPrimitiveIndex { value: MAX_WBITS })] - wbits: ArgPrimitiveIndex, + wbits: ArgPrimitiveIndex, } - /// Returns a bytes object containing compressed data. #[pyfunction] fn compress(args: PyFuncCompressArgs, vm: &VirtualMachine) -> PyResult { - let PyFuncCompressArgs { - data, - level, - ref wbits, - } = args; - let level = level.ok_or_else(|| new_zlib_error("Bad compression level", vm))?; - - let compress = InitOptions::new(wbits.value, vm)?.compress(level); - let mut encoder = ZlibEncoder::new_with_compress(Vec::new(), compress); - data.with_ref(|input_bytes| encoder.write_all(input_bytes).unwrap()); - let encoded_bytes = encoder.finish().unwrap(); - Ok(vm.ctx.new_bytes(encoded_bytes)) - } - - enum InitOptions { - Standard { - header: bool, - // [De]Compress::new_with_window_bits is only enabled for zlib; miniz_oxide doesn't - // support wbits (yet?) - wbits: u8, - }, - Gzip { - wbits: u8, - }, - } - - impl InitOptions { - fn new(wbits: i8, vm: &VirtualMachine) -> PyResult { - let header = wbits > 0; - let wbits = wbits.unsigned_abs(); - match wbits { - // TODO: wbits = 0 should be a valid option: - // > windowBits can also be zero to request that inflate use the window size in - // > the zlib header of the compressed stream. - // but flate2 doesn't expose it - // 0 => ... - 9..=15 => Ok(Self::Standard { header, wbits }), - 25..=31 => Ok(Self::Gzip { wbits: wbits - 16 }), - _ => Err(vm.new_value_error("Invalid initialization option")), - } - } - - fn decompress(self) -> Decompress { - match self { - Self::Standard { header, wbits } => Decompress::new_with_window_bits(header, wbits), - Self::Gzip { wbits } => Decompress::new_gzip(wbits), - } - } - fn compress(self, level: Compression) -> Compress { - match self { - Self::Standard { header, wbits } => { - Compress::new_with_window_bits(level, header, wbits) - } - Self::Gzip { wbits } => Compress::new_gzip(level, wbits), - } - } + let PyFuncCompressArgs { data, level, wbits } = args; + let level = level + .value() + .ok_or_else(|| new_zlib_error("Bad compression level", vm))?; + let encoded = data.with_ref(|data| backend::compress(data, level, wbits.value)); + encoded + .map(|data| vm.ctx.new_bytes(data)) + .map_err(|err| new_init_or_zlib_error(err, vm)) } #[derive(FromArgs)] @@ -154,12 +91,11 @@ mod zlib { #[pyarg(positional)] data: ArgBytesLike, #[pyarg(any, default = ArgPrimitiveIndex { value: MAX_WBITS })] - wbits: ArgPrimitiveIndex, + wbits: ArgPrimitiveIndex, #[pyarg(any, default = ArgPrimitiveIndex { value: DEF_BUF_SIZE })] bufsize: ArgPrimitiveIndex, } - /// Returns a bytes object containing the uncompressed data. #[pyfunction] fn decompress(args: PyFuncDecompressArgs, vm: &VirtualMachine) -> PyResult> { let PyFuncDecompressArgs { @@ -167,69 +103,88 @@ mod zlib { wbits, bufsize, } = args; - data.with_ref(|data| { - let mut d = InitOptions::new(wbits.value, vm)?.decompress(); - let (buf, stream_end) = _decompress(data, &mut d, bufsize.value, None, flush_sync) - .map_err(|e| new_zlib_error(e.to_string(), vm))?; - if !stream_end { - return Err(new_zlib_error( - "Error -5 while decompressing data: incomplete or truncated stream", - vm, - )); - } - Ok(buf) - }) + data.with_ref(|data| backend::decompress(data, wbits.value, bufsize.value)) + .map_err(|err| new_init_or_zlib_error(err, vm)) } #[derive(FromArgs)] struct DecompressobjArgs { #[pyarg(any, default = ArgPrimitiveIndex { value: MAX_WBITS })] - wbits: ArgPrimitiveIndex, + wbits: ArgPrimitiveIndex, #[pyarg(any, optional)] zdict: OptionalArg, } + fn owned_dict(zdict: OptionalArg) -> Option> { + zdict + .into_option() + .map(|dict| dict.with_ref(|data| data.to_vec())) + } + #[pyfunction] fn decompressobj(args: DecompressobjArgs, vm: &VirtualMachine) -> PyResult { - let mut decompress = InitOptions::new(args.wbits.value, vm)?.decompress(); - let zdict = args.zdict.into_option(); - if let Some(dict) = &zdict - && args.wbits.value < 0 - { - dict.with_ref(|d| decompress.set_dictionary(d)) - .map_err(|_| new_zlib_error("failed to set dictionary", vm))?; - } - let inner = PyDecompressInner { - decompress: Some(DecompressWithDict { decompress, zdict }), - eof: false, - unused_data: vm.ctx.empty_bytes.clone(), - unconsumed_tail: vm.ctx.empty_bytes.clone(), - }; + let decompress = backend::Decompressor::new(args.wbits.value, owned_dict(args.zdict)) + .map_err(|err| new_init_or_zlib_error(err, vm))?; Ok(PyDecompress { - inner: PyMutex::new(inner), + inner: PyMutex::new(PyDecompressInner { + decompress, + unused_data: vm.ctx.empty_bytes.clone(), + unconsumed_tail: vm.ctx.empty_bytes.clone(), + }), }) } - #[derive(Debug)] struct PyDecompressInner { - decompress: Option, - eof: bool, + decompress: backend::Decompressor, unused_data: PyBytesRef, unconsumed_tail: PyBytesRef, } + impl PyDecompressInner { + fn sync_visible_state(&mut self, vm: &VirtualMachine) { + if self.unused_data.as_bytes() != self.decompress.unused_data() { + self.unused_data = vm.ctx.new_bytes(self.decompress.unused_data().to_vec()); + } + if self.unconsumed_tail.as_bytes() != self.decompress.unconsumed_tail() { + self.unconsumed_tail = vm.ctx.new_bytes(self.decompress.unconsumed_tail().to_vec()); + } + } + } + #[pyattr] - #[pyclass(name = "Decompress")] - #[derive(Debug, PyPayload)] + #[pyclass(name = "Decompress", traverse)] + #[derive(PyPayload)] struct PyDecompress { + #[pytraverse(skip)] inner: PyMutex, } + impl fmt::Debug for PyDecompress { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "zlib.Decompress") + } + } + #[pyclass(flags(DISALLOW_INSTANTIATION))] impl PyDecompress { + fn copy_inner(&self, vm: &VirtualMachine) -> PyResult { + let inner = self.inner.lock(); + let decompress = inner + .decompress + .copy() + .map_err(|err| vm.new_value_error(err))?; + Ok(Self { + inner: PyMutex::new(PyDecompressInner { + decompress, + unused_data: inner.unused_data.clone(), + unconsumed_tail: inner.unconsumed_tail.clone(), + }), + }) + } + #[pygetset] fn eof(&self) -> bool { - self.inner.lock().eof + self.inner.lock().decompress.eof() } #[pygetset] @@ -242,55 +197,6 @@ mod zlib { self.inner.lock().unconsumed_tail.clone() } - fn decompress_inner( - inner: &mut PyDecompressInner, - data: &[u8], - bufsize: usize, - max_length: Option, - is_flush: bool, - vm: &VirtualMachine, - ) -> PyResult<(PyResult>, bool)> { - let Some(d) = &mut inner.decompress else { - return Err(new_zlib_error(USE_AFTER_FINISH_ERR, vm)); - }; - - let prev_in = d.total_in(); - let res = if is_flush { - // if is_flush: ignore zdict, finish if final chunk - let calc_flush = |final_chunk| { - if final_chunk { - FlushDecompress::Finish - } else { - FlushDecompress::None - } - }; - _decompress(data, &mut d.decompress, bufsize, max_length, calc_flush) - } else { - _decompress(data, d, bufsize, max_length, flush_sync) - } - .map_err(|e| new_zlib_error(e.to_string(), vm)); - let (ret, stream_end) = match res { - Ok((buf, stream_end)) => (Ok(buf), stream_end), - Err(err) => (Err(err), false), - }; - let consumed = (d.total_in() - prev_in) as usize; - - // save unused input - let unconsumed = &data[consumed..]; - if !unconsumed.is_empty() { - if stream_end { - let unused = [inner.unused_data.as_bytes(), unconsumed].concat(); - inner.unused_data = vm.ctx.new_pyref(unused); - } else { - inner.unconsumed_tail = vm.ctx.new_bytes(unconsumed.to_vec()); - } - } else if !inner.unconsumed_tail.is_empty() { - inner.unconsumed_tail = vm.ctx.empty_bytes.clone(); - } - - Ok((ret, stream_end)) - } - #[pymethod] fn decompress(&self, args: DecompressArgs, vm: &VirtualMachine) -> PyResult> { let max_length: usize = args @@ -301,14 +207,10 @@ mod zlib { let max_length = (max_length != 0).then_some(max_length); let data = &*args.data(); - let inner = &mut *self.inner.lock(); - - let (ret, stream_end) = - Self::decompress_inner(inner, data, DEF_BUF_SIZE, max_length, false, vm)?; - - inner.eof |= stream_end; - - ret + let mut inner = self.inner.lock(); + let result = inner.decompress.decompress(data, max_length); + inner.sync_visible_state(vm); + result.map_err(|err| new_zlib_error(err, vm)) } #[pymethod] @@ -321,268 +223,202 @@ mod zlib { OptionalArg::Missing => DEF_BUF_SIZE, }; - let inner = &mut *self.inner.lock(); - let data = core::mem::replace(&mut inner.unconsumed_tail, vm.ctx.empty_bytes.clone()); + let mut inner = self.inner.lock(); + let result = inner.decompress.flush(length); + inner.sync_visible_state(vm); + result.map_err(|err| new_zlib_error(err, vm)) + } - let (ret, _) = Self::decompress_inner(inner, &data, length, None, true, vm)?; + #[pymethod] + fn copy(&self, vm: &VirtualMachine) -> PyResult { + self.copy_inner(vm) + } - if inner.eof { - inner.decompress = None; - } + #[pymethod(name = "__copy__")] + fn copy_dunder(&self, vm: &VirtualMachine) -> PyResult { + self.copy_inner(vm) + } - ret + #[pymethod(name = "__deepcopy__")] + fn deepcopy(&self, _memo: PyObjectRef, vm: &VirtualMachine) -> PyResult { + self.copy_inner(vm) } } #[derive(FromArgs)] - #[allow(dead_code)] // FIXME: use args struct CompressobjArgs { #[pyarg(any, default = Level::new(Z_DEFAULT_COMPRESSION))] level: Level, - // only DEFLATED is valid right now, it's w/e #[pyarg(any, default = DEFLATED)] method: i32, #[pyarg(any, default = ArgPrimitiveIndex { value: MAX_WBITS })] - wbits: ArgPrimitiveIndex, + wbits: ArgPrimitiveIndex, #[pyarg(any, name = "memLevel", default = DEF_MEM_LEVEL)] mem_level: u8, #[pyarg(any, default = Z_DEFAULT_STRATEGY)] strategy: i32, #[pyarg(any, optional)] - zdict: Option, + zdict: OptionalArg, } #[pyfunction] fn compressobj(args: CompressobjArgs, vm: &VirtualMachine) -> PyResult { let CompressobjArgs { level, + method, wbits, + mem_level, + strategy, zdict, - .. } = args; - let level = level.ok_or_else(|| vm.new_value_error("invalid initialization option"))?; - #[allow(unused_mut)] - let mut compress = InitOptions::new(wbits.value, vm)?.compress(level); - if let Some(zdict) = zdict { - zdict.with_ref(|zdict| compress.set_dictionary(zdict).unwrap()); - } + let level = level + .value() + .ok_or_else(|| vm.new_value_error("Invalid initialization option"))?; + let zdict = owned_dict(zdict); + let compress = backend::Compressor::new( + level, + method, + wbits.value, + mem_level.into(), + strategy, + zdict.as_deref(), + ) + .map_err(|err| new_init_or_zlib_error(err, vm))?; Ok(PyCompress { - inner: PyMutex::new(CompressState::new(CompressInner::new(compress))), + inner: PyMutex::new(compress), }) } - #[derive(Debug)] - struct CompressInner { - compress: Compress, - } - #[pyattr] - #[pyclass(name = "Compress")] - #[derive(Debug, PyPayload)] + #[pyclass(name = "Compress", traverse)] + #[derive(PyPayload)] struct PyCompress { - inner: PyMutex>, + #[pytraverse(skip)] + inner: PyMutex, + } + + impl fmt::Debug for PyCompress { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "zlib.Compress") + } } #[pyclass(flags(DISALLOW_INSTANTIATION))] impl PyCompress { + fn copy_inner(&self, vm: &VirtualMachine) -> PyResult { + let compress = self + .inner + .lock() + .copy() + .map_err(|err| vm.new_value_error(err))?; + Ok(Self { + inner: PyMutex::new(compress), + }) + } + #[pymethod] fn compress(&self, data: ArgBytesLike, vm: &VirtualMachine) -> PyResult> { - let mut inner = self.inner.lock(); - data.with_ref(|b| inner.compress(b, vm)) + data.with_ref(|data| self.inner.lock().compress(data)) + .map_err(|err| new_zlib_error(err, vm)) } #[pymethod] fn flush(&self, mode: OptionalArg, vm: &VirtualMachine) -> PyResult> { - let mode = match mode.unwrap_or(Z_FINISH) { - Z_NO_FLUSH => return Ok(vec![]), - Z_PARTIAL_FLUSH => FlushCompress::Partial, - Z_SYNC_FLUSH => FlushCompress::Sync, - Z_FULL_FLUSH => FlushCompress::Full, - Z_FINISH => FlushCompress::Finish, - _ => return Err(new_zlib_error("invalid mode", vm)), - }; - self.inner.lock().flush(mode, vm) + self.inner + .lock() + .flush(mode.unwrap_or(Z_FINISH)) + .map_err(|err| new_zlib_error(err, vm)) } - // TODO: This is an optional feature of Compress - // #[pymethod] - // #[pymethod(name = "__copy__")] - // #[pymethod(name = "__deepcopy__")] - // fn copy(&self) -> Self { - // todo!("") - // } - } - - const CHUNKSIZE: usize = u32::MAX as usize; - - impl CompressInner { - const fn new(compress: Compress) -> Self { - Self { compress } + #[pymethod] + fn copy(&self, vm: &VirtualMachine) -> PyResult { + self.copy_inner(vm) } - } - - impl CompressStatusKind for Status { - const EOF: Self = Self::StreamEnd; - fn to_usize(self) -> usize { - self as usize + #[pymethod(name = "__copy__")] + fn copy_dunder(&self, vm: &VirtualMachine) -> PyResult { + self.copy_inner(vm) } - } - - impl CompressFlushKind for FlushCompress { - const NONE: Self = Self::None; - const FINISH: Self = Self::Finish; - fn to_usize(self) -> usize { - self as usize + #[pymethod(name = "__deepcopy__")] + fn deepcopy(&self, _memo: PyObjectRef, vm: &VirtualMachine) -> PyResult { + self.copy_inner(vm) } } - impl Compressor for CompressInner { - type Status = Status; - type Flush = FlushCompress; - const CHUNKSIZE: usize = CHUNKSIZE; - const DEF_BUF_SIZE: usize = DEF_BUF_SIZE; - - fn compress_vec( - &mut self, - input: &[u8], - output: &mut Vec, - flush: Self::Flush, - vm: &VirtualMachine, - ) -> PyResult { - self.compress - .compress_vec(input, output, flush) - .map_err(|_| new_zlib_error("error while compressing", vm)) - } - - fn total_in(&mut self) -> usize { - self.compress.total_in() as usize - } - - fn new_error(message: impl Into, vm: &VirtualMachine) -> PyBaseExceptionRef { - new_zlib_error(message, vm) - } + fn new_zlib_error(message: impl Into, vm: &VirtualMachine) -> PyBaseExceptionRef { + vm.new_exception_msg(vm.class("zlib", "error"), message.into().into()) } - fn new_zlib_error(message: impl Into, vm: &VirtualMachine) -> PyBaseExceptionRef { - let msg: String = message.into(); - vm.new_exception_msg(vm.class("zlib", "error"), msg.into()) + fn new_init_or_zlib_error( + error: backend::InitError, + vm: &VirtualMachine, + ) -> PyBaseExceptionRef { + match error { + backend::InitError::InvalidOption => { + vm.new_value_error("Invalid initialization option") + } + backend::InitError::Zlib(message) => new_zlib_error(message, vm), + } } - struct Level(Option); + struct Level(Option); impl Level { - fn new(level: i32) -> Self { - let compression = match level { - Z_DEFAULT_COMPRESSION => Compression::default(), - valid_level @ Z_NO_COMPRESSION..=Z_BEST_COMPRESSION => { - Compression::new(valid_level as u32) - } - _ => return Self(None), - }; - Self(Some(compression)) + const fn new(level: i32) -> Self { + if matches!( + level, + Z_DEFAULT_COMPRESSION | Z_NO_COMPRESSION..=Z_BEST_COMPRESSION + ) { + Self(Some(level)) + } else { + Self(None) + } } - fn ok_or_else( - self, - f: impl FnOnce() -> PyBaseExceptionRef, - ) -> PyResult { - self.0.ok_or_else(f) + const fn value(self) -> Option { + self.0 } } impl<'a> TryFromBorrowedObject<'a> for Level { fn try_from_borrowed_object(vm: &VirtualMachine, obj: &'a PyObject) -> PyResult { - let int: i32 = obj.try_index(vm)?.try_to_primitive(vm)?; - Ok(Self::new(int)) - } - } - - #[pyattr] - #[pyclass(name = "_ZlibDecompressor")] - #[derive(Debug, PyPayload)] - struct ZlibDecompressor { - inner: PyMutex>, - } - - #[derive(Debug)] - struct DecompressWithDict { - decompress: Decompress, - zdict: Option, - } - - impl DecompressStatus for Status { - fn is_stream_end(&self) -> bool { - *self == Self::StreamEnd + let level: i32 = obj.try_index(vm)?.try_to_primitive(vm)?; + Ok(Self::new(level)) } } - impl DecompressFlushKind for FlushDecompress { - const SYNC: Self = Self::Sync; + struct PyZlibDecompressorInner { + decompress: backend::ZlibDecompressor, + unused_data: PyBytesRef, } - impl Decompressor for Decompress { - type Flush = FlushDecompress; - type Status = Status; - type Error = flate2::DecompressError; - - fn total_in(&self) -> u64 { - self.total_in() - } - - fn decompress_vec( - &mut self, - input: &[u8], - output: &mut Vec, - flush: Self::Flush, - ) -> Result { - self.decompress_vec(input, output, flush) - } + #[pyattr] + #[pyclass(name = "_ZlibDecompressor", traverse)] + #[derive(PyPayload)] + struct ZlibDecompressor { + #[pytraverse(skip)] + inner: PyMutex, } - impl Decompressor for DecompressWithDict { - type Flush = FlushDecompress; - type Status = Status; - type Error = flate2::DecompressError; - - fn total_in(&self) -> u64 { - self.decompress.total_in() - } - - fn decompress_vec( - &mut self, - input: &[u8], - output: &mut Vec, - flush: Self::Flush, - ) -> Result { - self.decompress.decompress_vec(input, output, flush) - } - - fn maybe_set_dict(&mut self, err: Self::Error) -> Result<(), Self::Error> { - let zdict = err.needs_dictionary().and(self.zdict.as_ref()).ok_or(err)?; - self.decompress.set_dictionary(&zdict.borrow_buf())?; - Ok(()) + impl fmt::Debug for ZlibDecompressor { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "zlib._ZlibDecompressor") } } - // impl Deconstruct - impl Constructor for ZlibDecompressor { type Args = DecompressobjArgs; fn py_new(_cls: &Py, args: Self::Args, vm: &VirtualMachine) -> PyResult { - let mut decompress = InitOptions::new(args.wbits.value, vm)?.decompress(); - let zdict = args.zdict.into_option(); - if let Some(dict) = &zdict - && args.wbits.value < 0 - { - dict.with_ref(|d| decompress.set_dictionary(d)) - .map_err(|_| new_zlib_error("failed to set dictionary", vm))?; - } - let inner = DecompressState::new(DecompressWithDict { decompress, zdict }, vm); + let decompress = + backend::ZlibDecompressor::new(args.wbits.value, owned_dict(args.zdict)) + .map_err(|err| new_init_or_zlib_error(err, vm))?; Ok(Self { - inner: PyMutex::new(inner), + inner: PyMutex::new(PyZlibDecompressorInner { + decompress, + unused_data: vm.ctx.empty_bytes.clone(), + }), }) } } @@ -591,17 +427,17 @@ mod zlib { impl ZlibDecompressor { #[pygetset] fn eof(&self) -> bool { - self.inner.lock().eof() + self.inner.lock().decompress.eof() } #[pygetset] fn unused_data(&self) -> PyBytesRef { - self.inner.lock().unused_data() + self.inner.lock().unused_data.clone() } #[pygetset] fn needs_input(&self) -> bool { - self.inner.lock().needs_input() + self.inner.lock().decompress.needs_input() } #[pymethod] @@ -609,20 +445,15 @@ mod zlib { let max_length = args.max_length(); let data = &*args.data(); - let inner = &mut *self.inner.lock(); - - inner - .decompress(data, max_length, DEF_BUF_SIZE, vm) - .map_err(|e| match e { - DecompressError::Decompress(err) => new_zlib_error(err.to_string(), vm), - DecompressError::Eof(err) => err.to_pyexception(vm), - }) + let mut inner = self.inner.lock(); + let result = inner.decompress.decompress(data, max_length); + if inner.unused_data.as_bytes() != inner.decompress.unused_data() { + inner.unused_data = vm.ctx.new_bytes(inner.decompress.unused_data().to_vec()); + } + result.map_err(|err| match err { + backend::DecompressError::Zlib(err) => new_zlib_error(err, vm), + backend::DecompressError::Eof => vm.new_eof_error("End of stream already reached"), + }) } - - // TODO: Wait for getstate pyslot to be fixed - // #[pyslot] - // fn getstate(zelf: &PyObject, vm: &VirtualMachine) -> PyResult { - // Err(vm.new_type_error("cannot serialize '_ZlibDecompressor' object".to_owned())) - // } } } diff --git a/crates/wasm/Cargo.toml b/crates/wasm/Cargo.toml index 5a5e8d77fd6..aa7a95984c9 100644 --- a/crates/wasm/Cargo.toml +++ b/crates/wasm/Cargo.toml @@ -49,4 +49,4 @@ wasm-opt = false#["-O1"] workspace = true [package.metadata.cargo-shear] -ignored = ["serde", "rustpython-common"] +ignored = ["serde"]