From 9e5c27f71ce30ad22558875c1fda75d82054017a Mon Sep 17 00:00:00 2001 From: James David Clarke Date: Sat, 23 May 2026 18:10:33 +0100 Subject: [PATCH 01/20] Add _zstd stdlib module backing compression.zstd - Implement ZstdDict, ZstdCompressor, and ZstdDecompressor on top of the zstd-safe crate - Wire the module into stdlib_module_defs (gated off Android/wasm32) - Mark test_zstd_multithread_compress as expected failure when libzstd lacks multi-threading support --- Cargo.toml | 2 + Lib/test/test_zstd.py | 7 +- crates/stdlib/Cargo.toml | 2 + crates/stdlib/src/lib.rs | 5 + crates/stdlib/src/zstd.rs | 1712 +++++++++++++++++++++++++++++++++++++ 5 files changed, 1726 insertions(+), 2 deletions(-) create mode 100644 crates/stdlib/src/zstd.rs diff --git a/Cargo.toml b/Cargo.toml index 3b489a88687..067d4960beb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -323,6 +323,8 @@ x509-cert = "0.3.0" x509-parser = "0.18" xml = "1.3" writeable = "0.6" +zstd = { version = "0.13", default-features = false, features = ["zdict_builder"] } +zstd-safe = { version = "7", default-features = false, features = ["zdict_builder"] } # Lints diff --git a/Lib/test/test_zstd.py b/Lib/test/test_zstd.py index cf618534add..1b4b819d2ef 100644 --- a/Lib/test/test_zstd.py +++ b/Lib/test/test_zstd.py @@ -321,8 +321,11 @@ def test_unknown_compression_parameter(self): with self.assertRaisesRegex(ValueError, pattern): ZstdCompressor(options=option) - @unittest.skipIf(not SUPPORT_MULTITHREADING, - "zstd build doesn't support multi-threaded compression") + # TODO: RUSTPYTHON + @unittest.expectedFailureIf( + not SUPPORT_MULTITHREADING, + "TODO: RUSTPYTHON; libzstd built without multi-threaded compression support", + ) def test_zstd_multithread_compress(self): size = 40*_1M b = THIS_FILE_BYTES * (size // len(THIS_FILE_BYTES)) diff --git a/crates/stdlib/Cargo.toml b/crates/stdlib/Cargo.toml index 98a6677de16..9f8dc95ef4c 100644 --- a/crates/stdlib/Cargo.toml +++ b/crates/stdlib/Cargo.toml @@ -121,6 +121,8 @@ pkcs8 = { workspace = true, features = ["encryption", "pkcs5", "pem"], optional libsqlite3-sys = { workspace = true, features = ["bundled"], optional = true } xz = { workspace = true } xz-sys = { workspace = true } +zstd = { workspace = true } +zstd-safe = { workspace = true } [target.'cfg(windows)'.dependencies] paste = { workspace = true } diff --git a/crates/stdlib/src/lib.rs b/crates/stdlib/src/lib.rs index c41a7ea9faf..a987a1e29aa 100644 --- a/crates/stdlib/src/lib.rs +++ b/crates/stdlib/src/lib.rs @@ -27,6 +27,9 @@ mod lzma; mod zlib; +#[cfg(not(any(target_os = "android", target_arch = "wasm32")))] +mod zstd; + mod blake2; mod hashlib; mod md5; @@ -279,5 +282,7 @@ pub fn stdlib_module_defs(ctx: &Context) -> Vec<&'static builtins::PyModuleDef> ))] uuid::module_def(ctx), zlib::module_def(ctx), + #[cfg(not(any(target_os = "android", target_arch = "wasm32")))] + zstd::module_def(ctx), ] } diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs new file mode 100644 index 00000000000..6c9ad33396e --- /dev/null +++ b/crates/stdlib/src/zstd.rs @@ -0,0 +1,1712 @@ +// spell-checker:ignore cctx dctx CCTX DCTX ldm cdict ddict windowlog hashlog chainlog searchlog +// spell-checker:ignore minmatch dictid checksumflag dstream cstream pyobj zstandard btopt btultra +// spell-checker:ignore btlazy dfast nbworkers windowlogmax windowlog overlap targetcblock +// spell-checker:ignore srcsize zdict refprefix refcdict refddict pledgedsrcsize getframecontentsize +// spell-checker:ignore Zstd Zstandard pylib RFC + +//! The `_zstd` extension module. Backs the pure-Python `compression.zstd` +//! package by exposing the same classes, functions and constants that +//! CPython's `Modules/_zstd/` exposes. The Python wrapper at +//! `Lib/compression/zstd/__init__.py` imports from this module unconditionally, +//! so the names and call signatures here must stay in sync with CPython. +//! +//! Backend: the `zstd_safe` Rust crate (a thin safe wrapper over Facebook's +//! libzstd C library, which is what CPython links against). A handful of +//! routines that `zstd_safe` does not expose at a safe level (parameter +//! bounds, dictionary finalization) drop down to raw `zstd_sys` FFI calls. + +pub(crate) use _zstd::module_def; + +// The compression/decompression parameter and strategy constants below use +// CPython's `ZSTD_c_camelCase` / `ZSTD_d_camelCase` naming convention so the +// pure-Python `compression.zstd` package, which references them by those exact +// names, keeps working unchanged. +#[allow(non_upper_case_globals)] +#[pymodule] +mod _zstd { + use rustpython_common::lock::PyMutex; + use rustpython_vm::builtins::{ + PyBaseExceptionRef, PyBytesRef, PyDict, PyTupleRef, PyType, PyTypeRef, + }; + use rustpython_vm::function::{ArgBytesLike, OptionalArg}; + use rustpython_vm::types::{AsMapping, Constructor, Representable}; + use rustpython_vm::{ + AsObject, Context, Py, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, + }; + use std::ffi::c_int; + use zstd_safe::zstd_sys; + use zstd_safe::{CCtx, CParameter, DCtx, DParameter, InBuffer, OutBuffer}; + + /// Stores the `CompressionParameter` and `DecompressionParameter` enum + /// types that the pure-Python `compression.zstd` wrapper passes through + /// `set_parameter_types`. They let us produce CPython-compatible + /// `TypeError`s when callers pass the wrong enum kind to a (de)compressor + /// (a `CompressionParameter` key to a decompressor, or vice versa). Stored + /// in a `PyMutex` so the module init can set them once and constructors + /// can read them later from any thread. + static REGISTERED_PARAM_TYPES: PyMutex = + PyMutex::new(RegisteredParamTypes { + compression: None, + decompression: None, + }); + + struct RegisteredParamTypes { + compression: Option, + decompression: Option, + } + + // ========================================================================= + // Module-level constants + // ========================================================================= + + /// Default compression level used when none is requested. Matches + /// libzstd's `ZSTD_CLEVEL_DEFAULT`. + #[pyattr] + const ZSTD_CLEVEL_DEFAULT: i32 = zstd_sys::ZSTD_CLEVEL_DEFAULT as i32; + + // Compression parameter identifiers. Values match the `ZSTD_cParameter` + // enum in libzstd, which is what the public `CompressionParameter` IntEnum + // in `Lib/compression/zstd/__init__.py` derives its members from. + #[pyattr] + const ZSTD_c_compressionLevel: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_compressionLevel as i32; + #[pyattr] + const ZSTD_c_windowLog: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_windowLog as i32; + #[pyattr] + const ZSTD_c_hashLog: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_hashLog as i32; + #[pyattr] + const ZSTD_c_chainLog: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_chainLog as i32; + #[pyattr] + const ZSTD_c_searchLog: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_searchLog as i32; + #[pyattr] + const ZSTD_c_minMatch: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_minMatch as i32; + #[pyattr] + const ZSTD_c_targetLength: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_targetLength as i32; + #[pyattr] + const ZSTD_c_strategy: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_strategy as i32; + #[pyattr] + const ZSTD_c_enableLongDistanceMatching: i32 = + zstd_sys::ZSTD_cParameter::ZSTD_c_enableLongDistanceMatching as i32; + #[pyattr] + const ZSTD_c_ldmHashLog: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_ldmHashLog as i32; + #[pyattr] + const ZSTD_c_ldmMinMatch: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_ldmMinMatch as i32; + #[pyattr] + const ZSTD_c_ldmBucketSizeLog: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_ldmBucketSizeLog as i32; + #[pyattr] + const ZSTD_c_ldmHashRateLog: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_ldmHashRateLog as i32; + #[pyattr] + const ZSTD_c_contentSizeFlag: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_contentSizeFlag as i32; + #[pyattr] + const ZSTD_c_checksumFlag: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_checksumFlag as i32; + #[pyattr] + const ZSTD_c_dictIDFlag: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_dictIDFlag as i32; + #[pyattr] + const ZSTD_c_nbWorkers: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_nbWorkers as i32; + #[pyattr] + const ZSTD_c_jobSize: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_jobSize as i32; + #[pyattr] + const ZSTD_c_overlapLog: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_overlapLog as i32; + + // Decompression parameter identifiers. libzstd only exposes one non- + // experimental decompression parameter. + #[pyattr] + const ZSTD_d_windowLogMax: i32 = zstd_sys::ZSTD_dParameter::ZSTD_d_windowLogMax as i32; + + // Strategy enum members ordered from fastest to strongest. These power + // the `Strategy` IntEnum in `Lib/compression/zstd/__init__.py`. + #[pyattr] + const ZSTD_fast: i32 = zstd_sys::ZSTD_strategy::ZSTD_fast as i32; + #[pyattr] + const ZSTD_dfast: i32 = zstd_sys::ZSTD_strategy::ZSTD_dfast as i32; + #[pyattr] + const ZSTD_greedy: i32 = zstd_sys::ZSTD_strategy::ZSTD_greedy as i32; + #[pyattr] + const ZSTD_lazy: i32 = zstd_sys::ZSTD_strategy::ZSTD_lazy as i32; + #[pyattr] + const ZSTD_lazy2: i32 = zstd_sys::ZSTD_strategy::ZSTD_lazy2 as i32; + #[pyattr] + const ZSTD_btlazy2: i32 = zstd_sys::ZSTD_strategy::ZSTD_btlazy2 as i32; + #[pyattr] + const ZSTD_btopt: i32 = zstd_sys::ZSTD_strategy::ZSTD_btopt as i32; + #[pyattr] + const ZSTD_btultra: i32 = zstd_sys::ZSTD_strategy::ZSTD_btultra as i32; + #[pyattr] + const ZSTD_btultra2: i32 = zstd_sys::ZSTD_strategy::ZSTD_btultra2 as i32; + + /// Runtime version string of the linked libzstd, like `"1.5.7"`. + #[pyattr(once, name = "zstd_version")] + fn zstd_version(_vm: &VirtualMachine) -> String { + zstd_safe::version_string().to_string() + } + + /// Packed integer version: `MAJOR * 10000 + MINOR * 100 + RELEASE`. + /// The Python wrapper unpacks this into a `zstd_version_info` tuple. + #[pyattr(once, name = "zstd_version_number")] + fn zstd_version_number(_vm: &VirtualMachine) -> u32 { + zstd_safe::version_number() + } + + /// Recommended output buffer size for decompression. The `ZstdFile.read1` + /// wrapper uses this as its default read size. + #[pyattr(once, name = "ZSTD_DStreamOutSize")] + fn zstd_dstream_out_size(_vm: &VirtualMachine) -> usize { + DCtx::out_size() + } + + // Dictionary load type markers. The `ZstdDict.as_*` properties wrap the + // dictionary in a `(zdict, marker)` tuple so the compressor or decompressor + // constructor knows which load mode to apply. Numbering matches CPython's + // `Modules/_zstd/_zstdmodule.h::DictType`. + const DICT_TYPE_DIGESTED: i32 = 0; + const DICT_TYPE_UNDIGESTED: i32 = 1; + const DICT_TYPE_PREFIX: i32 = 2; + + // ========================================================================= + // ZstdError exception + // ========================================================================= + + /// `_zstd.ZstdError`: the exception type raised for libzstd errors and for + /// invalid usage (calling `decompress` on a closed decompressor, etc.). + #[pyattr(once, name = "ZstdError")] + fn zstd_error(vm: &VirtualMachine) -> PyTypeRef { + vm.ctx.new_exception_type( + "_zstd", + "ZstdError", + Some(vec![vm.ctx.exceptions.exception_type.to_owned()]), + ) + } + + fn new_zstd_error(message: impl Into, vm: &VirtualMachine) -> PyBaseExceptionRef { + let msg: String = message.into(); + vm.new_exception_msg(vm.class("_zstd", "ZstdError"), msg.into()) + } + + /// Convert a libzstd error code (the `usize` returned by most ZSTD_* + /// functions when `ZSTD_isError(code)` is non-zero) into a `ZstdError` + /// carrying the human-readable message from `ZSTD_getErrorName`. + fn catch_zstd_error(code: usize, vm: &VirtualMachine) -> PyBaseExceptionRef { + new_zstd_error(zstd_safe::get_error_name(code).to_string(), vm) + } + + // ========================================================================= + // Parameter helpers + // ========================================================================= + + /// Convert any int-like Python value to an `i32`, rejecting floats. CPython + /// uses `PyLong_AsInt` here which only accepts integers; `try_index` is + /// the equivalent of CPython's `__index__` slot lookup (rejects floats). + fn pyobj_to_i32(obj: &PyObjectRef, vm: &VirtualMachine) -> PyResult { + obj.try_index(vm)?.try_to_primitive(vm) + } + + /// Normalize an `OptionalArg` constructor parameter to an + /// `Option`. Treats both "argument missing" and "argument + /// explicitly None" the same way (consistent with CPython, which uses + /// `None` defaults throughout the `_zstd` API). + fn arg_or_none(arg: OptionalArg, vm: &VirtualMachine) -> Option { + arg.into_option().filter(|o| !vm.is_none(o)) + } + + /// If `key` is a registered `CompressionParameter` or `DecompressionParameter` + /// enum member and its kind does not match the caller's context, return a + /// `TypeError`. Used by both (de)compressor option-dict loops so passing a + /// `CompressionParameter` to a `ZstdDecompressor` (or vice versa) is + /// rejected with the CPython-compatible "must not be a XParameter" + /// message. If `set_parameter_types` has not been called yet, this is a + /// no-op. + fn check_wrong_param_kind( + key: &PyObjectRef, + is_compress: bool, + vm: &VirtualMachine, + ) -> PyResult<()> { + let guard = REGISTERED_PARAM_TYPES.lock(); + let (forbidden_type, forbidden_name, kind) = if is_compress { + (&guard.decompression, "DecompressionParameter", "compression") + } else { + (&guard.compression, "CompressionParameter", "decompression") + }; + if let Some(ref t) = *forbidden_type + && key.fast_isinstance(t) + { + return Err(vm.new_type_error(format!( + "{kind} options dictionary key must not be a {forbidden_name}" + ))); + } + Ok(()) + } + + /// Map a compression parameter id to its raw libzstd `ZSTD_cParameter` + /// enum value. Returns `None` for unknown ids so callers can surface a + /// targeted `ValueError`. Done with an explicit match rather than + /// `mem::transmute` so passing junk like `ZSTD_cParameter(42)` cannot + /// be triggered from Python. + fn c_param_enum(param: i32) -> Option { + use zstd_sys::ZSTD_cParameter as P; + Some(match param { + ZSTD_c_compressionLevel => P::ZSTD_c_compressionLevel, + ZSTD_c_windowLog => P::ZSTD_c_windowLog, + ZSTD_c_hashLog => P::ZSTD_c_hashLog, + ZSTD_c_chainLog => P::ZSTD_c_chainLog, + ZSTD_c_searchLog => P::ZSTD_c_searchLog, + ZSTD_c_minMatch => P::ZSTD_c_minMatch, + ZSTD_c_targetLength => P::ZSTD_c_targetLength, + ZSTD_c_strategy => P::ZSTD_c_strategy, + ZSTD_c_enableLongDistanceMatching => P::ZSTD_c_enableLongDistanceMatching, + ZSTD_c_ldmHashLog => P::ZSTD_c_ldmHashLog, + ZSTD_c_ldmMinMatch => P::ZSTD_c_ldmMinMatch, + ZSTD_c_ldmBucketSizeLog => P::ZSTD_c_ldmBucketSizeLog, + ZSTD_c_ldmHashRateLog => P::ZSTD_c_ldmHashRateLog, + ZSTD_c_contentSizeFlag => P::ZSTD_c_contentSizeFlag, + ZSTD_c_checksumFlag => P::ZSTD_c_checksumFlag, + ZSTD_c_dictIDFlag => P::ZSTD_c_dictIDFlag, + ZSTD_c_nbWorkers => P::ZSTD_c_nbWorkers, + ZSTD_c_jobSize => P::ZSTD_c_jobSize, + ZSTD_c_overlapLog => P::ZSTD_c_overlapLog, + _ => return None, + }) + } + + /// Map a decompression parameter id to its raw libzstd `ZSTD_dParameter` + /// enum value. See [`c_param_enum`] for rationale. + fn d_param_enum(param: i32) -> Option { + use zstd_sys::ZSTD_dParameter as P; + match param { + ZSTD_d_windowLogMax => Some(P::ZSTD_d_windowLogMax), + _ => None, + } + } + + /// Map a compression-parameter id and integer value to a `CParameter` + /// variant. Used by the compressor's `options=` constructor argument. + fn cparameter_from_int(param: i32, value: i32, vm: &VirtualMachine) -> PyResult { + let p = match param { + ZSTD_c_compressionLevel => CParameter::CompressionLevel(value), + ZSTD_c_windowLog => CParameter::WindowLog(value as u32), + ZSTD_c_hashLog => CParameter::HashLog(value as u32), + ZSTD_c_chainLog => CParameter::ChainLog(value as u32), + ZSTD_c_searchLog => CParameter::SearchLog(value as u32), + ZSTD_c_minMatch => CParameter::MinMatch(value as u32), + ZSTD_c_targetLength => CParameter::TargetLength(value as u32), + ZSTD_c_strategy => CParameter::Strategy(strategy_from_int(value).ok_or_else( + || new_zstd_error(format!("invalid strategy value: {value}"), vm), + )?), + ZSTD_c_enableLongDistanceMatching => CParameter::EnableLongDistanceMatching(value != 0), + ZSTD_c_ldmHashLog => CParameter::LdmHashLog(value as u32), + ZSTD_c_ldmMinMatch => CParameter::LdmMinMatch(value as u32), + ZSTD_c_ldmBucketSizeLog => CParameter::LdmBucketSizeLog(value as u32), + ZSTD_c_ldmHashRateLog => CParameter::LdmHashRateLog(value as u32), + ZSTD_c_contentSizeFlag => CParameter::ContentSizeFlag(value != 0), + ZSTD_c_checksumFlag => CParameter::ChecksumFlag(value != 0), + ZSTD_c_dictIDFlag => CParameter::DictIdFlag(value != 0), + ZSTD_c_nbWorkers => CParameter::NbWorkers(value as u32), + ZSTD_c_jobSize => CParameter::JobSize(value as u32), + ZSTD_c_overlapLog => CParameter::OverlapSizeLog(value as u32), + _ => { + return Err(vm.new_value_error(format!( + "invalid compression parameter 'unknown parameter (key {param})'" + ))); + } + }; + Ok(p) + } + + /// Map a decompression-parameter id and integer value to a `DParameter`. + /// Used by the decompressor's `options=` constructor argument. + fn dparameter_from_int(param: i32, value: i32, vm: &VirtualMachine) -> PyResult { + match param { + ZSTD_d_windowLogMax => Ok(DParameter::WindowLogMax(value as u32)), + _ => Err(vm.new_value_error(format!( + "invalid decompression parameter 'unknown parameter (key {param})'" + ))), + } + } + + /// Build the CPython-compatible " parameter 'name' received an + /// illegal value V; the valid range is [lo, hi]" `ValueError` for a + /// parameter that was rejected by libzstd or that fell outside the + /// documented bounds. + fn param_value_error_for( + param: i32, + value: i32, + is_compress: bool, + vm: &VirtualMachine, + ) -> PyBaseExceptionRef { + let kind = if is_compress { "compression" } else { "decompression" }; + let name = parameter_name(param, is_compress); + match lookup_param_bounds(param, is_compress) { + Some((lo, hi)) => vm.new_value_error(format!( + "{kind} parameter '{name}' received an illegal value {value}; \ + the valid range is [{lo}, {hi}]" + )), + None => vm.new_value_error(format!( + "{kind} parameter '{name}' received an illegal value {value}" + )), + } + } + + /// Return the valid `(lower, upper)` bounds for the libzstd compression + /// level. Used when validating the `level=` argument upfront because + /// libzstd silently clamps out-of-range values rather than surfacing + /// them as errors. + fn level_bounds() -> (i32, i32) { + lookup_param_bounds(ZSTD_c_compressionLevel, true) + .expect("compressionLevel always has valid bounds") + } + + /// Look up parameter bounds for a known compression or decompression + /// parameter id. Returns `None` if the id is not recognized (callers + /// validate the id separately). + fn lookup_param_bounds(param: i32, is_compress: bool) -> Option<(i32, i32)> { + // SAFETY: `ZSTD_*Param_getBounds` reads no memory beyond the enum + // discriminant; the helpers above validated that `param` maps to a + // real enum variant. + let bounds = if is_compress { + let p = c_param_enum(param)?; + unsafe { zstd_sys::ZSTD_cParam_getBounds(p) } + } else { + let p = d_param_enum(param)?; + unsafe { zstd_sys::ZSTD_dParam_getBounds(p) } + }; + // SAFETY: ZSTD_isError just inspects the bounds.error integer. + if unsafe { zstd_sys::ZSTD_isError(bounds.error) } != 0 { + return None; + } + Some((bounds.lowerBound, bounds.upperBound)) + } + + /// Map a parameter integer id back to the Python-visible enum member + /// name. Used for error messages that pin-point the parameter that went + /// out of range. Returns `"unknown"` for unrecognized ids. + fn parameter_name(param: i32, is_compress: bool) -> &'static str { + if is_compress { + match param { + ZSTD_c_compressionLevel => "compression_level", + ZSTD_c_windowLog => "window_log", + ZSTD_c_hashLog => "hash_log", + ZSTD_c_chainLog => "chain_log", + ZSTD_c_searchLog => "search_log", + ZSTD_c_minMatch => "min_match", + ZSTD_c_targetLength => "target_length", + ZSTD_c_strategy => "strategy", + ZSTD_c_enableLongDistanceMatching => "enable_long_distance_matching", + ZSTD_c_ldmHashLog => "ldm_hash_log", + ZSTD_c_ldmMinMatch => "ldm_min_match", + ZSTD_c_ldmBucketSizeLog => "ldm_bucket_size_log", + ZSTD_c_ldmHashRateLog => "ldm_hash_rate_log", + ZSTD_c_contentSizeFlag => "content_size_flag", + ZSTD_c_checksumFlag => "checksum_flag", + ZSTD_c_dictIDFlag => "dict_id_flag", + ZSTD_c_nbWorkers => "nb_workers", + ZSTD_c_jobSize => "job_size", + ZSTD_c_overlapLog => "overlap_log", + _ => "unknown", + } + } else { + match param { + ZSTD_d_windowLogMax => "window_log_max", + _ => "unknown", + } + } + } + + /// Map a strategy integer (as exposed via the `Strategy` IntEnum) back + /// to the underlying `ZSTD_strategy` C enum value. Done via an explicit + /// match for the same reason as [`c_param_enum`]: an untrusted int + /// might not correspond to any real enum variant. + fn strategy_from_int(v: i32) -> Option { + use zstd_sys::ZSTD_strategy as S; + Some(match v { + ZSTD_fast => S::ZSTD_fast, + ZSTD_dfast => S::ZSTD_dfast, + ZSTD_greedy => S::ZSTD_greedy, + ZSTD_lazy => S::ZSTD_lazy, + ZSTD_lazy2 => S::ZSTD_lazy2, + ZSTD_btlazy2 => S::ZSTD_btlazy2, + ZSTD_btopt => S::ZSTD_btopt, + ZSTD_btultra => S::ZSTD_btultra, + ZSTD_btultra2 => S::ZSTD_btultra2, + _ => return None, + }) + } + + /// Decode the `zstd_dict=` constructor argument. Accepts either a + /// `ZstdDict` instance (treated as the default digested form) or a + /// `(ZstdDict, marker)` tuple produced by one of `ZstdDict.as_*`. + fn parse_zstd_dict_arg( + obj: PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult<(PyRef, i32)> { + if let Ok(d) = obj.clone().downcast::() { + return Ok((d, DICT_TYPE_DIGESTED)); + } + if let Ok(tuple) = obj.clone().downcast::() { + let items = tuple.as_slice(); + // Reject any tuple shape that is not (ZstdDict, int_marker) so the + // test suite's bad-args coverage (`(zd, 1.0)`, `(zd,)`, `(zd, 3)`, + // etc.) raises TypeError. Marker bounds match the three documented + // `as_*` properties. + if items.len() != 2 { + return Err(vm.new_type_error("zstd_dict argument should be a ZstdDict object")); + } + let d = items[0] + .clone() + .downcast::() + .map_err(|_| vm.new_type_error("zstd_dict argument should be a ZstdDict object"))?; + // The marker must be a plain int (not float/etc); overflow on + // `2**1000` propagates as OverflowError via `try_index`. + let marker_obj = &items[1]; + let marker = pyobj_to_i32(marker_obj, vm).map_err(|e| { + // Preserve OverflowError; everything else becomes TypeError so + // callers see a consistent "should be a ZstdDict" message. + if e.fast_isinstance(vm.ctx.exceptions.overflow_error) { + e + } else { + vm.new_type_error("zstd_dict argument should be a ZstdDict object") + } + })?; + if !(DICT_TYPE_DIGESTED..=DICT_TYPE_PREFIX).contains(&marker) { + return Err(vm.new_type_error("zstd_dict argument should be a ZstdDict object")); + } + return Ok((d, marker)); + } + Err(vm.new_type_error("zstd_dict argument should be a ZstdDict object")) + } + + // ========================================================================= + // ZstdDict + // ========================================================================= + + #[derive(FromArgs)] + pub(super) struct ZstdDictArgs { + #[pyarg(positional)] + dict_content: ArgBytesLike, + #[pyarg(named, default = false)] + is_raw: bool, + } + + /// ZstdDict(dict_content, /, *, is_raw=False) + /// + /// Represents a Zstandard dictionary. Hold the raw dictionary bytes and + /// the parsed `dict_id`. The same instance may be shared by many + /// compressors and decompressors. + /// + /// *dict_content* is a bytes-like object containing the dictionary + /// content. If *is_raw* is False (the default), the content is expected + /// to start with the Zstandard dictionary magic number and have a valid + /// header. If *is_raw* is True, arbitrary bytes are accepted and the + /// dictionary is treated as raw back-reference content. + #[pyattr] + #[pyclass(name = "ZstdDict")] + #[derive(PyPayload)] + struct ZstdDict { + dict_content: PyBytesRef, + dict_id: u32, + } + + impl std::fmt::Debug for ZstdDict { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "", + self.dict_id, + self.dict_content.len() + ) + } + } + + impl Constructor for ZstdDict { + type Args = ZstdDictArgs; + + fn py_new(_cls: &Py, args: Self::Args, vm: &VirtualMachine) -> PyResult { + let dict_content = args.dict_content.with_ref(|b| b.to_vec()); + // libzstd's `get_dict_id_from_dict` returns `None` either when + // the content is too small to contain a valid header or when it + // does not carry the dictionary magic. Both are runtime errors + // when `is_raw=False`, matching CPython's behavior of raising + // `ValueError` on a non-conformant dictionary. + let parsed_id = zstd_safe::get_dict_id_from_dict(&dict_content) + .map(|n| n.get()) + .unwrap_or(0); + if !args.is_raw && parsed_id == 0 { + return Err(vm.new_value_error( + "ZSTD_DICT_MAGIC_NUMBER not found, dict_content cannot be a 'raw content' \ + dictionary. To create a raw content dictionary, pass is_raw=True.", + )); + } + // Raw dictionaries still get a non-zero `dict_id` whenever their + // contents happen to look like a valid dict (this is the + // documented behavior tested in `test_is_raw`). + Ok(Self { + dict_content: vm.ctx.new_bytes(dict_content), + dict_id: parsed_id, + }) + } + } + + impl Representable for ZstdDict { + #[inline] + fn repr_str(zelf: &Py, _vm: &VirtualMachine) -> PyResult { + Ok(format!( + "", + zelf.dict_id, + zelf.dict_content.len() + )) + } + } + + impl rustpython_vm::types::AsMapping for ZstdDict { + fn as_mapping() -> &'static rustpython_vm::protocol::PyMappingMethods { + static AS_MAPPING: rustpython_vm::protocol::PyMappingMethods = + rustpython_vm::protocol::PyMappingMethods { + length: Some(|mapping, _vm| { + Ok(ZstdDict::mapping_downcast(mapping).dict_content.len()) + }), + subscript: None, + ass_subscript: None, + }; + &AS_MAPPING + } + } + + #[pyclass(with(Constructor, Representable, AsMapping))] + impl ZstdDict { + /// The content of a Zstandard dictionary, as a bytes object. + #[pygetset] + fn dict_content(&self) -> PyBytesRef { + self.dict_content.clone() + } + + /// The Zstandard dictionary ID. + /// + /// A non-zero value identifies a standard-format dictionary. A value + /// of zero either means the dictionary is in raw-content form or + /// that the dictionary ID was deliberately omitted from the header. + /// Dictionary IDs fall in the range 0 to 2**32 - 1. + #[pygetset] + fn dict_id(&self) -> u32 { + self.dict_id + } + + /// Load this ZstdDict as a digested dictionary. + /// + /// Returns a `(self, DICT_TYPE_DIGESTED)` tuple that can be passed + /// as the `zstd_dict=` argument of `ZstdCompressor` or + /// `ZstdDecompressor`. Digested form lets libzstd cache the parsed + /// dictionary for reuse across compression contexts. + #[pygetset] + fn as_digested_dict(zelf: PyRef, vm: &VirtualMachine) -> PyTupleRef { + vm.ctx + .new_tuple(vec![zelf.into(), vm.ctx.new_int(DICT_TYPE_DIGESTED).into()]) + } + + /// Load this ZstdDict as an undigested dictionary. + /// + /// Returns a `(self, DICT_TYPE_UNDIGESTED)` tuple. In undigested + /// form the raw dictionary bytes are copied into each + /// (de)compressor, which is slower than digested but preserves the + /// original dictionary parameters across loads. + #[pygetset] + fn as_undigested_dict(zelf: PyRef, vm: &VirtualMachine) -> PyTupleRef { + vm.ctx.new_tuple(vec![ + zelf.into(), + vm.ctx.new_int(DICT_TYPE_UNDIGESTED).into(), + ]) + } + + /// Load this ZstdDict as a prefix dictionary. + /// + /// Returns a `(self, DICT_TYPE_PREFIX)` tuple. Prefix loading treats + /// the dictionary bytes as a one-shot back-reference window for the + /// next frame only; subsequent frames will not see the prefix. + #[pygetset] + fn as_prefix(zelf: PyRef, vm: &VirtualMachine) -> PyTupleRef { + vm.ctx + .new_tuple(vec![zelf.into(), vm.ctx.new_int(DICT_TYPE_PREFIX).into()]) + } + } + + // ========================================================================= + // ZstdCompressor + // ========================================================================= + + // The three flush modes for `ZstdCompressor.compress()`, mirrored as + // class attributes via `extend_class` below. Values are positional and + // chosen to match what CPython exposes. + const COMP_MODE_CONTINUE: i32 = 0; + const COMP_MODE_FLUSH_BLOCK: i32 = 1; + const COMP_MODE_FLUSH_FRAME: i32 = 2; + + /// Internal state of a `ZstdCompressor`. Holds the libzstd context, the + /// last mode used (for `last_mode` and `set_pledged_input_size` + /// validation), and the dictionary handles that the context may reference + /// internally. Field order matters here: Rust drops in declaration order, + /// so `cctx` is freed first; the held `CDict` (if any) and the source + /// `PyRef` go away afterwards, which is the safe order for + /// teardown. + struct CompressorState { + cctx: CCtx<'static>, + /// Cached digested dictionary. The CCtx references this via + /// `ref_cdict`, so it must outlive the CCtx (handled by Rust's + /// field drop order: `cctx` drops first). + _cdict: Option>, + /// Keeps the ZstdDict's bytes alive for `ref_prefix` mode. + _dict: Option>, + last_mode: i32, + } + + /// ZstdCompressor(level=None, options=None, zstd_dict=None) + /// + /// Create a compressor object for incremental compression of data into + /// the Zstandard format. + /// + /// *level* is an int specifying the compression level. Valid range is + /// roughly [-(1<<17), 22]; 0 means "use the default level". Cannot be + /// combined with *options*. + /// + /// *options* is a dict mapping `CompressionParameter` enum members to + /// integer values for fine-grained control. Cannot be combined with + /// *level*. + /// + /// *zstd_dict* is an optional `ZstdDict` (or one of its `as_*` + /// variants) for dictionary-assisted compression. + /// + /// Use `compress()` to feed data into the compressor and `flush()` to + /// finalize the current frame. + #[pyattr] + #[pyclass(name = "ZstdCompressor")] + #[derive(PyPayload)] + struct ZstdCompressor { + state: PyMutex, + } + + impl std::fmt::Debug for ZstdCompressor { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "_zstd.ZstdCompressor") + } + } + + #[derive(FromArgs)] + pub(super) struct ZstdCompressorArgs { + #[pyarg(any, optional)] + level: OptionalArg, + #[pyarg(any, optional)] + options: OptionalArg, + #[pyarg(any, optional)] + zstd_dict: OptionalArg, + } + + /// Translate the public `mode` int to the libzstd `ZSTD_EndDirective` + /// the streaming API takes. + fn end_directive_from_mode( + mode: i32, + vm: &VirtualMachine, + ) -> PyResult { + match mode { + COMP_MODE_CONTINUE => Ok(zstd_sys::ZSTD_EndDirective::ZSTD_e_continue), + COMP_MODE_FLUSH_BLOCK => Ok(zstd_sys::ZSTD_EndDirective::ZSTD_e_flush), + COMP_MODE_FLUSH_FRAME => Ok(zstd_sys::ZSTD_EndDirective::ZSTD_e_end), + _ => Err(vm.new_value_error(format!( + "mode argument wrong value, it should be one of \ + ZstdCompressor.CONTINUE ({COMP_MODE_CONTINUE}), \ + ZstdCompressor.FLUSH_BLOCK ({COMP_MODE_FLUSH_BLOCK}), or \ + ZstdCompressor.FLUSH_FRAME ({COMP_MODE_FLUSH_FRAME})" + ))), + } + } + + impl Constructor for ZstdCompressor { + type Args = ZstdCompressorArgs; + + fn py_new(_cls: &Py, args: Self::Args, vm: &VirtualMachine) -> PyResult { + let level_opt = arg_or_none(args.level, vm); + let options_opt = arg_or_none(args.options, vm); + let dict_opt = arg_or_none(args.zstd_dict, vm); + + if level_opt.is_some() && options_opt.is_some() { + return Err(vm.new_runtime_error("Only one of level or options should be used")); + } + + let mut cctx = CCtx::<'static>::create(); + + if let Some(level_obj) = level_opt { + let level = parse_compression_level(&level_obj, vm)?; + cctx.set_parameter(CParameter::CompressionLevel(level)) + .map_err(|_| param_value_error_for(ZSTD_c_compressionLevel, level, true, vm))?; + } + + if let Some(options_obj) = options_opt { + apply_options(&mut cctx, options_obj, true, vm)?; + } + + let (held_cdict, held_dict) = load_compressor_dict(&mut cctx, dict_opt, vm)?; + + Ok(Self { + state: PyMutex::new(CompressorState { + cctx, + _cdict: held_cdict, + _dict: held_dict, + last_mode: COMP_MODE_FLUSH_FRAME, + }), + }) + } + } + + /// Parse and validate a compression `level` argument. libzstd silently + /// clamps out-of-range levels rather than erroring, but CPython surfaces + /// them as `ValueError`, and bigints become `ValueError` (not + /// `OverflowError`) for the same reason. This helper centralizes both + /// conversions so the constructor stays linear. + fn parse_compression_level(obj: &PyObjectRef, vm: &VirtualMachine) -> PyResult { + let (lo, hi) = level_bounds(); + let level = pyobj_to_i32(obj, vm).map_err(|e| { + if e.fast_isinstance(vm.ctx.exceptions.overflow_error) { + vm.new_value_error(format!( + "illegal compression level; the valid range is [{lo}, {hi}]" + )) + } else { + e + } + })?; + if level < lo || level > hi { + return Err(vm.new_value_error(format!( + "illegal compression level {level}; the valid range is [{lo}, {hi}]" + ))); + } + Ok(level) + } + + /// Drain an `options=` dict onto either a `CCtx` or a `DCtx`. Validates + /// each key/value pair (rejects wrong enum kind, rejects floats, rejects + /// out-of-range values) so the constructor's flow stays a single line. + fn apply_options( + ctx: &mut dyn ParamSetter, + options_obj: PyObjectRef, + is_compress: bool, + vm: &VirtualMachine, + ) -> PyResult<()> { + let dict = options_obj + .downcast::() + .map_err(|_| vm.new_type_error("options must be a dict"))?; + for (k, v) in dict { + // Reject the wrong enum kind ("CompressionParameter" vs + // "DecompressionParameter") before doing any numeric coercion, + // so the error carries the documented type name rather than a + // generic out-of-range message. + check_wrong_param_kind(&k, is_compress, vm)?; + let key_int = pyobj_to_i32(&k, vm)?; + let val_int = pyobj_to_i32(&v, vm)?; + // libzstd silently clamps out-of-range values for some + // parameters (notably compression_level) rather than rejecting + // them, so validate against the documented bounds upfront. + if let Some((lo, hi)) = lookup_param_bounds(key_int, is_compress) + && (val_int < lo || val_int > hi) + { + return Err(param_value_error_for(key_int, val_int, is_compress, vm)); + } + ctx.apply(key_int, val_int, vm)?; + } + Ok(()) + } + + /// Trait wrapper over `CCtx::set_parameter` and `DCtx::set_parameter` so + /// `apply_options` can drive either context without duplicated code. + /// `set_parameter` translates the (id, value) pair to the appropriate + /// `CParameter` / `DParameter` enum variant, then forwards to libzstd. + trait ParamSetter { + fn apply(&mut self, param: i32, value: i32, vm: &VirtualMachine) -> PyResult<()>; + } + + impl ParamSetter for CCtx<'static> { + fn apply(&mut self, param: i32, value: i32, vm: &VirtualMachine) -> PyResult<()> { + let p = cparameter_from_int(param, value, vm)?; + self.set_parameter(p) + .map_err(|_| param_value_error_for(param, value, true, vm))?; + Ok(()) + } + } + + impl ParamSetter for DCtx<'static> { + fn apply(&mut self, param: i32, value: i32, vm: &VirtualMachine) -> PyResult<()> { + let p = dparameter_from_int(param, value, vm)?; + self.set_parameter(p) + .map_err(|_| param_value_error_for(param, value, false, vm))?; + Ok(()) + } + } + + /// Trait that captures the only differences between how the compressor + /// and decompressor consume a dictionary: the name of the type that + /// appears in error messages, the eager-validation constructor for the + /// digested variant, and the three ways of attaching it to the context. + trait DictLoader<'a> { + type Digested; + const KIND_NAME: &'static str; + fn try_create_digested(bytes: &[u8]) -> Option; + fn ref_digested(&mut self, dict: &Self::Digested) -> zstd_safe::SafeResult; + fn load_undigested(&mut self, bytes: &[u8]) -> zstd_safe::SafeResult; + fn ref_prefix_static(&mut self, bytes: &'static [u8]) -> zstd_safe::SafeResult; + } + + impl DictLoader<'static> for CCtx<'static> { + type Digested = zstd_safe::CDict<'static>; + const KIND_NAME: &'static str = "ZSTD_CDict"; + fn try_create_digested(bytes: &[u8]) -> Option { + zstd_safe::CDict::try_create(bytes, ZSTD_CLEVEL_DEFAULT) + } + fn ref_digested(&mut self, dict: &Self::Digested) -> zstd_safe::SafeResult { + self.ref_cdict(dict) + } + fn load_undigested(&mut self, bytes: &[u8]) -> zstd_safe::SafeResult { + self.load_dictionary(bytes) + } + fn ref_prefix_static(&mut self, bytes: &'static [u8]) -> zstd_safe::SafeResult { + self.ref_prefix(bytes) + } + } + + impl DictLoader<'static> for DCtx<'static> { + type Digested = zstd_safe::DDict<'static>; + const KIND_NAME: &'static str = "ZSTD_DDict"; + fn try_create_digested(bytes: &[u8]) -> Option { + zstd_safe::DDict::try_create(bytes) + } + fn ref_digested(&mut self, dict: &Self::Digested) -> zstd_safe::SafeResult { + self.ref_ddict(dict) + } + fn load_undigested(&mut self, bytes: &[u8]) -> zstd_safe::SafeResult { + self.load_dictionary(bytes) + } + fn ref_prefix_static(&mut self, bytes: &'static [u8]) -> zstd_safe::SafeResult { + self.ref_prefix(bytes) + } + } + + /// Common path for attaching a dictionary to either context type. Returns + /// the digested `CDict`/`DDict` (if the caller used digested mode) plus + /// the `PyRef` we hold to keep the bytes alive for `ref_prefix`. + fn load_dict>( + ctx: &mut L, + dict_obj: Option, + vm: &VirtualMachine, + ) -> PyResult<(Option, Option>)> { + let Some(dict_obj) = dict_obj else { + return Ok((None, None)); + }; + let (zdict, marker) = parse_zstd_dict_arg(dict_obj, vm)?; + let bad_dict_err = || -> PyBaseExceptionRef { + new_zstd_error( + format!( + "Failed to load the {} instance from corrupted Zstandard dictionary content.", + L::KIND_NAME + ), + vm, + ) + }; + let dict_bytes = zdict.dict_content.as_bytes(); + let mut digested = None; + match marker { + DICT_TYPE_PREFIX => { + // SAFETY: the returned `PyRef` keeps `zdict` alive + // (and therefore `dict_bytes`) for the lifetime of the + // context. libzstd `ref_prefix` stores a raw pointer into + // `dict_bytes`; as long as the (de)compressor outlives no + // longer than the held PyRef, the pointer stays valid. + let static_bytes: &'static [u8] = unsafe { + std::slice::from_raw_parts(dict_bytes.as_ptr(), dict_bytes.len()) + }; + ctx.ref_prefix_static(static_bytes) + .map_err(|_| bad_dict_err())?; + } + DICT_TYPE_DIGESTED => { + // Build the digested dict eagerly so a corrupted dictionary + // surfaces as a `ZstdError` at construction time, not when + // the first compress/decompress call runs. + let d = L::try_create_digested(dict_bytes).ok_or_else(bad_dict_err)?; + ctx.ref_digested(&d).map_err(|_| bad_dict_err())?; + digested = Some(d); + } + _ => { + // Undigested: copy the bytes into the context. Validation + // happens lazily at the first stream call in this mode. + ctx.load_undigested(dict_bytes).map_err(|_| bad_dict_err())?; + } + } + Ok((digested, Some(zdict))) + } + + fn load_compressor_dict( + cctx: &mut CCtx<'static>, + dict_obj: Option, + vm: &VirtualMachine, + ) -> PyResult<(Option>, Option>)> { + load_dict::>(cctx, dict_obj, vm) + } + + fn load_decompressor_dict( + dctx: &mut DCtx<'static>, + dict_obj: Option, + vm: &VirtualMachine, + ) -> PyResult<(Option>, Option>)> { + load_dict::>(dctx, dict_obj, vm) + } + + + /// Drive `compress_stream2` until the input is fully consumed and, for + /// flush/end directives, the internal buffers report zero remaining bytes. + /// Grows the output `Vec` by `CStreamOutSize` chunks as needed. + fn do_compress( + state: &mut CompressorState, + data: &[u8], + end_op: zstd_sys::ZSTD_EndDirective, + vm: &VirtualMachine, + ) -> PyResult> { + let mut output = Vec::new(); + let mut input = InBuffer::around(data); + let chunk_size = CCtx::out_size().max(1); + let is_end = end_op != zstd_sys::ZSTD_EndDirective::ZSTD_e_continue; + + loop { + let prev_len = output.len(); + output.reserve(chunk_size); + let remaining = { + let mut out_buf = OutBuffer::around_pos(&mut output, prev_len); + state.cctx.compress_stream2(&mut out_buf, &mut input, end_op) + } + .map_err(|c| catch_zstd_error(c, vm))?; + let consumed_all = input.pos == input.src.len(); + // Stop when input is fully consumed and, for flush/end directives, + // libzstd reports that all internal buffers have been drained + // (remaining == 0). Otherwise loop; the next `reserve` will grow + // the output buffer if we hit the previous capacity. + if consumed_all && (!is_end || remaining == 0) { + break; + } + } + Ok(output) + } + + #[pyclass(with(Constructor))] + impl ZstdCompressor { + /// compress(data, mode=ZstdCompressor.CONTINUE) + /// + /// Provide data to the compressor and return a chunk of compressed + /// output bytes (which may be empty if not enough data has been + /// buffered yet). + /// + /// *mode* controls when to flush the output: + /// * `CONTINUE` (default) keeps buffering, producing the densest + /// compression. The compressor may return an empty byte string. + /// * `FLUSH_BLOCK` closes the current zstd block so all data + /// provided so far can be immediately decompressed. Frequent + /// use reduces the compression ratio. + /// * `FLUSH_FRAME` closes the current frame entirely. The next + /// `compress()` call starts a new frame. + #[pymethod] + fn compress(&self, args: CompressMethodArgs, vm: &VirtualMachine) -> PyResult> { + let mode = args.mode.unwrap_or(COMP_MODE_CONTINUE); + let end_op = end_directive_from_mode(mode, vm)?; + let data = args.data.with_ref(|b| b.to_vec()); + let mut state = self.state.lock(); + let out = do_compress(&mut state, &data, end_op, vm)?; + state.last_mode = mode; + Ok(out) + } + + /// flush(mode=ZstdCompressor.FLUSH_FRAME) + /// + /// Finish the compression process and return all the remaining + /// compressed bytes. *mode* must be one of `FLUSH_FRAME` (close the + /// current frame; new frames may be started afterward) or + /// `FLUSH_BLOCK` (close the current block but stay inside the + /// frame). `CONTINUE` is not a valid value for `flush()`. + #[pymethod] + fn flush(&self, args: FlushMethodArgs, vm: &VirtualMachine) -> PyResult> { + let mode = args.mode.unwrap_or(COMP_MODE_FLUSH_FRAME); + if mode != COMP_MODE_FLUSH_BLOCK && mode != COMP_MODE_FLUSH_FRAME { + return Err(vm.new_value_error(format!( + "mode argument wrong value, it should be \ + ZstdCompressor.FLUSH_FRAME ({COMP_MODE_FLUSH_FRAME}) or \ + ZstdCompressor.FLUSH_BLOCK ({COMP_MODE_FLUSH_BLOCK})" + ))); + } + let end_op = end_directive_from_mode(mode, vm)?; + let mut state = self.state.lock(); + let out = do_compress(&mut state, &[], end_op, vm)?; + state.last_mode = mode; + Ok(out) + } + + /// set_pledged_input_size(size, /) + /// + /// Set the uncompressed content size to be written into the frame + /// header of the next frame. Useful when the size is known up front: + /// libzstd writes it into the header (unless the content-size flag + /// is disabled) and verifies that the actual input matches at frame + /// end. Pass `None` to mean "unknown". + /// + /// May only be called when `last_mode == FLUSH_FRAME` (i.e. between + /// frames; not in the middle of one). + #[pymethod] + fn set_pledged_input_size(&self, size: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { + let mut state = self.state.lock(); + if state.last_mode != COMP_MODE_FLUSH_FRAME { + return Err(vm.new_value_error( + "set_pledged_input_size() method must be called when last_mode == FLUSH_FRAME", + )); + } + // Python passes `None` to mean "unknown"; libzstd represents that + // internally as `ZSTD_CONTENTSIZE_UNKNOWN` (`u64::MAX`), and + // `zstd_safe` translates `None` accordingly. Reject any concrete + // int outside the valid range up front so callers see the + // documented `ValueError`, not a libzstd-level error. The valid + // upper bound is `ZSTD_CONTENTSIZE_ERROR - 1`, i.e. `u64::MAX - 2`. + let pledged: Option = if vm.is_none(&size) { + None + } else { + const UPPER_BOUND: u64 = u64::MAX - 2; + let int_ref = size.try_index(vm)?; + // Try fitting into u64; OverflowError covers negatives and + // values above `u64::MAX`. + let v: u64 = int_ref.try_to_primitive(vm).map_err(|_| { + vm.new_value_error(format!( + "size argument should be a positive int less than {UPPER_BOUND}" + )) + })?; + if v > UPPER_BOUND { + return Err(vm.new_value_error(format!( + "size argument should be a positive int less than {UPPER_BOUND}" + ))); + } + Some(v) + }; + state + .cctx + .set_pledged_src_size(pledged) + .map_err(|c| catch_zstd_error(c, vm))?; + Ok(()) + } + + /// The mode argument passed to the most recent `compress()` or + /// `flush()` call. + /// + /// The value is one of `CONTINUE`, `FLUSH_BLOCK`, or `FLUSH_FRAME`. + /// The initial value is `FLUSH_FRAME`, which signals that no frame + /// is currently in progress (it is safe to call + /// `set_pledged_input_size()` in this state). + #[pygetset] + fn last_mode(&self) -> i32 { + self.state.lock().last_mode + } + + /// Install class-level constants `CONTINUE`, `FLUSH_BLOCK`, and + /// `FLUSH_FRAME` so callers can reference them as + /// `ZstdCompressor.FLUSH_FRAME` (as the Python `ZstdFile` wrapper + /// does). + #[extend_class] + fn extend_class(ctx: &Context, class: &'static Py) { + class.set_attr( + ctx.intern_str("CONTINUE"), + ctx.new_int(COMP_MODE_CONTINUE).into(), + ); + class.set_attr( + ctx.intern_str("FLUSH_BLOCK"), + ctx.new_int(COMP_MODE_FLUSH_BLOCK).into(), + ); + class.set_attr( + ctx.intern_str("FLUSH_FRAME"), + ctx.new_int(COMP_MODE_FLUSH_FRAME).into(), + ); + } + } + + #[derive(FromArgs)] + pub(super) struct CompressMethodArgs { + #[pyarg(positional)] + data: ArgBytesLike, + #[pyarg(any, optional)] + mode: Option, + } + + #[derive(FromArgs)] + pub(super) struct FlushMethodArgs { + #[pyarg(any, optional)] + mode: Option, + } + + // ========================================================================= + // ZstdDecompressor + // ========================================================================= + + /// Internal state of a `ZstdDecompressor`. The CPython decompressor is + /// single-frame: once we hit end-of-frame, additional bytes go into + /// `unused_data` and further `decompress` calls raise `EOFError`. Field + /// drop order matters here for the same reason as in `CompressorState`: + /// the `dctx` is freed first and must give up its internal pointers + /// before any referenced `DDict`/`PyRef` is dropped. + struct DecompressorState { + dctx: DCtx<'static>, + /// Cached decompression dictionary referenced by the DCtx. + _ddict: Option>, + _dict: Option>, + eof: bool, + needs_input: bool, + /// Bytes that arrived after the end of the first frame. + unused_data: PyBytesRef, + /// Input bytes buffered because the previous `decompress` call ran + /// into its `max_length` cap before consuming them all. + input_buffer: Vec, + } + + /// ZstdDecompressor(zstd_dict=None, options=None) + /// + /// Create a decompressor object for incremental decompression of a + /// single Zstandard frame. Once the frame is fully decompressed the + /// decompressor is at EOF; trailing bytes are exposed via + /// `unused_data`. For multi-frame input use the module-level + /// `compression.zstd.decompress()` or `ZstdFile`. + /// + /// *zstd_dict* is an optional `ZstdDict` (or one of its `as_*` + /// variants) matching the dictionary used during compression. + /// + /// *options* is a dict mapping `DecompressionParameter` enum members + /// to integer values for fine-grained control. + #[pyattr] + #[pyclass(name = "ZstdDecompressor")] + #[derive(PyPayload)] + struct ZstdDecompressor { + state: PyMutex, + } + + impl std::fmt::Debug for ZstdDecompressor { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "_zstd.ZstdDecompressor") + } + } + + #[derive(FromArgs)] + pub(super) struct ZstdDecompressorArgs { + #[pyarg(any, optional)] + zstd_dict: OptionalArg, + #[pyarg(any, optional)] + options: OptionalArg, + } + + impl Constructor for ZstdDecompressor { + type Args = ZstdDecompressorArgs; + + fn py_new(_cls: &Py, args: Self::Args, vm: &VirtualMachine) -> PyResult { + let dict_opt = arg_or_none(args.zstd_dict, vm); + let options_opt = arg_or_none(args.options, vm); + + let mut dctx = DCtx::<'static>::create(); + + if let Some(options_obj) = options_opt { + apply_options(&mut dctx, options_obj, false, vm)?; + } + + let (held_ddict, held_dict) = load_decompressor_dict(&mut dctx, dict_opt, vm)?; + + Ok(Self { + state: PyMutex::new(DecompressorState { + dctx, + _ddict: held_ddict, + _dict: held_dict, + eof: false, + needs_input: true, + unused_data: vm.ctx.empty_bytes.clone(), + input_buffer: Vec::new(), + }), + }) + } + } + + #[derive(FromArgs)] + pub(super) struct DecompressMethodArgs { + #[pyarg(positional)] + data: ArgBytesLike, + #[pyarg(any, default = -1)] + max_length: isize, + } + + /// Drive `decompress_stream` until either the frame ends, the input is + /// exhausted with no more output coming, or `max_length` bytes have been + /// produced. Sets the various `state` flags to reflect the new situation. + /// + /// Loop control: we always keep going while either the input still has + /// bytes to feed OR the previous call filled the output buffer (which + /// indicates libzstd had more to emit but ran out of room). We only stop + /// short of a frame boundary when the input is exhausted AND libzstd had + /// room left in the output buffer, which means it is genuinely waiting + /// for more compressed bytes. + fn do_decompress( + state: &mut DecompressorState, + new_data: &[u8], + max_length: Option, + vm: &VirtualMachine, + ) -> PyResult> { + // Combine any buffered leftover input with the new data so the + // decompressor sees one contiguous stream. `Cow` avoids the + // allocation when there is no leftover. + let work_data: std::borrow::Cow<'_, [u8]> = if state.input_buffer.is_empty() { + std::borrow::Cow::Borrowed(new_data) + } else { + let mut combined = Vec::with_capacity(state.input_buffer.len() + new_data.len()); + combined.extend_from_slice(&state.input_buffer); + combined.extend_from_slice(new_data); + std::borrow::Cow::Owned(combined) + }; + + let mut input = InBuffer::around(&*work_data); + let mut output: Vec = Vec::new(); + let chunk_size = DCtx::out_size().max(1); + // Reusable scratch buffer for each decompress_stream call. We need an + // exact-size output buffer because `Vec::reserve` may over-allocate; + // OutBuffer reports the full Vec capacity to libzstd, which would + // then happily write past `max_length`. + let mut scratch: Vec = vec![0u8; chunk_size]; + let mut hit_max = false; + let mut iteration = 0usize; + + loop { + iteration += 1; + // Honor `max_length`: stop growing the output buffer once we have + // produced enough. Special-case the first iteration when the cap + // is zero so a zero-output frame (skippable frame, empty content + // frame) can still complete; we hand libzstd a 1-byte scratch + // and discard the byte if it ends up writing one. + let grow = match max_length { + Some(maxl) if output.len() >= maxl && iteration > 1 => { + hit_max = true; + break; + } + Some(maxl) if output.len() >= maxl => 1, + Some(maxl) => (maxl - output.len()).min(chunk_size), + None => chunk_size, + }; + let result; + let written; + { + let slot = &mut scratch[..grow]; + let mut out_buf = OutBuffer::around(slot as &mut [u8]); + result = state.dctx.decompress_stream(&mut out_buf, &mut input); + written = out_buf.pos(); + } + output.extend_from_slice(&scratch[..written]); + match result { + Ok(0) => { + // Frame fully decompressed; the decompressor is at EOF. + state.eof = true; + break; + } + Ok(_) => { + let output_was_full = written == grow; + let input_consumed = input.pos == input.src.len(); + + if let Some(maxl) = max_length + && output.len() >= maxl + && iteration > 1 + { + hit_max = true; + break; + } + + // Input is gone and libzstd had room to write but did not, + // which means the frame is incomplete and the caller has + // to supply more input. + if input_consumed && !output_was_full { + break; + } + } + Err(code) => return Err(catch_zstd_error(code, vm)), + } + } + + // If `max_length == 0` opened a courtesy iteration that produced more + // bytes than the caller asked for, truncate. Should not happen with + // the scratch slicing above, but keep the safety net. + if let Some(maxl) = max_length + && output.len() > maxl + { + output.truncate(maxl); + hit_max = true; + } + + let consumed = input.pos; + let remaining = &work_data[consumed..]; + + if state.eof { + if !remaining.is_empty() { + state.unused_data = vm.ctx.new_bytes(remaining.to_vec()); + } + state.input_buffer.clear(); + state.needs_input = false; + } else if hit_max { + // Output cap reached with input still pending. Buffer the rest + // and report `needs_input == false` so the caller knows to call + // `decompress(b'', max_length=...)` to drain it. + state.input_buffer = remaining.to_vec(); + state.needs_input = false; + } else if max_length == Some(0) { + // Caller explicitly asked for zero output bytes. Keep whatever + // input is left around for the next call and signal that they + // do not need to feed more right now. CPython's decompressor + // treats `max_length=0` as "stop here without losing state". + state.input_buffer = remaining.to_vec(); + state.needs_input = false; + } else { + // All input consumed but the frame is not complete; the caller + // should provide more data on the next call. + state.input_buffer.clear(); + state.needs_input = true; + } + + Ok(output) + } + + #[pyclass(with(Constructor))] + impl ZstdDecompressor { + /// decompress(data, max_length=-1) + /// + /// Decompress *data* and return at most *max_length* bytes of + /// uncompressed output. A *max_length* of -1 (the default) means + /// "no limit". After the current frame is fully decompressed, `eof` + /// becomes `True` and any further `decompress()` calls raise + /// `EOFError`; bytes that arrived after the frame go into + /// `unused_data` for the caller to forward. + /// + /// If the output cap is reached with input still pending, the + /// remaining input is buffered internally and `needs_input` is set + /// to `False`, signalling that the next call can be made with an + /// empty *data* argument to keep draining. + #[pymethod] + fn decompress(&self, args: DecompressMethodArgs, vm: &VirtualMachine) -> PyResult> { + let data_vec = args.data.with_ref(|b| b.to_vec()); + let max_length = if args.max_length < 0 { + None + } else { + Some(args.max_length as usize) + }; + let mut state = self.state.lock(); + if state.eof { + return Err(vm.new_exception_msg( + vm.ctx.exceptions.eof_error.to_owned(), + "Already at the end of a Zstandard frame.".to_owned().into(), + )); + } + do_decompress(&mut state, &data_vec, max_length, vm) + } + + /// True if the end-of-frame marker has been reached. + #[pygetset] + fn eof(&self) -> bool { + self.state.lock().eof + } + + /// True when the decompressor expects more compressed input. + /// + /// Becomes False after a `decompress()` call that hits its + /// `max_length` cap with input still buffered internally; in that + /// state the caller should call `decompress(b'', max_length=...)` + /// to drain the buffered output before providing more input. + #[pygetset] + fn needs_input(&self) -> bool { + self.state.lock().needs_input + } + + /// Bytes left over after the end-of-frame marker. + /// + /// Empty until the decompressor reaches EOF. Read-only. + #[pygetset] + fn unused_data(&self) -> PyBytesRef { + self.state.lock().unused_data.clone() + } + } + + // ========================================================================= + // Module-level functions + // ========================================================================= + + /// get_frame_size(frame_buffer, /) + /// + /// Return the size in bytes of the first complete Zstandard frame in + /// *frame_buffer* (which must contain the entire frame, header through + /// trailer). Raises `ZstdError` if *frame_buffer* is shorter than the + /// complete frame. + #[pyfunction] + fn get_frame_size(frame_buffer: ArgBytesLike, vm: &VirtualMachine) -> PyResult { + let buf = frame_buffer.with_ref(|b| b.to_vec()); + zstd_safe::find_frame_compressed_size(&buf).map_err(|_| { + new_zstd_error( + "Error when finding the compressed size of a Zstandard frame. \ + Ensure the frame_buffer argument starts from the beginning of a frame, \ + and its length not less than this complete frame.", + vm, + ) + }) + } + + /// get_frame_info(frame_buffer, /) + /// + /// Return `(decompressed_size, dictionary_id)` for the first Zstandard + /// frame in *frame_buffer*. `decompressed_size` is `None` when the + /// frame header does not record it (small or streaming frames). + /// `dictionary_id` is 0 if no dictionary was used. *frame_buffer* must + /// contain at least the 6-18 byte frame header. The pure-Python + /// wrapper boxes this tuple into a `FrameInfo` named tuple. + #[pyfunction] + fn get_frame_info( + frame_buffer: ArgBytesLike, + vm: &VirtualMachine, + ) -> PyResult<(PyObjectRef, u32)> { + let buf = frame_buffer.with_ref(|b| b.to_vec()); + let content_size = zstd_safe::get_frame_content_size(&buf).map_err(|_| { + new_zstd_error( + "Error when getting information from the header of a Zstandard frame. \ + Ensure the frame_buffer argument starts from the beginning of a frame, \ + and its length not less than the frame header (6~18 bytes).", + vm, + ) + })?; + let content_size_obj: PyObjectRef = match content_size { + Some(n) => vm.ctx.new_int(n).into(), + None => vm.ctx.none(), + }; + let dict_id = zstd_safe::get_dict_id_from_frame(&buf) + .map(|n| n.get()) + .unwrap_or(0); + Ok((content_size_obj, dict_id)) + } + + #[derive(FromArgs)] + pub(super) struct TrainDictArgs { + /// Concatenated sample bytes. Must be a `bytes` object, not + /// `bytearray` or another buffer type, to match CPython's strict + /// type-checking on this argument. + #[pyarg(positional)] + samples_bytes: PyBytesRef, + /// A tuple of integer sample sizes that partition `samples_bytes`. + /// Lists and other iterables are not accepted. + #[pyarg(positional)] + samples_sizes: PyTupleRef, + /// Maximum size of the returned dictionary, in bytes. Must be a + /// positive `int`. + #[pyarg(positional)] + dict_size: PyObjectRef, + } + + /// Collect the elements of `tuple` into a `Vec`, validating that + /// each element is a non-negative int that fits in `usize`. Used by both + /// `train_dict` and `finalize_dict` for the `samples_sizes` argument. + /// + /// Floats (and any object whose `__index__` slot is missing) raise + /// `TypeError`; values that do not fit `usize` raise `ValueError` so + /// the test suite's `(2**1000,)` / `(-1,)` coverage holds. + fn parse_sample_sizes(tuple: PyTupleRef, vm: &VirtualMachine) -> PyResult> { + let items = tuple.as_slice(); + let mut out = Vec::with_capacity(items.len()); + for item in items { + let idx = item.try_index(vm)?; + let v: usize = idx.try_to_primitive(vm).map_err(|_| { + vm.new_value_error("sample size out of range for size_t".to_owned()) + })?; + out.push(v); + } + Ok(out) + } + + /// Convert a Python `int` to a positive `isize`. Rejects floats (via + /// `try_index`) and non-positive values; bigints that don't fit `isize` + /// propagate as `OverflowError`. Used for the `dict_size` argument of + /// `train_dict` and `finalize_dict`, which must always be a strictly + /// positive int. + fn parse_positive_dict_size(obj: &PyObjectRef, vm: &VirtualMachine) -> PyResult { + let idx = obj.try_index(vm)?; + // `try_to_primitive::` raises `OverflowError` on bigints out of + // range; pass that through verbatim so the test suite's + // `assertRaises(OverflowError)` coverage matches. + let v: isize = idx.try_to_primitive(vm)?; + if v <= 0 { + return Err(vm.new_value_error("dict_size must be positive")); + } + Ok(v as usize) + } + + /// train_dict(samples_bytes, samples_sizes, dict_size, /) + /// + /// Train a Zstandard dictionary from sample data and return the raw + /// dictionary bytes. The pure-Python `train_dict` wrapper then wraps + /// the result in a `ZstdDict`. + /// + /// `samples_bytes` is the concatenation of all sample contents. + /// `samples_sizes` is a tuple giving the length of each sample, in the + /// same order they appear in `samples_bytes`. `dict_size` is the + /// maximum size of the dictionary in bytes. + #[pyfunction] + fn train_dict(args: TrainDictArgs, vm: &VirtualMachine) -> PyResult> { + let dict_size = parse_positive_dict_size(&args.dict_size, vm)?; + let samples_buffer = args.samples_bytes.as_bytes().to_vec(); + let sizes = parse_sample_sizes(args.samples_sizes, vm)?; + let total: usize = sizes.iter().sum(); + if total != samples_buffer.len() { + return Err( + vm.new_value_error("The samples size tuple doesn't match the concatenation's size") + ); + } + let mut dict_buffer: Vec = Vec::with_capacity(dict_size); + zstd_safe::train_from_buffer(&mut dict_buffer, &samples_buffer, &sizes) + .map_err(|c| catch_zstd_error(c, vm))?; + Ok(dict_buffer) + } + + #[derive(FromArgs)] + pub(super) struct FinalizeDictArgs { + /// Raw "starting" dictionary content to finalize. Must be `bytes` + /// (not `bytearray`) to match CPython. + #[pyarg(positional)] + custom_dict_bytes: PyBytesRef, + /// Concatenated sample bytes used to derive the dictionary's + /// statistics tables. Must be `bytes`. + #[pyarg(positional)] + samples_bytes: PyBytesRef, + /// Tuple of integer sample sizes partitioning `samples_bytes`. + #[pyarg(positional)] + samples_sizes: PyTupleRef, + /// Maximum size of the finalized dictionary, in bytes. Positive int. + #[pyarg(positional)] + dict_size: PyObjectRef, + /// Compression level the dictionary will be tuned for. Must be int. + #[pyarg(positional)] + compression_level: PyObjectRef, + } + + /// finalize_dict(custom_dict_bytes, samples_bytes, samples_sizes, dict_size, compression_level, /) + /// + /// Take a hand-crafted custom dictionary plus a set of samples and + /// produce a standard-format Zstandard dictionary by appending the + /// header and statistics tables. `zstd_safe` does not wrap this + /// function, so we drop down to raw `zstd_sys` FFI for it. + #[pyfunction] + fn finalize_dict(args: FinalizeDictArgs, vm: &VirtualMachine) -> PyResult> { + let dict_size = parse_positive_dict_size(&args.dict_size, vm)?; + let compression_level = pyobj_to_i32(&args.compression_level, vm)?; + let custom_dict = args.custom_dict_bytes.as_bytes().to_vec(); + let samples_buffer = args.samples_bytes.as_bytes().to_vec(); + let sizes = parse_sample_sizes(args.samples_sizes, vm)?; + let total: usize = sizes.iter().sum(); + if total != samples_buffer.len() { + return Err( + vm.new_value_error("The samples size tuple doesn't match the concatenation's size") + ); + } + + let mut dict_buffer: Vec = vec![0u8; dict_size]; + let params = zstd_sys::ZDICT_params_t { + compressionLevel: compression_level, + notificationLevel: 0, + dictID: 0, + }; + + // SAFETY: All pointers point into Rust-owned, properly sized buffers + // that outlive the FFI call. ZDICT_finalizeDictionary just reads from + // the sample/dict buffers and writes into `dict_buffer`. + let written = unsafe { + zstd_sys::ZDICT_finalizeDictionary( + dict_buffer.as_mut_ptr() as *mut _, + dict_buffer.len(), + custom_dict.as_ptr() as *const _, + custom_dict.len(), + samples_buffer.as_ptr() as *const _, + sizes.as_ptr(), + sizes.len() as u32, + params, + ) + }; + // SAFETY: ZDICT_isError just inspects the integer return code. + if unsafe { zstd_sys::ZDICT_isError(written) } != 0 { + // SAFETY: ZDICT_getErrorName returns a static NUL-terminated + // C string from libzstd's internal error table. + let err_ptr = unsafe { zstd_sys::ZDICT_getErrorName(written) }; + let msg = if err_ptr.is_null() { + "zstd dictionary finalization failed".to_string() + } else { + unsafe { std::ffi::CStr::from_ptr(err_ptr) } + .to_string_lossy() + .into_owned() + }; + return Err(new_zstd_error(msg, vm)); + } + dict_buffer.truncate(written); + Ok(dict_buffer) + } + + #[derive(FromArgs)] + pub(super) struct ParamBoundsArgs { + #[pyarg(positional)] + parameter: i32, + #[pyarg(named)] + is_compress: bool, + } + + /// get_param_bounds(parameter, /, *, is_compress) + /// + /// Return the inclusive `(lower, upper)` integer bounds for a + /// compression parameter (when *is_compress* is True) or a + /// decompression parameter (when False). Used by + /// `CompressionParameter.bounds()` and `DecompressionParameter.bounds()` + /// in the pure-Python wrapper. + #[pyfunction] + fn get_param_bounds(args: ParamBoundsArgs, vm: &VirtualMachine) -> PyResult<(c_int, c_int)> { + let bounds = if args.is_compress { + // Validate by mapping the int to a CParameter variant first; if + // the int is not a known parameter id we surface a clear error + // instead of letting libzstd return a generic "parameter unsupported". + let _ = cparameter_from_int(args.parameter, 0, vm)?; + // SAFETY: we just confirmed `args.parameter` is a value of the + // `ZSTD_cParameter` enum. + let param_enum = unsafe { + std::mem::transmute::(args.parameter as u32) + }; + unsafe { zstd_sys::ZSTD_cParam_getBounds(param_enum) } + } else { + let _ = dparameter_from_int(args.parameter, 0, vm)?; + // SAFETY: same justification as the compress branch above. + let param_enum = unsafe { + std::mem::transmute::(args.parameter as u32) + }; + unsafe { zstd_sys::ZSTD_dParam_getBounds(param_enum) } + }; + // SAFETY: ZSTD_isError just inspects the integer error code. + if unsafe { zstd_sys::ZSTD_isError(bounds.error) } != 0 { + return Err(catch_zstd_error(bounds.error, vm)); + } + Ok((bounds.lowerBound, bounds.upperBound)) + } + + /// set_parameter_types(c_parameter_type, d_parameter_type, /) + /// + /// Register the `CompressionParameter` and `DecompressionParameter` + /// `IntEnum` types defined in the pure-Python wrapper. Once registered, + /// the (de)compressor constructors validate that callers do not mix the + /// two enum kinds (e.g. passing a `CompressionParameter` key to a + /// `ZstdDecompressor`), raising `TypeError` if they do. + /// + /// The Python wrapper at `Lib/compression/zstd/__init__.py:242` calls + /// this exactly once at import time. + #[pyfunction] + fn set_parameter_types( + c_parameter_type: PyObjectRef, + d_parameter_type: PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult<()> { + let c_type = c_parameter_type + .downcast::() + .map_err(|_| vm.new_type_error("c_parameter_type must be a type object"))?; + let d_type = d_parameter_type + .downcast::() + .map_err(|_| vm.new_type_error("d_parameter_type must be a type object"))?; + let mut guard = REGISTERED_PARAM_TYPES.lock(); + guard.compression = Some(c_type); + guard.decompression = Some(d_type); + Ok(()) + } +} From 1def1b4ec74dc96971b0d539e426416321c68e63 Mon Sep 17 00:00:00 2001 From: James David Clarke Date: Sat, 23 May 2026 18:33:36 +0100 Subject: [PATCH 02/20] Fix iOS build by removing static PyMutex `PyMutex` is built on `RawCellMutex` on single-threaded targets (iOS, Android, wasm32) and is not `Sync`, so a `static PyMutex<...>` fails to compile. `PyTypeRef` is also not `Send` by default, which compounds the issue. Replace the static parameter-type registry with class-name comparison (`"CompressionParameter"` / `"DecompressionParameter"`). The two names originate from the `compression.zstd` module we ship, so identity vs. name comparison is equivalent in practice. `set_parameter_types` now just validates that its arguments are type objects. Verified locally: - `cargo build -p rustpython-stdlib` (default features): clean. - `cargo check -p rustpython-stdlib --no-default-features --features host_env` (simulates iOS by dropping the `threading` feature): clean. - `cargo run -- -m test test_zstd`: 119 tests pass, 1 expected failure. Co-Authored-By: Claude Opus 4.7 (1M context) --- crates/stdlib/src/zstd.rs | 81 ++++++++++++++++++--------------------- 1 file changed, 37 insertions(+), 44 deletions(-) diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs index 6c9ad33396e..7f0ef4990f5 100644 --- a/crates/stdlib/src/zstd.rs +++ b/crates/stdlib/src/zstd.rs @@ -37,23 +37,15 @@ mod _zstd { use zstd_safe::zstd_sys; use zstd_safe::{CCtx, CParameter, DCtx, DParameter, InBuffer, OutBuffer}; - /// Stores the `CompressionParameter` and `DecompressionParameter` enum - /// types that the pure-Python `compression.zstd` wrapper passes through - /// `set_parameter_types`. They let us produce CPython-compatible - /// `TypeError`s when callers pass the wrong enum kind to a (de)compressor - /// (a `CompressionParameter` key to a decompressor, or vice versa). Stored - /// in a `PyMutex` so the module init can set them once and constructors - /// can read them later from any thread. - static REGISTERED_PARAM_TYPES: PyMutex = - PyMutex::new(RegisteredParamTypes { - compression: None, - decompression: None, - }); - - struct RegisteredParamTypes { - compression: Option, - decompression: Option, - } + /// Class names of the `CompressionParameter` and `DecompressionParameter` + /// `IntEnum`s defined by the pure-Python wrapper at + /// `Lib/compression/zstd/__init__.py`. We identify these by name rather + /// than identity because storing `PyTypeRef` in a Rust `static` requires + /// `Sync`, which `PyMutex` does not provide on single-threaded targets + /// (Android, iOS, wasm32). Name comparison is sufficient since both names + /// originate from a module we ship and control. + const C_PARAMETER_TYPE_NAME: &str = "CompressionParameter"; + const D_PARAMETER_TYPE_NAME: &str = "DecompressionParameter"; // ========================================================================= // Module-level constants @@ -207,29 +199,29 @@ mod _zstd { arg.into_option().filter(|o| !vm.is_none(o)) } - /// If `key` is a registered `CompressionParameter` or `DecompressionParameter` - /// enum member and its kind does not match the caller's context, return a - /// `TypeError`. Used by both (de)compressor option-dict loops so passing a - /// `CompressionParameter` to a `ZstdDecompressor` (or vice versa) is - /// rejected with the CPython-compatible "must not be a XParameter" - /// message. If `set_parameter_types` has not been called yet, this is a - /// no-op. + /// If `key`'s class is the wrong kind of parameter enum for the caller's + /// context (a `CompressionParameter` passed to a decompressor, or vice + /// versa), return a `TypeError`. Used by both (de)compressor option-dict + /// loops so the resulting message names the actual type rather than the + /// generic "out of range" wording. + /// + /// We identify the enums by class name (see [`C_PARAMETER_TYPE_NAME`]), + /// which avoids storing a `PyTypeRef` in a Rust `static` (impossible on + /// single-threaded builds whose `PyMutex` is not `Sync`). fn check_wrong_param_kind( key: &PyObjectRef, is_compress: bool, vm: &VirtualMachine, ) -> PyResult<()> { - let guard = REGISTERED_PARAM_TYPES.lock(); - let (forbidden_type, forbidden_name, kind) = if is_compress { - (&guard.decompression, "DecompressionParameter", "compression") + let class_name = key.class().name(); + let (forbidden, kind) = if is_compress { + (D_PARAMETER_TYPE_NAME, "compression") } else { - (&guard.compression, "CompressionParameter", "decompression") + (C_PARAMETER_TYPE_NAME, "decompression") }; - if let Some(ref t) = *forbidden_type - && key.fast_isinstance(t) - { + if &*class_name == forbidden { return Err(vm.new_type_error(format!( - "{kind} options dictionary key must not be a {forbidden_name}" + "{kind} options dictionary key must not be a {forbidden}" ))); } Ok(()) @@ -1684,29 +1676,30 @@ mod _zstd { /// set_parameter_types(c_parameter_type, d_parameter_type, /) /// - /// Register the `CompressionParameter` and `DecompressionParameter` - /// `IntEnum` types defined in the pure-Python wrapper. Once registered, - /// the (de)compressor constructors validate that callers do not mix the - /// two enum kinds (e.g. passing a `CompressionParameter` key to a - /// `ZstdDecompressor`), raising `TypeError` if they do. + /// Validate that the `CompressionParameter` and `DecompressionParameter` + /// types passed in by the pure-Python wrapper are actual type objects. + /// The (de)compressor constructors use the *names* of these classes + /// (see [`check_wrong_param_kind`]) rather than identity, so this + /// function does not need to retain the type objects beyond the call. /// - /// The Python wrapper at `Lib/compression/zstd/__init__.py:242` calls - /// this exactly once at import time. + /// CPython retains them and uses them for nicer error messages naming + /// the specific enum member; matching that behavior would require a + /// `Sync` cell-of-`PyTypeRef`, which is not available on single-threaded + /// targets. The pure-Python wrapper at + /// `Lib/compression/zstd/__init__.py:242` calls this exactly once at + /// import time. #[pyfunction] fn set_parameter_types( c_parameter_type: PyObjectRef, d_parameter_type: PyObjectRef, vm: &VirtualMachine, ) -> PyResult<()> { - let c_type = c_parameter_type + let _ = c_parameter_type .downcast::() .map_err(|_| vm.new_type_error("c_parameter_type must be a type object"))?; - let d_type = d_parameter_type + let _ = d_parameter_type .downcast::() .map_err(|_| vm.new_type_error("d_parameter_type must be a type object"))?; - let mut guard = REGISTERED_PARAM_TYPES.lock(); - guard.compression = Some(c_type); - guard.decompression = Some(d_type); Ok(()) } } From ee26123e6d73565c5c0b432e2cf23f64db9a82f0 Mon Sep 17 00:00:00 2001 From: James David Clarke Date: Sat, 23 May 2026 18:56:31 +0100 Subject: [PATCH 03/20] Fix clippy errors in zstd module RustPython's clippy lint set rejects `std::` imports for items that also live in `core::` or `alloc::`. Switch every such reference in `zstd.rs`: - `std::ffi::{c_int, CStr}` -> `core::ffi::{c_int, CStr}` - `std::fmt::*` -> `core::fmt::*` - `std::slice::from_raw_parts` -> `core::slice::from_raw_parts` - `std::borrow::Cow` -> `alloc::borrow::Cow` Also clear up other clippy findings under `-D warnings`: - Drop the redundant `obj.clone()` on the last use of `obj` in `parse_zstd_dict_arg`. - Collapse `.map(|n| n.get()).unwrap_or(0)` to `.map_or(0, |n| n.get())` in the two dict-id readers. - Replace `&*work_data` with `&work_data` (auto-deref). - Factor the `(Option, Option>)` return shape into a `DictLoadResult` type alias to satisfy `type_complexity`. - Replace the remaining `std::mem::transmute` / `ZSTD_dParameter` in `get_param_bounds` with the existing safe `c_param_enum` / `d_param_enum` helpers; surfaces a clear "invalid parameter" `ValueError` instead of relying on UB-adjacent transmutes for unknown ints. Verified locally: - `cargo clippy -p rustpython-stdlib --no-deps -- -Dwarnings`: clean (only the pre-existing `socket.rs::sock_wait` warning remains). - `cargo run -- -m test test_zstd`: 119 pass, 1 expected failure. Co-Authored-By: Claude Opus 4.7 (1M context) --- crates/stdlib/src/zstd.rs | 84 +++++++++++++++++++++------------------ 1 file changed, 45 insertions(+), 39 deletions(-) diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs index 7f0ef4990f5..95793ccee5b 100644 --- a/crates/stdlib/src/zstd.rs +++ b/crates/stdlib/src/zstd.rs @@ -33,7 +33,7 @@ mod _zstd { use rustpython_vm::{ AsObject, Context, Py, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, }; - use std::ffi::c_int; + use core::ffi::c_int; use zstd_safe::zstd_sys; use zstd_safe::{CCtx, CParameter, DCtx, DParameter, InBuffer, OutBuffer}; @@ -428,10 +428,13 @@ mod _zstd { obj: PyObjectRef, vm: &VirtualMachine, ) -> PyResult<(PyRef, i32)> { + // The first downcast clones `obj` because we fall through to the + // tuple branch if it fails. The second downcast (the tuple one) is + // the last use of `obj`, so we let it move directly. if let Ok(d) = obj.clone().downcast::() { return Ok((d, DICT_TYPE_DIGESTED)); } - if let Ok(tuple) = obj.clone().downcast::() { + if let Ok(tuple) = obj.downcast::() { let items = tuple.as_slice(); // Reject any tuple shape that is not (ZstdDict, int_marker) so the // test suite's bad-args coverage (`(zd, 1.0)`, `(zd,)`, `(zd, 3)`, @@ -495,8 +498,8 @@ mod _zstd { dict_id: u32, } - impl std::fmt::Debug for ZstdDict { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + impl core::fmt::Debug for ZstdDict { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { write!( f, "", @@ -516,9 +519,8 @@ mod _zstd { // does not carry the dictionary magic. Both are runtime errors // when `is_raw=False`, matching CPython's behavior of raising // `ValueError` on a non-conformant dictionary. - let parsed_id = zstd_safe::get_dict_id_from_dict(&dict_content) - .map(|n| n.get()) - .unwrap_or(0); + let parsed_id = + zstd_safe::get_dict_id_from_dict(&dict_content).map_or(0, |n| n.get()); if !args.is_raw && parsed_id == 0 { return Err(vm.new_value_error( "ZSTD_DICT_MAGIC_NUMBER not found, dict_content cannot be a 'raw content' \ @@ -671,8 +673,8 @@ mod _zstd { state: PyMutex, } - impl std::fmt::Debug for ZstdCompressor { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + impl core::fmt::Debug for ZstdCompressor { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { write!(f, "_zstd.ZstdCompressor") } } @@ -873,6 +875,11 @@ mod _zstd { } } + /// Return value of `load_dict`: the digested `CDict`/`DDict` (if any) + /// and the `PyRef` we hold to keep the dictionary bytes alive + /// while `ref_prefix` may point into them. + type DictLoadResult = PyResult<(Option, Option>)>; + /// Common path for attaching a dictionary to either context type. Returns /// the digested `CDict`/`DDict` (if the caller used digested mode) plus /// the `PyRef` we hold to keep the bytes alive for `ref_prefix`. @@ -880,7 +887,7 @@ mod _zstd { ctx: &mut L, dict_obj: Option, vm: &VirtualMachine, - ) -> PyResult<(Option, Option>)> { + ) -> DictLoadResult { let Some(dict_obj) = dict_obj else { return Ok((None, None)); }; @@ -904,7 +911,7 @@ mod _zstd { // `dict_bytes`; as long as the (de)compressor outlives no // longer than the held PyRef, the pointer stays valid. let static_bytes: &'static [u8] = unsafe { - std::slice::from_raw_parts(dict_bytes.as_ptr(), dict_bytes.len()) + core::slice::from_raw_parts(dict_bytes.as_ptr(), dict_bytes.len()) }; ctx.ref_prefix_static(static_bytes) .map_err(|_| bad_dict_err())?; @@ -930,7 +937,7 @@ mod _zstd { cctx: &mut CCtx<'static>, dict_obj: Option, vm: &VirtualMachine, - ) -> PyResult<(Option>, Option>)> { + ) -> DictLoadResult> { load_dict::>(cctx, dict_obj, vm) } @@ -938,7 +945,7 @@ mod _zstd { dctx: &mut DCtx<'static>, dict_obj: Option, vm: &VirtualMachine, - ) -> PyResult<(Option>, Option>)> { + ) -> DictLoadResult> { load_dict::>(dctx, dict_obj, vm) } @@ -1169,8 +1176,8 @@ mod _zstd { state: PyMutex, } - impl std::fmt::Debug for ZstdDecompressor { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + impl core::fmt::Debug for ZstdDecompressor { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { write!(f, "_zstd.ZstdDecompressor") } } @@ -1239,16 +1246,16 @@ mod _zstd { // Combine any buffered leftover input with the new data so the // decompressor sees one contiguous stream. `Cow` avoids the // allocation when there is no leftover. - let work_data: std::borrow::Cow<'_, [u8]> = if state.input_buffer.is_empty() { - std::borrow::Cow::Borrowed(new_data) + let work_data: alloc::borrow::Cow<'_, [u8]> = if state.input_buffer.is_empty() { + alloc::borrow::Cow::Borrowed(new_data) } else { let mut combined = Vec::with_capacity(state.input_buffer.len() + new_data.len()); combined.extend_from_slice(&state.input_buffer); combined.extend_from_slice(new_data); - std::borrow::Cow::Owned(combined) + alloc::borrow::Cow::Owned(combined) }; - let mut input = InBuffer::around(&*work_data); + let mut input = InBuffer::around(&work_data); let mut output: Vec = Vec::new(); let chunk_size = DCtx::out_size().max(1); // Reusable scratch buffer for each decompress_stream call. We need an @@ -1463,9 +1470,7 @@ mod _zstd { Some(n) => vm.ctx.new_int(n).into(), None => vm.ctx.none(), }; - let dict_id = zstd_safe::get_dict_id_from_frame(&buf) - .map(|n| n.get()) - .unwrap_or(0); + let dict_id = zstd_safe::get_dict_id_from_frame(&buf).map_or(0, |n| n.get()); Ok((content_size_obj, dict_id)) } @@ -1621,7 +1626,7 @@ mod _zstd { let msg = if err_ptr.is_null() { "zstd dictionary finalization failed".to_string() } else { - unsafe { std::ffi::CStr::from_ptr(err_ptr) } + unsafe { core::ffi::CStr::from_ptr(err_ptr) } .to_string_lossy() .into_owned() }; @@ -1648,24 +1653,25 @@ mod _zstd { /// in the pure-Python wrapper. #[pyfunction] fn get_param_bounds(args: ParamBoundsArgs, vm: &VirtualMachine) -> PyResult<(c_int, c_int)> { + let unknown = || -> PyBaseExceptionRef { + let kind = if args.is_compress { "compression" } else { "decompression" }; + vm.new_value_error(format!( + "invalid {kind} parameter 'unknown parameter (key {})'", + args.parameter + )) + }; + // Validate the id via the same safe enum-lookup helpers used in + // `lookup_param_bounds`, then call libzstd directly so we can + // distinguish a libzstd-reported error from our own "unknown". let bounds = if args.is_compress { - // Validate by mapping the int to a CParameter variant first; if - // the int is not a known parameter id we surface a clear error - // instead of letting libzstd return a generic "parameter unsupported". - let _ = cparameter_from_int(args.parameter, 0, vm)?; - // SAFETY: we just confirmed `args.parameter` is a value of the - // `ZSTD_cParameter` enum. - let param_enum = unsafe { - std::mem::transmute::(args.parameter as u32) - }; - unsafe { zstd_sys::ZSTD_cParam_getBounds(param_enum) } + let p = c_param_enum(args.parameter).ok_or_else(unknown)?; + // SAFETY: `c_param_enum` returned `Some`, so `p` is a real + // `ZSTD_cParameter` discriminant. + unsafe { zstd_sys::ZSTD_cParam_getBounds(p) } } else { - let _ = dparameter_from_int(args.parameter, 0, vm)?; - // SAFETY: same justification as the compress branch above. - let param_enum = unsafe { - std::mem::transmute::(args.parameter as u32) - }; - unsafe { zstd_sys::ZSTD_dParam_getBounds(param_enum) } + let p = d_param_enum(args.parameter).ok_or_else(unknown)?; + // SAFETY: same as above. + unsafe { zstd_sys::ZSTD_dParam_getBounds(p) } }; // SAFETY: ZSTD_isError just inspects the integer error code. if unsafe { zstd_sys::ZSTD_isError(bounds.error) } != 0 { From fbce9c996364d0cc7b4ddb91bb01bbb71cf77f94 Mon Sep 17 00:00:00 2001 From: James David Clarke Date: Sat, 23 May 2026 18:59:59 +0100 Subject: [PATCH 04/20] Retrigger CI The previous run's `Lint` job failed with a `401 Unauthorized` from reviewdog hitting the GitHub API (a transient CI/permissions issue, unrelated to this PR's diff). Pushing an empty commit to re-run CI. Co-Authored-By: Claude Opus 4.7 (1M context) From 63103b65c7fe2287b46ebb5220204610100597f9 Mon Sep 17 00:00:00 2001 From: James David Clarke Date: Sat, 23 May 2026 19:09:49 +0100 Subject: [PATCH 05/20] chore: format zstd module {"subject": "Format zstd module and extend spell-check ignores", "body": "- Apply rustfmt formatting to zstd.rs\n- Add additional spell-checker ignore entries for zstd identifiers"} --- crates/stdlib/src/zstd.rs | 40 +++++++++++++++++++++++++-------------- 1 file changed, 26 insertions(+), 14 deletions(-) diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs index 95793ccee5b..dc5a3db73a7 100644 --- a/crates/stdlib/src/zstd.rs +++ b/crates/stdlib/src/zstd.rs @@ -3,6 +3,8 @@ // spell-checker:ignore btlazy dfast nbworkers windowlogmax windowlog overlap targetcblock // spell-checker:ignore srcsize zdict refprefix refcdict refddict pledgedsrcsize getframecontentsize // spell-checker:ignore Zstd Zstandard pylib RFC +// spell-checker:ignore CLEVEL zstdmodule cparameter dparameter maxl +// spell-checker:ignore cctx dctx CCTX DCTX ldm cdict ddict windowlog hashlog chainlog searchlog CLEVEL //! The `_zstd` extension module. Backs the pure-Python `compression.zstd` //! package by exposing the same classes, functions and constants that @@ -24,6 +26,7 @@ pub(crate) use _zstd::module_def; #[allow(non_upper_case_globals)] #[pymodule] mod _zstd { + use core::ffi::c_int; use rustpython_common::lock::PyMutex; use rustpython_vm::builtins::{ PyBaseExceptionRef, PyBytesRef, PyDict, PyTupleRef, PyType, PyTypeRef, @@ -33,7 +36,6 @@ mod _zstd { use rustpython_vm::{ AsObject, Context, Py, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, }; - use core::ffi::c_int; use zstd_safe::zstd_sys; use zstd_safe::{CCtx, CParameter, DCtx, DParameter, InBuffer, OutBuffer}; @@ -279,9 +281,11 @@ mod _zstd { ZSTD_c_searchLog => CParameter::SearchLog(value as u32), ZSTD_c_minMatch => CParameter::MinMatch(value as u32), ZSTD_c_targetLength => CParameter::TargetLength(value as u32), - ZSTD_c_strategy => CParameter::Strategy(strategy_from_int(value).ok_or_else( - || new_zstd_error(format!("invalid strategy value: {value}"), vm), - )?), + ZSTD_c_strategy => { + CParameter::Strategy(strategy_from_int(value).ok_or_else(|| { + new_zstd_error(format!("invalid strategy value: {value}"), vm) + })?) + } ZSTD_c_enableLongDistanceMatching => CParameter::EnableLongDistanceMatching(value != 0), ZSTD_c_ldmHashLog => CParameter::LdmHashLog(value as u32), ZSTD_c_ldmMinMatch => CParameter::LdmMinMatch(value as u32), @@ -323,7 +327,11 @@ mod _zstd { is_compress: bool, vm: &VirtualMachine, ) -> PyBaseExceptionRef { - let kind = if is_compress { "compression" } else { "decompression" }; + let kind = if is_compress { + "compression" + } else { + "decompression" + }; let name = parameter_name(param, is_compress); match lookup_param_bounds(param, is_compress) { Some((lo, hi)) => vm.new_value_error(format!( @@ -519,8 +527,7 @@ mod _zstd { // does not carry the dictionary magic. Both are runtime errors // when `is_raw=False`, matching CPython's behavior of raising // `ValueError` on a non-conformant dictionary. - let parsed_id = - zstd_safe::get_dict_id_from_dict(&dict_content).map_or(0, |n| n.get()); + let parsed_id = zstd_safe::get_dict_id_from_dict(&dict_content).map_or(0, |n| n.get()); if !args.is_raw && parsed_id == 0 { return Err(vm.new_value_error( "ZSTD_DICT_MAGIC_NUMBER not found, dict_content cannot be a 'raw content' \ @@ -910,9 +917,8 @@ mod _zstd { // context. libzstd `ref_prefix` stores a raw pointer into // `dict_bytes`; as long as the (de)compressor outlives no // longer than the held PyRef, the pointer stays valid. - let static_bytes: &'static [u8] = unsafe { - core::slice::from_raw_parts(dict_bytes.as_ptr(), dict_bytes.len()) - }; + let static_bytes: &'static [u8] = + unsafe { core::slice::from_raw_parts(dict_bytes.as_ptr(), dict_bytes.len()) }; ctx.ref_prefix_static(static_bytes) .map_err(|_| bad_dict_err())?; } @@ -927,7 +933,8 @@ mod _zstd { _ => { // Undigested: copy the bytes into the context. Validation // happens lazily at the first stream call in this mode. - ctx.load_undigested(dict_bytes).map_err(|_| bad_dict_err())?; + ctx.load_undigested(dict_bytes) + .map_err(|_| bad_dict_err())?; } } Ok((digested, Some(zdict))) @@ -949,7 +956,6 @@ mod _zstd { load_dict::>(dctx, dict_obj, vm) } - /// Drive `compress_stream2` until the input is fully consumed and, for /// flush/end directives, the internal buffers report zero remaining bytes. /// Grows the output `Vec` by `CStreamOutSize` chunks as needed. @@ -969,7 +975,9 @@ mod _zstd { output.reserve(chunk_size); let remaining = { let mut out_buf = OutBuffer::around_pos(&mut output, prev_len); - state.cctx.compress_stream2(&mut out_buf, &mut input, end_op) + state + .cctx + .compress_stream2(&mut out_buf, &mut input, end_op) } .map_err(|c| catch_zstd_error(c, vm))?; let consumed_all = input.pos == input.src.len(); @@ -1654,7 +1662,11 @@ mod _zstd { #[pyfunction] fn get_param_bounds(args: ParamBoundsArgs, vm: &VirtualMachine) -> PyResult<(c_int, c_int)> { let unknown = || -> PyBaseExceptionRef { - let kind = if args.is_compress { "compression" } else { "decompression" }; + let kind = if args.is_compress { + "compression" + } else { + "decompression" + }; vm.new_value_error(format!( "invalid {kind} parameter 'unknown parameter (key {})'", args.parameter From 77370cb79bafdf610690ddbad31a8752194f4588 Mon Sep 17 00:00:00 2001 From: James Clarke Date: Sun, 24 May 2026 20:09:51 +0100 Subject: [PATCH 06/20] Update crates/stdlib/src/zstd.rs Co-authored-by: Shahar Naveh <50263213+ShaharNaveh@users.noreply.github.com> --- crates/stdlib/src/zstd.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs index dc5a3db73a7..339e9e3aa6c 100644 --- a/crates/stdlib/src/zstd.rs +++ b/crates/stdlib/src/zstd.rs @@ -1512,7 +1512,7 @@ mod _zstd { for item in items { let idx = item.try_index(vm)?; let v: usize = idx.try_to_primitive(vm).map_err(|_| { - vm.new_value_error("sample size out of range for size_t".to_owned()) + vm.new_value_error("sample size out of range for size_t") })?; out.push(v); } From ace862fc92dcd5e61bd4d9973d07992369fac1f7 Mon Sep 17 00:00:00 2001 From: James David Clarke Date: Sun, 24 May 2026 20:51:35 +0100 Subject: [PATCH 07/20] chore: remove rustdoc comments from zstd module {"subject": "Remove rustdoc comments from zstd module", "body": "- Strip doc comments from internal items in _zstd module\n- Convert a few API-doc comments into regular implementation notes where appropriate"} --- crates/stdlib/src/zstd.rs | 203 +++----------------------------------- 1 file changed, 13 insertions(+), 190 deletions(-) diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs index 339e9e3aa6c..d6a16efb6dd 100644 --- a/crates/stdlib/src/zstd.rs +++ b/crates/stdlib/src/zstd.rs @@ -53,8 +53,6 @@ mod _zstd { // Module-level constants // ========================================================================= - /// Default compression level used when none is requested. Matches - /// libzstd's `ZSTD_CLEVEL_DEFAULT`. #[pyattr] const ZSTD_CLEVEL_DEFAULT: i32 = zstd_sys::ZSTD_CLEVEL_DEFAULT as i32; @@ -127,21 +125,16 @@ mod _zstd { #[pyattr] const ZSTD_btultra2: i32 = zstd_sys::ZSTD_strategy::ZSTD_btultra2 as i32; - /// Runtime version string of the linked libzstd, like `"1.5.7"`. #[pyattr(once, name = "zstd_version")] fn zstd_version(_vm: &VirtualMachine) -> String { zstd_safe::version_string().to_string() } - /// Packed integer version: `MAJOR * 10000 + MINOR * 100 + RELEASE`. - /// The Python wrapper unpacks this into a `zstd_version_info` tuple. #[pyattr(once, name = "zstd_version_number")] fn zstd_version_number(_vm: &VirtualMachine) -> u32 { zstd_safe::version_number() } - /// Recommended output buffer size for decompression. The `ZstdFile.read1` - /// wrapper uses this as its default read size. #[pyattr(once, name = "ZSTD_DStreamOutSize")] fn zstd_dstream_out_size(_vm: &VirtualMachine) -> usize { DCtx::out_size() @@ -159,8 +152,6 @@ mod _zstd { // ZstdError exception // ========================================================================= - /// `_zstd.ZstdError`: the exception type raised for libzstd errors and for - /// invalid usage (calling `decompress` on a closed decompressor, etc.). #[pyattr(once, name = "ZstdError")] fn zstd_error(vm: &VirtualMachine) -> PyTypeRef { vm.ctx.new_exception_type( @@ -487,17 +478,6 @@ mod _zstd { is_raw: bool, } - /// ZstdDict(dict_content, /, *, is_raw=False) - /// - /// Represents a Zstandard dictionary. Hold the raw dictionary bytes and - /// the parsed `dict_id`. The same instance may be shared by many - /// compressors and decompressors. - /// - /// *dict_content* is a bytes-like object containing the dictionary - /// content. If *is_raw* is False (the default), the content is expected - /// to start with the Zstandard dictionary magic number and have a valid - /// header. If *is_raw* is True, arbitrary bytes are accepted and the - /// dictionary is treated as raw back-reference content. #[pyattr] #[pyclass(name = "ZstdDict")] #[derive(PyPayload)] @@ -571,41 +551,22 @@ mod _zstd { #[pyclass(with(Constructor, Representable, AsMapping))] impl ZstdDict { - /// The content of a Zstandard dictionary, as a bytes object. #[pygetset] fn dict_content(&self) -> PyBytesRef { self.dict_content.clone() } - /// The Zstandard dictionary ID. - /// - /// A non-zero value identifies a standard-format dictionary. A value - /// of zero either means the dictionary is in raw-content form or - /// that the dictionary ID was deliberately omitted from the header. - /// Dictionary IDs fall in the range 0 to 2**32 - 1. #[pygetset] fn dict_id(&self) -> u32 { self.dict_id } - /// Load this ZstdDict as a digested dictionary. - /// - /// Returns a `(self, DICT_TYPE_DIGESTED)` tuple that can be passed - /// as the `zstd_dict=` argument of `ZstdCompressor` or - /// `ZstdDecompressor`. Digested form lets libzstd cache the parsed - /// dictionary for reuse across compression contexts. #[pygetset] fn as_digested_dict(zelf: PyRef, vm: &VirtualMachine) -> PyTupleRef { vm.ctx .new_tuple(vec![zelf.into(), vm.ctx.new_int(DICT_TYPE_DIGESTED).into()]) } - /// Load this ZstdDict as an undigested dictionary. - /// - /// Returns a `(self, DICT_TYPE_UNDIGESTED)` tuple. In undigested - /// form the raw dictionary bytes are copied into each - /// (de)compressor, which is slower than digested but preserves the - /// original dictionary parameters across loads. #[pygetset] fn as_undigested_dict(zelf: PyRef, vm: &VirtualMachine) -> PyTupleRef { vm.ctx.new_tuple(vec![ @@ -614,11 +575,6 @@ mod _zstd { ]) } - /// Load this ZstdDict as a prefix dictionary. - /// - /// Returns a `(self, DICT_TYPE_PREFIX)` tuple. Prefix loading treats - /// the dictionary bytes as a one-shot back-reference window for the - /// next frame only; subsequent frames will not see the prefix. #[pygetset] fn as_prefix(zelf: PyRef, vm: &VirtualMachine) -> PyTupleRef { vm.ctx @@ -655,24 +611,6 @@ mod _zstd { last_mode: i32, } - /// ZstdCompressor(level=None, options=None, zstd_dict=None) - /// - /// Create a compressor object for incremental compression of data into - /// the Zstandard format. - /// - /// *level* is an int specifying the compression level. Valid range is - /// roughly [-(1<<17), 22]; 0 means "use the default level". Cannot be - /// combined with *options*. - /// - /// *options* is a dict mapping `CompressionParameter` enum members to - /// integer values for fine-grained control. Cannot be combined with - /// *level*. - /// - /// *zstd_dict* is an optional `ZstdDict` (or one of its `as_*` - /// variants) for dictionary-assisted compression. - /// - /// Use `compress()` to feed data into the compressor and `flush()` to - /// finalize the current frame. #[pyattr] #[pyclass(name = "ZstdCompressor")] #[derive(PyPayload)] @@ -994,20 +932,6 @@ mod _zstd { #[pyclass(with(Constructor))] impl ZstdCompressor { - /// compress(data, mode=ZstdCompressor.CONTINUE) - /// - /// Provide data to the compressor and return a chunk of compressed - /// output bytes (which may be empty if not enough data has been - /// buffered yet). - /// - /// *mode* controls when to flush the output: - /// * `CONTINUE` (default) keeps buffering, producing the densest - /// compression. The compressor may return an empty byte string. - /// * `FLUSH_BLOCK` closes the current zstd block so all data - /// provided so far can be immediately decompressed. Frequent - /// use reduces the compression ratio. - /// * `FLUSH_FRAME` closes the current frame entirely. The next - /// `compress()` call starts a new frame. #[pymethod] fn compress(&self, args: CompressMethodArgs, vm: &VirtualMachine) -> PyResult> { let mode = args.mode.unwrap_or(COMP_MODE_CONTINUE); @@ -1019,13 +943,6 @@ mod _zstd { Ok(out) } - /// flush(mode=ZstdCompressor.FLUSH_FRAME) - /// - /// Finish the compression process and return all the remaining - /// compressed bytes. *mode* must be one of `FLUSH_FRAME` (close the - /// current frame; new frames may be started afterward) or - /// `FLUSH_BLOCK` (close the current block but stay inside the - /// frame). `CONTINUE` is not a valid value for `flush()`. #[pymethod] fn flush(&self, args: FlushMethodArgs, vm: &VirtualMachine) -> PyResult> { let mode = args.mode.unwrap_or(COMP_MODE_FLUSH_FRAME); @@ -1043,16 +960,6 @@ mod _zstd { Ok(out) } - /// set_pledged_input_size(size, /) - /// - /// Set the uncompressed content size to be written into the frame - /// header of the next frame. Useful when the size is known up front: - /// libzstd writes it into the header (unless the content-size flag - /// is disabled) and verifies that the actual input matches at frame - /// end. Pass `None` to mean "unknown". - /// - /// May only be called when `last_mode == FLUSH_FRAME` (i.e. between - /// frames; not in the middle of one). #[pymethod] fn set_pledged_input_size(&self, size: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { let mut state = self.state.lock(); @@ -1093,13 +1000,6 @@ mod _zstd { Ok(()) } - /// The mode argument passed to the most recent `compress()` or - /// `flush()` call. - /// - /// The value is one of `CONTINUE`, `FLUSH_BLOCK`, or `FLUSH_FRAME`. - /// The initial value is `FLUSH_FRAME`, which signals that no frame - /// is currently in progress (it is safe to call - /// `set_pledged_input_size()` in this state). #[pygetset] fn last_mode(&self) -> i32 { self.state.lock().last_mode @@ -1164,19 +1064,6 @@ mod _zstd { input_buffer: Vec, } - /// ZstdDecompressor(zstd_dict=None, options=None) - /// - /// Create a decompressor object for incremental decompression of a - /// single Zstandard frame. Once the frame is fully decompressed the - /// decompressor is at EOF; trailing bytes are exposed via - /// `unused_data`. For multi-frame input use the module-level - /// `compression.zstd.decompress()` or `ZstdFile`. - /// - /// *zstd_dict* is an optional `ZstdDict` (or one of its `as_*` - /// variants) matching the dictionary used during compression. - /// - /// *options* is a dict mapping `DecompressionParameter` enum members - /// to integer values for fine-grained control. #[pyattr] #[pyclass(name = "ZstdDecompressor")] #[derive(PyPayload)] @@ -1372,19 +1259,6 @@ mod _zstd { #[pyclass(with(Constructor))] impl ZstdDecompressor { - /// decompress(data, max_length=-1) - /// - /// Decompress *data* and return at most *max_length* bytes of - /// uncompressed output. A *max_length* of -1 (the default) means - /// "no limit". After the current frame is fully decompressed, `eof` - /// becomes `True` and any further `decompress()` calls raise - /// `EOFError`; bytes that arrived after the frame go into - /// `unused_data` for the caller to forward. - /// - /// If the output cap is reached with input still pending, the - /// remaining input is buffered internally and `needs_input` is set - /// to `False`, signalling that the next call can be made with an - /// empty *data* argument to keep draining. #[pymethod] fn decompress(&self, args: DecompressMethodArgs, vm: &VirtualMachine) -> PyResult> { let data_vec = args.data.with_ref(|b| b.to_vec()); @@ -1403,26 +1277,16 @@ mod _zstd { do_decompress(&mut state, &data_vec, max_length, vm) } - /// True if the end-of-frame marker has been reached. #[pygetset] fn eof(&self) -> bool { self.state.lock().eof } - /// True when the decompressor expects more compressed input. - /// - /// Becomes False after a `decompress()` call that hits its - /// `max_length` cap with input still buffered internally; in that - /// state the caller should call `decompress(b'', max_length=...)` - /// to drain the buffered output before providing more input. #[pygetset] fn needs_input(&self) -> bool { self.state.lock().needs_input } - /// Bytes left over after the end-of-frame marker. - /// - /// Empty until the decompressor reaches EOF. Read-only. #[pygetset] fn unused_data(&self) -> PyBytesRef { self.state.lock().unused_data.clone() @@ -1433,12 +1297,6 @@ mod _zstd { // Module-level functions // ========================================================================= - /// get_frame_size(frame_buffer, /) - /// - /// Return the size in bytes of the first complete Zstandard frame in - /// *frame_buffer* (which must contain the entire frame, header through - /// trailer). Raises `ZstdError` if *frame_buffer* is shorter than the - /// complete frame. #[pyfunction] fn get_frame_size(frame_buffer: ArgBytesLike, vm: &VirtualMachine) -> PyResult { let buf = frame_buffer.with_ref(|b| b.to_vec()); @@ -1452,14 +1310,6 @@ mod _zstd { }) } - /// get_frame_info(frame_buffer, /) - /// - /// Return `(decompressed_size, dictionary_id)` for the first Zstandard - /// frame in *frame_buffer*. `decompressed_size` is `None` when the - /// frame header does not record it (small or streaming frames). - /// `dictionary_id` is 0 if no dictionary was used. *frame_buffer* must - /// contain at least the 6-18 byte frame header. The pure-Python - /// wrapper boxes this tuple into a `FrameInfo` named tuple. #[pyfunction] fn get_frame_info( frame_buffer: ArgBytesLike, @@ -1511,9 +1361,9 @@ mod _zstd { let mut out = Vec::with_capacity(items.len()); for item in items { let idx = item.try_index(vm)?; - let v: usize = idx.try_to_primitive(vm).map_err(|_| { - vm.new_value_error("sample size out of range for size_t") - })?; + let v: usize = idx + .try_to_primitive(vm) + .map_err(|_| vm.new_value_error("sample size out of range for size_t"))?; out.push(v); } Ok(out) @@ -1536,16 +1386,6 @@ mod _zstd { Ok(v as usize) } - /// train_dict(samples_bytes, samples_sizes, dict_size, /) - /// - /// Train a Zstandard dictionary from sample data and return the raw - /// dictionary bytes. The pure-Python `train_dict` wrapper then wraps - /// the result in a `ZstdDict`. - /// - /// `samples_bytes` is the concatenation of all sample contents. - /// `samples_sizes` is a tuple giving the length of each sample, in the - /// same order they appear in `samples_bytes`. `dict_size` is the - /// maximum size of the dictionary in bytes. #[pyfunction] fn train_dict(args: TrainDictArgs, vm: &VirtualMachine) -> PyResult> { let dict_size = parse_positive_dict_size(&args.dict_size, vm)?; @@ -1584,12 +1424,8 @@ mod _zstd { compression_level: PyObjectRef, } - /// finalize_dict(custom_dict_bytes, samples_bytes, samples_sizes, dict_size, compression_level, /) - /// - /// Take a hand-crafted custom dictionary plus a set of samples and - /// produce a standard-format Zstandard dictionary by appending the - /// header and statistics tables. `zstd_safe` does not wrap this - /// function, so we drop down to raw `zstd_sys` FFI for it. + // `zstd_safe` does not wrap `ZDICT_finalizeDictionary`, so we drop down to + // raw `zstd_sys` FFI for it. #[pyfunction] fn finalize_dict(args: FinalizeDictArgs, vm: &VirtualMachine) -> PyResult> { let dict_size = parse_positive_dict_size(&args.dict_size, vm)?; @@ -1652,13 +1488,6 @@ mod _zstd { is_compress: bool, } - /// get_param_bounds(parameter, /, *, is_compress) - /// - /// Return the inclusive `(lower, upper)` integer bounds for a - /// compression parameter (when *is_compress* is True) or a - /// decompression parameter (when False). Used by - /// `CompressionParameter.bounds()` and `DecompressionParameter.bounds()` - /// in the pure-Python wrapper. #[pyfunction] fn get_param_bounds(args: ParamBoundsArgs, vm: &VirtualMachine) -> PyResult<(c_int, c_int)> { let unknown = || -> PyBaseExceptionRef { @@ -1692,20 +1521,14 @@ mod _zstd { Ok((bounds.lowerBound, bounds.upperBound)) } - /// set_parameter_types(c_parameter_type, d_parameter_type, /) - /// - /// Validate that the `CompressionParameter` and `DecompressionParameter` - /// types passed in by the pure-Python wrapper are actual type objects. - /// The (de)compressor constructors use the *names* of these classes - /// (see [`check_wrong_param_kind`]) rather than identity, so this - /// function does not need to retain the type objects beyond the call. - /// - /// CPython retains them and uses them for nicer error messages naming - /// the specific enum member; matching that behavior would require a - /// `Sync` cell-of-`PyTypeRef`, which is not available on single-threaded - /// targets. The pure-Python wrapper at - /// `Lib/compression/zstd/__init__.py:242` calls this exactly once at - /// import time. + // The (de)compressor constructors identify the parameter enums by class + // *name* (see [`check_wrong_param_kind`]) rather than identity, so this + // function does not need to retain the type objects beyond the call. + // CPython retains them to produce nicer error messages naming the specific + // enum member; matching that would require a `Sync` cell-of-`PyTypeRef`, + // which is not available on single-threaded targets. The pure-Python + // wrapper at `Lib/compression/zstd/__init__.py:242` calls this exactly + // once at import time. #[pyfunction] fn set_parameter_types( c_parameter_type: PyObjectRef, From 07aa356e084bc55f77b6fb5217cd34f5f52a3ad1 Mon Sep 17 00:00:00 2001 From: James David Clarke Date: Mon, 25 May 2026 09:28:38 +0100 Subject: [PATCH 08/20] commit message {"subject": "Mark zstd load_dict as unsafe and document invariant", "body": "- Add # Safety docs requiring PyRef to outlive the context\n- Wrap calls in load_compressor_dict and load_decompressor_dict with SAFETY comments"} --- crates/stdlib/src/zstd.rs | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs index d6a16efb6dd..de220ecce6a 100644 --- a/crates/stdlib/src/zstd.rs +++ b/crates/stdlib/src/zstd.rs @@ -828,7 +828,13 @@ mod _zstd { /// Common path for attaching a dictionary to either context type. Returns /// the digested `CDict`/`DDict` (if the caller used digested mode) plus /// the `PyRef` we hold to keep the bytes alive for `ref_prefix`. - fn load_dict>( + /// # Safety + /// + /// The caller must ensure that the returned `PyRef` (if any) + /// outlives the context `ctx`. In `ref_prefix` mode, libzstd stores a + /// raw pointer into the dictionary bytes; the `PyRef` keeps those bytes + /// alive for the required lifetime. + unsafe fn load_dict>( ctx: &mut L, dict_obj: Option, vm: &VirtualMachine, @@ -883,7 +889,9 @@ mod _zstd { dict_obj: Option, vm: &VirtualMachine, ) -> DictLoadResult> { - load_dict::>(cctx, dict_obj, vm) + // SAFETY: The returned `PyRef` is stored alongside the + // `CCtx` in `CompressorState`, so it outlives the context. + unsafe { load_dict::>(cctx, dict_obj, vm) } } fn load_decompressor_dict( @@ -891,7 +899,9 @@ mod _zstd { dict_obj: Option, vm: &VirtualMachine, ) -> DictLoadResult> { - load_dict::>(dctx, dict_obj, vm) + // SAFETY: The returned `PyRef` is stored alongside the + // `DCtx` in `DecompressorState`, so it outlives the context. + unsafe { load_dict::>(dctx, dict_obj, vm) } } /// Drive `compress_stream2` until the input is fully consumed and, for From f7598d6091742826d0805f7bf014b99ca5034dc8 Mon Sep 17 00:00:00 2001 From: James David Clarke Date: Mon, 25 May 2026 12:02:48 +0100 Subject: [PATCH 09/20] commit message { "subject": "refactor(zstd): release GIL during (de)compress and tighten dict-load safety", "body": "- Wrap compress/decompress loops in vm.allow_threads to release the GIL\n- Replace load_*_dict helpers with build_*_state that assemble the full state, making the load_dict safety invariants structural\n- Switch constructor args from OptionalArg to OptionalOption and use try_to_value instead of the pyobj_to_i32/arg_or_none helpers\n- Add check_sample_sizes_match with overflow-safe summing for train_dict/finalize_dict" } --- crates/stdlib/src/zstd.rs | 375 +++++++++++++++++++++----------------- 1 file changed, 208 insertions(+), 167 deletions(-) diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs index de220ecce6a..58f5f807996 100644 --- a/crates/stdlib/src/zstd.rs +++ b/crates/stdlib/src/zstd.rs @@ -31,7 +31,7 @@ mod _zstd { use rustpython_vm::builtins::{ PyBaseExceptionRef, PyBytesRef, PyDict, PyTupleRef, PyType, PyTypeRef, }; - use rustpython_vm::function::{ArgBytesLike, OptionalArg}; + use rustpython_vm::function::{ArgBytesLike, OptionalOption}; use rustpython_vm::types::{AsMapping, Constructor, Representable}; use rustpython_vm::{ AsObject, Context, Py, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, @@ -177,21 +177,6 @@ mod _zstd { // Parameter helpers // ========================================================================= - /// Convert any int-like Python value to an `i32`, rejecting floats. CPython - /// uses `PyLong_AsInt` here which only accepts integers; `try_index` is - /// the equivalent of CPython's `__index__` slot lookup (rejects floats). - fn pyobj_to_i32(obj: &PyObjectRef, vm: &VirtualMachine) -> PyResult { - obj.try_index(vm)?.try_to_primitive(vm) - } - - /// Normalize an `OptionalArg` constructor parameter to an - /// `Option`. Treats both "argument missing" and "argument - /// explicitly None" the same way (consistent with CPython, which uses - /// `None` defaults throughout the `_zstd` API). - fn arg_or_none(arg: OptionalArg, vm: &VirtualMachine) -> Option { - arg.into_option().filter(|o| !vm.is_none(o)) - } - /// If `key`'s class is the wrong kind of parameter enum for the caller's /// context (a `CompressionParameter` passed to a decompressor, or vice /// versa), return a `TypeError`. Used by both (de)compressor option-dict @@ -449,7 +434,7 @@ mod _zstd { // The marker must be a plain int (not float/etc); overflow on // `2**1000` propagates as OverflowError via `try_index`. let marker_obj = &items[1]; - let marker = pyobj_to_i32(marker_obj, vm).map_err(|e| { + let marker: i32 = marker_obj.try_to_value(vm).map_err(|e| { // Preserve OverflowError; everything else becomes TypeError so // callers see a consistent "should be a ZstdDict" message. if e.fast_isinstance(vm.ctx.exceptions.overflow_error) { @@ -627,11 +612,11 @@ mod _zstd { #[derive(FromArgs)] pub(super) struct ZstdCompressorArgs { #[pyarg(any, optional)] - level: OptionalArg, + level: OptionalOption, #[pyarg(any, optional)] - options: OptionalArg, + options: OptionalOption, #[pyarg(any, optional)] - zstd_dict: OptionalArg, + zstd_dict: OptionalOption, } /// Translate the public `mode` int to the libzstd `ZSTD_EndDirective` @@ -657,9 +642,9 @@ mod _zstd { type Args = ZstdCompressorArgs; fn py_new(_cls: &Py, args: Self::Args, vm: &VirtualMachine) -> PyResult { - let level_opt = arg_or_none(args.level, vm); - let options_opt = arg_or_none(args.options, vm); - let dict_opt = arg_or_none(args.zstd_dict, vm); + let level_opt = args.level.flatten(); + let options_opt = args.options.flatten(); + let dict_opt = args.zstd_dict.flatten(); if level_opt.is_some() && options_opt.is_some() { return Err(vm.new_runtime_error("Only one of level or options should be used")); @@ -677,15 +662,9 @@ mod _zstd { apply_options(&mut cctx, options_obj, true, vm)?; } - let (held_cdict, held_dict) = load_compressor_dict(&mut cctx, dict_opt, vm)?; - + let state = build_compressor_state(cctx, dict_opt, COMP_MODE_FLUSH_FRAME, vm)?; Ok(Self { - state: PyMutex::new(CompressorState { - cctx, - _cdict: held_cdict, - _dict: held_dict, - last_mode: COMP_MODE_FLUSH_FRAME, - }), + state: PyMutex::new(state), }) } } @@ -697,7 +676,7 @@ mod _zstd { /// conversions so the constructor stays linear. fn parse_compression_level(obj: &PyObjectRef, vm: &VirtualMachine) -> PyResult { let (lo, hi) = level_bounds(); - let level = pyobj_to_i32(obj, vm).map_err(|e| { + let level: i32 = obj.try_to_value(vm).map_err(|e| { if e.fast_isinstance(vm.ctx.exceptions.overflow_error) { vm.new_value_error(format!( "illegal compression level; the valid range is [{lo}, {hi}]" @@ -732,8 +711,8 @@ mod _zstd { // so the error carries the documented type name rather than a // generic out-of-range message. check_wrong_param_kind(&k, is_compress, vm)?; - let key_int = pyobj_to_i32(&k, vm)?; - let val_int = pyobj_to_i32(&v, vm)?; + let key_int: i32 = k.try_to_value(vm)?; + let val_int: i32 = v.try_to_value(vm)?; // libzstd silently clamps out-of-range values for some // parameters (notably compression_level) rather than rejecting // them, so validate against the documented bounds upfront. @@ -827,13 +806,23 @@ mod _zstd { /// Common path for attaching a dictionary to either context type. Returns /// the digested `CDict`/`DDict` (if the caller used digested mode) plus - /// the `PyRef` we hold to keep the bytes alive for `ref_prefix`. + /// the `PyRef` whose bytes libzstd's `ref_prefix` may point into. + /// /// # Safety /// - /// The caller must ensure that the returned `PyRef` (if any) - /// outlives the context `ctx`. In `ref_prefix` mode, libzstd stores a - /// raw pointer into the dictionary bytes; the `PyRef` keeps those bytes - /// alive for the required lifetime. + /// libzstd stores the dictionary as a raw pointer that bypasses Rust's + /// lifetime tracking. The caller must keep both returned values alive at + /// least as long as `ctx`: + /// + /// - In `digested` mode, `ctx` holds a raw pointer to the returned + /// `L::Digested`; dropping it before `ctx` is use-after-free. + /// - In `prefix` mode, `ctx` holds a raw pointer into the bytes owned by + /// the returned `PyRef`; dropping the `PyRef` before `ctx` + /// is use-after-free. + /// + /// In `undigested` mode the bytes are copied into `ctx`, so neither + /// return value carries a safety obligation — but the caller cannot tell + /// the modes apart, so it must keep both alive regardless. unsafe fn load_dict>( ctx: &mut L, dict_obj: Option, @@ -856,11 +845,13 @@ mod _zstd { let mut digested = None; match marker { DICT_TYPE_PREFIX => { - // SAFETY: the returned `PyRef` keeps `zdict` alive - // (and therefore `dict_bytes`) for the lifetime of the - // context. libzstd `ref_prefix` stores a raw pointer into - // `dict_bytes`; as long as the (de)compressor outlives no - // longer than the held PyRef, the pointer stays valid. + // SAFETY: we extend `dict_bytes`' lifetime to `'static` only + // to thread it through `ref_prefix_static`'s signature. + // `load_dict`'s own safety contract requires the caller to + // keep the returned `PyRef` (which owns + // `dict_bytes`) alive at least as long as `ctx`, so the + // raw pointer libzstd stores stays valid for the required + // window. let static_bytes: &'static [u8] = unsafe { core::slice::from_raw_parts(dict_bytes.as_ptr(), dict_bytes.len()) }; ctx.ref_prefix_static(static_bytes) @@ -884,24 +875,53 @@ mod _zstd { Ok((digested, Some(zdict))) } - fn load_compressor_dict( - cctx: &mut CCtx<'static>, + /// Build a fully-initialized `CompressorState` from a freshly-created + /// `CCtx` and an optional dictionary argument. This is the safe interface + /// that `unsafe fn load_dict` was waiting for: by assembling the struct + /// here, both invariants `load_dict` documents become structural and a + /// safe-Rust caller cannot split the pieces apart. + fn build_compressor_state( + mut cctx: CCtx<'static>, dict_obj: Option, + last_mode: i32, vm: &VirtualMachine, - ) -> DictLoadResult> { - // SAFETY: The returned `PyRef` is stored alongside the - // `CCtx` in `CompressorState`, so it outlives the context. - unsafe { load_dict::>(cctx, dict_obj, vm) } + ) -> PyResult { + // SAFETY: `load_dict` requires its two return values to outlive `ctx`. + // We satisfy that by moving `cctx` and both return values into + // `CompressorState` in one expression — Rust drops the struct's + // fields in declaration order, so on teardown `cctx` is dropped + // first, releasing its raw pointers before `_cdict` (digested mode) + // and `_dict` (prefix mode) are freed. `CompressorState` is private + // to this module and is never destructured, so no safe caller can + // reorder the drops. + let (cdict, dict) = unsafe { load_dict::>(&mut cctx, dict_obj, vm) }?; + Ok(CompressorState { + cctx, + _cdict: cdict, + _dict: dict, + last_mode, + }) } - fn load_decompressor_dict( - dctx: &mut DCtx<'static>, + /// Build a fully-initialized `DecompressorState`. See + /// [`build_compressor_state`] for the safety reasoning; + /// `DecompressorState`'s field order plays the same role here. + fn build_decompressor_state( + mut dctx: DCtx<'static>, dict_obj: Option, vm: &VirtualMachine, - ) -> DictLoadResult> { - // SAFETY: The returned `PyRef` is stored alongside the - // `DCtx` in `DecompressorState`, so it outlives the context. - unsafe { load_dict::>(dctx, dict_obj, vm) } + ) -> PyResult { + // SAFETY: see [`build_compressor_state`]. + let (ddict, dict) = unsafe { load_dict::>(&mut dctx, dict_obj, vm) }?; + Ok(DecompressorState { + dctx, + _ddict: ddict, + _dict: dict, + eof: false, + needs_input: true, + unused_data: vm.ctx.empty_bytes.clone(), + input_buffer: Vec::new(), + }) } /// Drive `compress_stream2` until the input is fully consumed and, for @@ -913,31 +933,39 @@ mod _zstd { end_op: zstd_sys::ZSTD_EndDirective, vm: &VirtualMachine, ) -> PyResult> { - let mut output = Vec::new(); - let mut input = InBuffer::around(data); - let chunk_size = CCtx::out_size().max(1); + // Release the GIL for the duration of the compression loop. Safety: + // `data` is an immutable borrow of a local `Vec` in the caller, + // `state` is held under the compressor's `PyMutex` (no other Python + // thread can touch it), the `_dict` bytes referenced by libzstd are + // an immutable `Vec` inside a `PyRef` (other readers fine), + // and the output `Vec` is local to this function. No Python object + // access happens inside the closure — error codes are surfaced as + // `usize` and converted into exceptions after re-attaching. let is_end = end_op != zstd_sys::ZSTD_EndDirective::ZSTD_e_continue; - - loop { - let prev_len = output.len(); - output.reserve(chunk_size); - let remaining = { - let mut out_buf = OutBuffer::around_pos(&mut output, prev_len); - state - .cctx - .compress_stream2(&mut out_buf, &mut input, end_op) - } - .map_err(|c| catch_zstd_error(c, vm))?; - let consumed_all = input.pos == input.src.len(); - // Stop when input is fully consumed and, for flush/end directives, - // libzstd reports that all internal buffers have been drained - // (remaining == 0). Otherwise loop; the next `reserve` will grow - // the output buffer if we hit the previous capacity. - if consumed_all && (!is_end || remaining == 0) { - break; + let chunk_size = CCtx::out_size().max(1); + let result: Result, usize> = vm.allow_threads(|| { + let mut output = Vec::new(); + let mut input = InBuffer::around(data); + loop { + let prev_len = output.len(); + output.reserve(chunk_size); + let remaining = { + let mut out_buf = OutBuffer::around_pos(&mut output, prev_len); + state + .cctx + .compress_stream2(&mut out_buf, &mut input, end_op) + }?; + let consumed_all = input.pos == input.src.len(); + // Stop when input is fully consumed and, for flush/end + // directives, libzstd reports that all internal buffers have + // been drained (remaining == 0). Otherwise loop; the next + // `reserve` will grow the output if we hit the previous cap. + if consumed_all && (!is_end || remaining == 0) { + break Ok(output); + } } - } - Ok(output) + }); + result.map_err(|c| catch_zstd_error(c, vm)) } #[pyclass(with(Constructor))] @@ -1090,17 +1118,17 @@ mod _zstd { #[derive(FromArgs)] pub(super) struct ZstdDecompressorArgs { #[pyarg(any, optional)] - zstd_dict: OptionalArg, + zstd_dict: OptionalOption, #[pyarg(any, optional)] - options: OptionalArg, + options: OptionalOption, } impl Constructor for ZstdDecompressor { type Args = ZstdDecompressorArgs; fn py_new(_cls: &Py, args: Self::Args, vm: &VirtualMachine) -> PyResult { - let dict_opt = arg_or_none(args.zstd_dict, vm); - let options_opt = arg_or_none(args.options, vm); + let dict_opt = args.zstd_dict.flatten(); + let options_opt = args.options.flatten(); let mut dctx = DCtx::<'static>::create(); @@ -1108,18 +1136,9 @@ mod _zstd { apply_options(&mut dctx, options_obj, false, vm)?; } - let (held_ddict, held_dict) = load_decompressor_dict(&mut dctx, dict_opt, vm)?; - + let state = build_decompressor_state(dctx, dict_opt, vm)?; Ok(Self { - state: PyMutex::new(DecompressorState { - dctx, - _ddict: held_ddict, - _dict: held_dict, - eof: false, - needs_input: true, - unused_data: vm.ctx.empty_bytes.clone(), - input_buffer: Vec::new(), - }), + state: PyMutex::new(state), }) } } @@ -1160,70 +1179,81 @@ mod _zstd { alloc::borrow::Cow::Owned(combined) }; - let mut input = InBuffer::around(&work_data); - let mut output: Vec = Vec::new(); let chunk_size = DCtx::out_size().max(1); - // Reusable scratch buffer for each decompress_stream call. We need an - // exact-size output buffer because `Vec::reserve` may over-allocate; - // OutBuffer reports the full Vec capacity to libzstd, which would - // then happily write past `max_length`. - let mut scratch: Vec = vec![0u8; chunk_size]; - let mut hit_max = false; - let mut iteration = 0usize; - - loop { - iteration += 1; - // Honor `max_length`: stop growing the output buffer once we have - // produced enough. Special-case the first iteration when the cap - // is zero so a zero-output frame (skippable frame, empty content - // frame) can still complete; we hand libzstd a 1-byte scratch - // and discard the byte if it ends up writing one. - let grow = match max_length { - Some(maxl) if output.len() >= maxl && iteration > 1 => { - hit_max = true; - break; - } - Some(maxl) if output.len() >= maxl => 1, - Some(maxl) => (maxl - output.len()).min(chunk_size), - None => chunk_size, - }; - let result; - let written; - { - let slot = &mut scratch[..grow]; - let mut out_buf = OutBuffer::around(slot as &mut [u8]); - result = state.dctx.decompress_stream(&mut out_buf, &mut input); - written = out_buf.pos(); - } - output.extend_from_slice(&scratch[..written]); - match result { - Ok(0) => { - // Frame fully decompressed; the decompressor is at EOF. - state.eof = true; - break; - } - Ok(_) => { - let output_was_full = written == grow; - let input_consumed = input.pos == input.src.len(); - - if let Some(maxl) = max_length - && output.len() >= maxl - && iteration > 1 - { + // Release the GIL for the streaming loop. Safety: see `do_compress`; + // the closure captures only Rust-owned buffers and `&mut state` (held + // under the decompressor's `PyMutex`), and surfaces error codes as + // `usize` so we can build the exception after re-attaching. + let loop_result: Result<(Vec, bool, usize), usize> = vm.allow_threads(|| { + let mut input = InBuffer::around(&work_data); + let mut output: Vec = Vec::new(); + // Reusable scratch buffer for each decompress_stream call. We need + // an exact-size output buffer because `Vec::reserve` may + // over-allocate; `OutBuffer` reports the full Vec capacity to + // libzstd, which would then happily write past `max_length`. + let mut scratch: Vec = vec![0u8; chunk_size]; + let mut hit_max = false; + let mut iteration = 0usize; + + let outcome = loop { + iteration += 1; + // Honor `max_length`: stop growing the output buffer once + // we have produced enough. Special-case the first iteration + // when the cap is zero so a zero-output frame (skippable + // frame, empty content frame) can still complete; we hand + // libzstd a 1-byte scratch and discard the byte if it ends + // up writing one. + let grow = match max_length { + Some(maxl) if output.len() >= maxl && iteration > 1 => { hit_max = true; - break; + break Ok(()); } - - // Input is gone and libzstd had room to write but did not, - // which means the frame is incomplete and the caller has - // to supply more input. - if input_consumed && !output_was_full { - break; + Some(maxl) if output.len() >= maxl => 1, + Some(maxl) => (maxl - output.len()).min(chunk_size), + None => chunk_size, + }; + let result; + let written; + { + let slot = &mut scratch[..grow]; + let mut out_buf = OutBuffer::around(slot as &mut [u8]); + result = state.dctx.decompress_stream(&mut out_buf, &mut input); + written = out_buf.pos(); + } + output.extend_from_slice(&scratch[..written]); + match result { + Ok(0) => { + // Frame fully decompressed; the decompressor is at EOF. + state.eof = true; + break Ok(()); } + Ok(_) => { + let output_was_full = written == grow; + let input_consumed = input.pos == input.src.len(); + + if let Some(maxl) = max_length + && output.len() >= maxl + && iteration > 1 + { + hit_max = true; + break Ok(()); + } + + // Input is gone and libzstd had room to write but did + // not, which means the frame is incomplete and the + // caller has to supply more input. + if input_consumed && !output_was_full { + break Ok(()); + } + } + Err(code) => break Err(code), } - Err(code) => return Err(catch_zstd_error(code, vm)), - } - } + }; + outcome.map(|()| (output, hit_max, input.pos)) + }); + + let (mut output, mut hit_max, consumed) = + loop_result.map_err(|c| catch_zstd_error(c, vm))?; // If `max_length == 0` opened a courtesy iteration that produced more // bytes than the caller asked for, truncate. Should not happen with @@ -1235,7 +1265,6 @@ mod _zstd { hit_max = true; } - let consumed = input.pos; let remaining = &work_data[consumed..]; if state.eof { @@ -1396,17 +1425,34 @@ mod _zstd { Ok(v as usize) } + /// Sum the per-sample sizes and check they exactly cover `expected_total`, + /// rejecting overflow. A safe-fn interface that calls into libzstd with + /// these sizes must not let a wrapping sum sneak past the equality check, + /// since libzstd would then read past the samples buffer. + fn check_sample_sizes_match( + sizes: &[usize], + expected_total: usize, + vm: &VirtualMachine, + ) -> PyResult<()> { + let mismatch = || -> PyBaseExceptionRef { + vm.new_value_error("The samples size tuple doesn't match the concatenation's size") + }; + let total = sizes + .iter() + .try_fold(0usize, |a, &b| a.checked_add(b)) + .ok_or_else(mismatch)?; + if total != expected_total { + return Err(mismatch()); + } + Ok(()) + } + #[pyfunction] fn train_dict(args: TrainDictArgs, vm: &VirtualMachine) -> PyResult> { let dict_size = parse_positive_dict_size(&args.dict_size, vm)?; let samples_buffer = args.samples_bytes.as_bytes().to_vec(); let sizes = parse_sample_sizes(args.samples_sizes, vm)?; - let total: usize = sizes.iter().sum(); - if total != samples_buffer.len() { - return Err( - vm.new_value_error("The samples size tuple doesn't match the concatenation's size") - ); - } + check_sample_sizes_match(&sizes, samples_buffer.len(), vm)?; let mut dict_buffer: Vec = Vec::with_capacity(dict_size); zstd_safe::train_from_buffer(&mut dict_buffer, &samples_buffer, &sizes) .map_err(|c| catch_zstd_error(c, vm))?; @@ -1439,16 +1485,11 @@ mod _zstd { #[pyfunction] fn finalize_dict(args: FinalizeDictArgs, vm: &VirtualMachine) -> PyResult> { let dict_size = parse_positive_dict_size(&args.dict_size, vm)?; - let compression_level = pyobj_to_i32(&args.compression_level, vm)?; + let compression_level: i32 = args.compression_level.try_to_value(vm)?; let custom_dict = args.custom_dict_bytes.as_bytes().to_vec(); let samples_buffer = args.samples_bytes.as_bytes().to_vec(); let sizes = parse_sample_sizes(args.samples_sizes, vm)?; - let total: usize = sizes.iter().sum(); - if total != samples_buffer.len() { - return Err( - vm.new_value_error("The samples size tuple doesn't match the concatenation's size") - ); - } + check_sample_sizes_match(&sizes, samples_buffer.len(), vm)?; let mut dict_buffer: Vec = vec![0u8; dict_size]; let params = zstd_sys::ZDICT_params_t { From 9b1e29eb76e4cd251f8ddc67d0caeacd0cb8d84a Mon Sep 17 00:00:00 2001 From: James David Clarke Date: Wed, 3 Jun 2026 14:48:21 +0100 Subject: [PATCH 10/20] Identify zstd parameter enums by type identity, not name - Have set_parameter_types stash the CompressionParameter/DecompressionParameter classes as private _zstd module attributes instead of validating and discarding them - check_wrong_param_kind now compares key classes by identity against the registered type, matching CPython's Py_TYPE check, and skips when unregistered - Error message names the actual key type as an attribute --- crates/stdlib/src/zstd.rs | 87 ++++++++++++++++++--------------------- 1 file changed, 41 insertions(+), 46 deletions(-) diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs index 58f5f807996..c17edae6559 100644 --- a/crates/stdlib/src/zstd.rs +++ b/crates/stdlib/src/zstd.rs @@ -39,16 +39,6 @@ mod _zstd { use zstd_safe::zstd_sys; use zstd_safe::{CCtx, CParameter, DCtx, DParameter, InBuffer, OutBuffer}; - /// Class names of the `CompressionParameter` and `DecompressionParameter` - /// `IntEnum`s defined by the pure-Python wrapper at - /// `Lib/compression/zstd/__init__.py`. We identify these by name rather - /// than identity because storing `PyTypeRef` in a Rust `static` requires - /// `Sync`, which `PyMutex` does not provide on single-threaded targets - /// (Android, iOS, wasm32). Name comparison is sufficient since both names - /// originate from a module we ship and control. - const C_PARAMETER_TYPE_NAME: &str = "CompressionParameter"; - const D_PARAMETER_TYPE_NAME: &str = "DecompressionParameter"; - // ========================================================================= // Module-level constants // ========================================================================= @@ -177,29 +167,30 @@ mod _zstd { // Parameter helpers // ========================================================================= - /// If `key`'s class is the wrong kind of parameter enum for the caller's - /// context (a `CompressionParameter` passed to a decompressor, or vice - /// versa), return a `TypeError`. Used by both (de)compressor option-dict - /// loops so the resulting message names the actual type rather than the - /// generic "out of range" wording. + /// Reject an options-dict `key` whose class is the parameter enum that is + /// invalid for the caller's context (a `CompressionParameter` passed to a + /// decompressor, or vice versa) with a `TypeError` naming the type. /// - /// We identify the enums by class name (see [`C_PARAMETER_TYPE_NAME`]), - /// which avoids storing a `PyTypeRef` in a Rust `static` (impossible on - /// single-threaded builds whose `PyMutex` is not `Sync`). + /// `forbidden` is that invalid enum class, resolved once by the caller from + /// the types the pure-Python wrapper registers via [`set_parameter_types`]; + /// `None` (the wrapper never ran) skips the check, matching CPython's NULL + /// module-state pointers. The comparison is by identity, mirroring + /// CPython's `Py_TYPE(key) == ...` check. fn check_wrong_param_kind( key: &PyObjectRef, - is_compress: bool, + forbidden: Option<&PyObjectRef>, + kind: &str, vm: &VirtualMachine, ) -> PyResult<()> { - let class_name = key.class().name(); - let (forbidden, kind) = if is_compress { - (D_PARAMETER_TYPE_NAME, "compression") - } else { - (C_PARAMETER_TYPE_NAME, "decompression") + let Some(forbidden) = forbidden else { + return Ok(()); }; - if &*class_name == forbidden { + if key.class().is(forbidden) { + // `key`'s class is `forbidden` here, so name it directly (the same + // string CPython formats from `Py_TYPE(key)->tp_name`). + let name = key.class().name(); return Err(vm.new_type_error(format!( - "{kind} options dictionary key must not be a {forbidden}" + "{kind} options dictionary key must not be a {name} attribute" ))); } Ok(()) @@ -705,12 +696,21 @@ mod _zstd { let dict = options_obj .downcast::() .map_err(|_| vm.new_type_error("options must be a dict"))?; + // Resolve the parameter enum class that is invalid for this direction + // (registered by the pure-Python wrapper via set_parameter_types) once, + // so a key from the wrong family yields a clear TypeError naming the + // type. `None` when the wrapper never ran (e.g. `_zstd` used directly). + let (wrong_kind_attr, kind) = if is_compress { + ("_decompression_parameter_type", "compression") + } else { + ("_compression_parameter_type", "decompression") + }; + let wrong_kind = vm.get_attribute_opt(vm.import("_zstd", 0)?, wrong_kind_attr)?; for (k, v) in dict { - // Reject the wrong enum kind ("CompressionParameter" vs - // "DecompressionParameter") before doing any numeric coercion, - // so the error carries the documented type name rather than a + // Reject a key from the wrong parameter family before any numeric + // coercion, so the error names the type rather than giving a // generic out-of-range message. - check_wrong_param_kind(&k, is_compress, vm)?; + check_wrong_param_kind(&k, wrong_kind.as_ref(), kind, vm)?; let key_int: i32 = k.try_to_value(vm)?; let val_int: i32 = v.try_to_value(vm)?; // libzstd silently clamps out-of-range values for some @@ -1572,26 +1572,21 @@ mod _zstd { Ok((bounds.lowerBound, bounds.upperBound)) } - // The (de)compressor constructors identify the parameter enums by class - // *name* (see [`check_wrong_param_kind`]) rather than identity, so this - // function does not need to retain the type objects beyond the call. - // CPython retains them to produce nicer error messages naming the specific - // enum member; matching that would require a `Sync` cell-of-`PyTypeRef`, - // which is not available on single-threaded targets. The pure-Python - // wrapper at `Lib/compression/zstd/__init__.py:242` calls this exactly - // once at import time. + // Register the `CompressionParameter` / `DecompressionParameter` enum + // classes defined by the pure-Python wrapper so [`check_wrong_param_kind`] + // can reject a key from the wrong parameter family by identity. The types + // are stashed as private `_zstd` module attributes — the RustPython + // equivalent of the module state CPython keeps these in. The wrapper in + // `Lib/compression/zstd/__init__.py` calls this exactly once at import. #[pyfunction] fn set_parameter_types( - c_parameter_type: PyObjectRef, - d_parameter_type: PyObjectRef, + c_parameter_type: PyTypeRef, + d_parameter_type: PyTypeRef, vm: &VirtualMachine, ) -> PyResult<()> { - let _ = c_parameter_type - .downcast::() - .map_err(|_| vm.new_type_error("c_parameter_type must be a type object"))?; - let _ = d_parameter_type - .downcast::() - .map_err(|_| vm.new_type_error("d_parameter_type must be a type object"))?; + let module = vm.import("_zstd", 0)?; + module.set_attr("_compression_parameter_type", c_parameter_type, vm)?; + module.set_attr("_decompression_parameter_type", d_parameter_type, vm)?; Ok(()) } } From e6c5b19d207afb36b7ca006fe0eef25f9ab3da48 Mon Sep 17 00:00:00 2001 From: James David Clarke Date: Wed, 3 Jun 2026 15:24:11 +0100 Subject: [PATCH 11/20] Match CPython error behavior in _zstd ZstdCompressor - Raise TypeError (not RuntimeError), matching CPython, when both `level` and `options` are passed to ZstdCompressor - Parse set_pledged_input_size's argument before taking the state lock so a re-entrant __index__ can't deadlock, and fix the off-by-one in its bound message to match CPython Co-Authored-By: Claude Opus 4.8 (1M context) --- crates/stdlib/src/zstd.rs | 51 ++++++++++++++++++++++----------------- 1 file changed, 29 insertions(+), 22 deletions(-) diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs index c17edae6559..e6f96b60a1a 100644 --- a/crates/stdlib/src/zstd.rs +++ b/crates/stdlib/src/zstd.rs @@ -638,7 +638,7 @@ mod _zstd { let dict_opt = args.zstd_dict.flatten(); if level_opt.is_some() && options_opt.is_some() { - return Err(vm.new_runtime_error("Only one of level or options should be used")); + return Err(vm.new_type_error("Only one of level or options should be used.")); } let mut cctx = CCtx::<'static>::create(); @@ -1000,37 +1000,44 @@ mod _zstd { #[pymethod] fn set_pledged_input_size(&self, size: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { - let mut state = self.state.lock(); - if state.last_mode != COMP_MODE_FLUSH_FRAME { - return Err(vm.new_value_error( - "set_pledged_input_size() method must be called when last_mode == FLUSH_FRAME", - )); - } + // Parse the argument *before* taking the lock: `try_index` can run a + // Python `__index__`, and doing that while holding `self.state` would + // let a re-entrant call into this compressor deadlock. CPython + // likewise converts the argument before touching the compressor. + // // Python passes `None` to mean "unknown"; libzstd represents that // internally as `ZSTD_CONTENTSIZE_UNKNOWN` (`u64::MAX`), and - // `zstd_safe` translates `None` accordingly. Reject any concrete - // int outside the valid range up front so callers see the - // documented `ValueError`, not a libzstd-level error. The valid - // upper bound is `ZSTD_CONTENTSIZE_ERROR - 1`, i.e. `u64::MAX - 2`. + // `zstd_safe` translates `None` accordingly. libzstd also reserves + // `ZSTD_CONTENTSIZE_ERROR` (`u64::MAX - 1`), so a concrete size must + // be strictly less than that; reject anything else up front so + // callers see the documented `ValueError`, not a libzstd-level error. let pledged: Option = if vm.is_none(&size) { None } else { - const UPPER_BOUND: u64 = u64::MAX - 2; - let int_ref = size.try_index(vm)?; - // Try fitting into u64; OverflowError covers negatives and - // values above `u64::MAX`. - let v: u64 = int_ref.try_to_primitive(vm).map_err(|_| { + const LIMIT: u64 = u64::MAX - 1; + let err = || { vm.new_value_error(format!( - "size argument should be a positive int less than {UPPER_BOUND}" + "size argument should be a positive int less than {LIMIT}" )) - })?; - if v > UPPER_BOUND { - return Err(vm.new_value_error(format!( - "size argument should be a positive int less than {UPPER_BOUND}" - ))); + }; + // `try_to_primitive` fails (OverflowError) for negatives and for + // values above `u64::MAX`; the explicit check covers the rest of + // the reserved range. + let v: u64 = size + .try_index(vm)? + .try_to_primitive(vm) + .map_err(|_| err())?; + if v >= LIMIT { + return Err(err()); } Some(v) }; + let mut state = self.state.lock(); + if state.last_mode != COMP_MODE_FLUSH_FRAME { + return Err(vm.new_value_error( + "set_pledged_input_size() method must be called when last_mode == FLUSH_FRAME", + )); + } state .cctx .set_pledged_src_size(pledged) From dbe4626d0bebd7af878e1815ae32ed4d8680b5cb Mon Sep 17 00:00:00 2001 From: James David Clarke Date: Sun, 9 Aug 2026 16:31:50 +0100 Subject: [PATCH 12/20] Address review: drop decorative separators and unused zstd crate Remove the remaining `// ===...` section separators from zstd.rs per the project comment guidelines, drop the unused `zstd` crate dependency (only zstd-safe/zstd-sys are used) and declare zstd-safe's std feature explicitly instead of relying on the zstd crate's feature unification, and convert a doc comment on #[extend_class] to a regular comment. --- Cargo.toml | 3 +-- crates/stdlib/Cargo.toml | 1 - crates/stdlib/src/zstd.rs | 36 ++++-------------------------------- 3 files changed, 5 insertions(+), 35 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 067d4960beb..d605f8b906d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -323,8 +323,7 @@ x509-cert = "0.3.0" x509-parser = "0.18" xml = "1.3" writeable = "0.6" -zstd = { version = "0.13", default-features = false, features = ["zdict_builder"] } -zstd-safe = { version = "7", default-features = false, features = ["zdict_builder"] } +zstd-safe = { version = "7", default-features = false, features = ["std", "zdict_builder"] } # Lints diff --git a/crates/stdlib/Cargo.toml b/crates/stdlib/Cargo.toml index 9f8dc95ef4c..80130e6e3f9 100644 --- a/crates/stdlib/Cargo.toml +++ b/crates/stdlib/Cargo.toml @@ -121,7 +121,6 @@ pkcs8 = { workspace = true, features = ["encryption", "pkcs5", "pem"], optional libsqlite3-sys = { workspace = true, features = ["bundled"], optional = true } xz = { workspace = true } xz-sys = { workspace = true } -zstd = { workspace = true } zstd-safe = { workspace = true } [target.'cfg(windows)'.dependencies] diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs index e6f96b60a1a..2c7e1558d02 100644 --- a/crates/stdlib/src/zstd.rs +++ b/crates/stdlib/src/zstd.rs @@ -39,10 +39,6 @@ mod _zstd { use zstd_safe::zstd_sys; use zstd_safe::{CCtx, CParameter, DCtx, DParameter, InBuffer, OutBuffer}; - // ========================================================================= - // Module-level constants - // ========================================================================= - #[pyattr] const ZSTD_CLEVEL_DEFAULT: i32 = zstd_sys::ZSTD_CLEVEL_DEFAULT as i32; @@ -138,10 +134,6 @@ mod _zstd { const DICT_TYPE_UNDIGESTED: i32 = 1; const DICT_TYPE_PREFIX: i32 = 2; - // ========================================================================= - // ZstdError exception - // ========================================================================= - #[pyattr(once, name = "ZstdError")] fn zstd_error(vm: &VirtualMachine) -> PyTypeRef { vm.ctx.new_exception_type( @@ -163,10 +155,6 @@ mod _zstd { new_zstd_error(zstd_safe::get_error_name(code).to_string(), vm) } - // ========================================================================= - // Parameter helpers - // ========================================================================= - /// Reject an options-dict `key` whose class is the parameter enum that is /// invalid for the caller's context (a `CompressionParameter` passed to a /// decompressor, or vice versa) with a `TypeError` naming the type. @@ -442,10 +430,6 @@ mod _zstd { Err(vm.new_type_error("zstd_dict argument should be a ZstdDict object")) } - // ========================================================================= - // ZstdDict - // ========================================================================= - #[derive(FromArgs)] pub(super) struct ZstdDictArgs { #[pyarg(positional)] @@ -558,10 +542,6 @@ mod _zstd { } } - // ========================================================================= - // ZstdCompressor - // ========================================================================= - // The three flush modes for `ZstdCompressor.compress()`, mirrored as // class attributes via `extend_class` below. Values are positional and // chosen to match what CPython exposes. @@ -1050,10 +1030,10 @@ mod _zstd { self.state.lock().last_mode } - /// Install class-level constants `CONTINUE`, `FLUSH_BLOCK`, and - /// `FLUSH_FRAME` so callers can reference them as - /// `ZstdCompressor.FLUSH_FRAME` (as the Python `ZstdFile` wrapper - /// does). + // Install class-level constants `CONTINUE`, `FLUSH_BLOCK`, and + // `FLUSH_FRAME` so callers can reference them as + // `ZstdCompressor.FLUSH_FRAME` (as the Python `ZstdFile` wrapper + // does). #[extend_class] fn extend_class(ctx: &Context, class: &'static Py) { class.set_attr( @@ -1085,10 +1065,6 @@ mod _zstd { mode: Option, } - // ========================================================================= - // ZstdDecompressor - // ========================================================================= - /// Internal state of a `ZstdDecompressor`. The CPython decompressor is /// single-frame: once we hit end-of-frame, additional bytes go into /// `unused_data` and further `decompress` calls raise `EOFError`. Field @@ -1339,10 +1315,6 @@ mod _zstd { } } - // ========================================================================= - // Module-level functions - // ========================================================================= - #[pyfunction] fn get_frame_size(frame_buffer: ArgBytesLike, vm: &VirtualMachine) -> PyResult { let buf = frame_buffer.with_ref(|b| b.to_vec()); From 9d497f4ee2fac08b34b03aa448812edf0ecaa453 Mon Sep 17 00:00:00 2001 From: James David Clarke Date: Sun, 9 Aug 2026 17:26:37 +0100 Subject: [PATCH 13/20] Switch _zstd backend from zstd-safe to pure-Rust libzstd-rs-sys Replace the C-library bindings (zstd-safe/zstd-sys) with Trifecta Tech Foundation's libzstd-rs-sys (pinned git rev), a c2rust translation of upstream libzstd. The module now calls its C-API-compatible functions through small owning RAII wrappers (CCtx/DCtx/CDict/DDict) with the same drop-order invariants as before. The new backend builds zstdmt_compress in, so nb_workers bounds are (0, 256) and test_zstd_multithread_compress passes for real; drop the now-obsolete expectedFailureIf marker. --- Cargo.lock | 10 + Cargo.toml | 2 +- Lib/test/test_zstd.py | 5 - crates/stdlib/Cargo.toml | 2 +- crates/stdlib/src/zstd.rs | 824 +++++++++++++++++++++++++------------- 5 files changed, 549 insertions(+), 294 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index fa5e7b52236..7d38d0762be 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2040,6 +2040,15 @@ dependencies = [ "zlib-rs", ] +[[package]] +name = "libzstd-rs-sys" +version = "0.0.1-prerelease.2" +source = "git+https://github.com/trifectatechfoundation/libzstd-rs-sys?rev=99fb4c096a21a9b37feebe180333a2ff8645e23b#99fb4c096a21a9b37feebe180333a2ff8645e23b" +dependencies = [ + "cc", + "libc", +] + [[package]] name = "linux-raw-sys" version = "0.12.1" @@ -3578,6 +3587,7 @@ dependencies = [ "libc", "libsqlite3-sys", "libz-rs-sys", + "libzstd-rs-sys", "malachite-bigint", "md-5", "memchr", diff --git a/Cargo.toml b/Cargo.toml index d605f8b906d..0446793ec1e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -241,6 +241,7 @@ xz = "0.4.6-rc.0" xz-sys = "0.4.6-rc.0" libsqlite3-sys = "0.38" libz-rs-sys = "0.6" +libzstd-rs-sys = { git = "https://github.com/trifectatechfoundation/libzstd-rs-sys", rev = "99fb4c096a21a9b37feebe180333a2ff8645e23b" } lock_api = "0.4" log = "0.4.30" lz4_flex = "0.13" @@ -323,7 +324,6 @@ x509-cert = "0.3.0" x509-parser = "0.18" xml = "1.3" writeable = "0.6" -zstd-safe = { version = "7", default-features = false, features = ["std", "zdict_builder"] } # Lints diff --git a/Lib/test/test_zstd.py b/Lib/test/test_zstd.py index 1b4b819d2ef..44a5bbb4b19 100644 --- a/Lib/test/test_zstd.py +++ b/Lib/test/test_zstd.py @@ -321,11 +321,6 @@ def test_unknown_compression_parameter(self): with self.assertRaisesRegex(ValueError, pattern): ZstdCompressor(options=option) - # TODO: RUSTPYTHON - @unittest.expectedFailureIf( - not SUPPORT_MULTITHREADING, - "TODO: RUSTPYTHON; libzstd built without multi-threaded compression support", - ) def test_zstd_multithread_compress(self): size = 40*_1M b = THIS_FILE_BYTES * (size // len(THIS_FILE_BYTES)) diff --git a/crates/stdlib/Cargo.toml b/crates/stdlib/Cargo.toml index 80130e6e3f9..605d62702e1 100644 --- a/crates/stdlib/Cargo.toml +++ b/crates/stdlib/Cargo.toml @@ -119,9 +119,9 @@ 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 } +libzstd-rs-sys = { workspace = true } xz = { workspace = true } xz-sys = { workspace = true } -zstd-safe = { workspace = true } [target.'cfg(windows)'.dependencies] paste = { workspace = true } diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs index 2c7e1558d02..8d9e194bfd6 100644 --- a/crates/stdlib/src/zstd.rs +++ b/crates/stdlib/src/zstd.rs @@ -3,7 +3,7 @@ // spell-checker:ignore btlazy dfast nbworkers windowlogmax windowlog overlap targetcblock // spell-checker:ignore srcsize zdict refprefix refcdict refddict pledgedsrcsize getframecontentsize // spell-checker:ignore Zstd Zstandard pylib RFC -// spell-checker:ignore CLEVEL zstdmodule cparameter dparameter maxl +// spell-checker:ignore CLEVEL zstdmodule cparameter dparameter maxl c2rust // spell-checker:ignore cctx dctx CCTX DCTX ldm cdict ddict windowlog hashlog chainlog searchlog CLEVEL //! The `_zstd` extension module. Backs the pure-Python `compression.zstd` @@ -12,10 +12,9 @@ //! `Lib/compression/zstd/__init__.py` imports from this module unconditionally, //! so the names and call signatures here must stay in sync with CPython. //! -//! Backend: the `zstd_safe` Rust crate (a thin safe wrapper over Facebook's -//! libzstd C library, which is what CPython links against). A handful of -//! routines that `zstd_safe` does not expose at a safe level (parameter -//! bounds, dictionary finalization) drop down to raw `zstd_sys` FFI calls. +//! Backend: the `libzstd-rs-sys` crate (Trifecta Tech Foundation's pure-Rust +//! c2rust translation of Facebook's libzstd, the C library CPython links +//! against), used directly through its `unsafe extern "C"` API surface. pub(crate) use _zstd::module_def; @@ -26,7 +25,23 @@ pub(crate) use _zstd::module_def; #[allow(non_upper_case_globals)] #[pymodule] mod _zstd { - use core::ffi::c_int; + use core::ffi::{CStr, c_int}; + use libzstd_rs_sys::lib::compress::zstd_compress::{ZSTD_e_continue, ZSTD_e_end, ZSTD_e_flush}; + use libzstd_rs_sys::lib::zdict::ZDICT_finalizeDictionary; + use libzstd_rs_sys::{ + ZDICT_getErrorName, ZDICT_isError, ZDICT_params_t, ZDICT_trainFromBuffer, ZSTD_CCtx, + ZSTD_CCtx_loadDictionary, ZSTD_CCtx_refCDict, ZSTD_CCtx_refPrefix, ZSTD_CCtx_setParameter, + ZSTD_CCtx_setPledgedSrcSize, ZSTD_CDict, ZSTD_CONTENTSIZE_ERROR, ZSTD_CONTENTSIZE_UNKNOWN, + ZSTD_CStreamOutSize, ZSTD_DCtx, ZSTD_DCtx_loadDictionary, ZSTD_DCtx_refDDict, + ZSTD_DCtx_refPrefix, ZSTD_DCtx_setParameter, ZSTD_DDict, ZSTD_DStreamOutSize, + ZSTD_EndDirective, ZSTD_cParam_getBounds, ZSTD_cParameter, ZSTD_compressStream2, + ZSTD_createCCtx, ZSTD_createCDict, ZSTD_createDCtx, ZSTD_createDDict, + ZSTD_dParam_getBounds, ZSTD_dParameter, ZSTD_decompressStream, + ZSTD_findFrameCompressedSize, ZSTD_freeCCtx, ZSTD_freeCDict, ZSTD_freeDCtx, ZSTD_freeDDict, + ZSTD_getDictID_fromDict, ZSTD_getDictID_fromFrame, ZSTD_getErrorName, + ZSTD_getFrameContentSize, ZSTD_inBuffer, ZSTD_isError, ZSTD_outBuffer, ZSTD_strategy, + ZSTD_versionNumber, ZSTD_versionString, + }; use rustpython_common::lock::PyMutex; use rustpython_vm::builtins::{ PyBaseExceptionRef, PyBytesRef, PyDict, PyTupleRef, PyType, PyTypeRef, @@ -36,94 +51,98 @@ mod _zstd { use rustpython_vm::{ AsObject, Context, Py, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, }; - use zstd_safe::zstd_sys; - use zstd_safe::{CCtx, CParameter, DCtx, DParameter, InBuffer, OutBuffer}; #[pyattr] - const ZSTD_CLEVEL_DEFAULT: i32 = zstd_sys::ZSTD_CLEVEL_DEFAULT as i32; + const ZSTD_CLEVEL_DEFAULT: i32 = libzstd_rs_sys::ZSTD_CLEVEL_DEFAULT; // Compression parameter identifiers. Values match the `ZSTD_cParameter` // enum in libzstd, which is what the public `CompressionParameter` IntEnum // in `Lib/compression/zstd/__init__.py` derives its members from. + // libzstd-rs-sys models the enum as a newtype whose inner value is + // crate-private, so the ids — frozen C ABI values from zstd.h — are + // spelled out as literals here. #[pyattr] - const ZSTD_c_compressionLevel: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_compressionLevel as i32; + const ZSTD_c_compressionLevel: i32 = 100; #[pyattr] - const ZSTD_c_windowLog: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_windowLog as i32; + const ZSTD_c_windowLog: i32 = 101; #[pyattr] - const ZSTD_c_hashLog: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_hashLog as i32; + const ZSTD_c_hashLog: i32 = 102; #[pyattr] - const ZSTD_c_chainLog: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_chainLog as i32; + const ZSTD_c_chainLog: i32 = 103; #[pyattr] - const ZSTD_c_searchLog: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_searchLog as i32; + const ZSTD_c_searchLog: i32 = 104; #[pyattr] - const ZSTD_c_minMatch: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_minMatch as i32; + const ZSTD_c_minMatch: i32 = 105; #[pyattr] - const ZSTD_c_targetLength: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_targetLength as i32; + const ZSTD_c_targetLength: i32 = 106; #[pyattr] - const ZSTD_c_strategy: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_strategy as i32; + const ZSTD_c_strategy: i32 = 107; #[pyattr] - const ZSTD_c_enableLongDistanceMatching: i32 = - zstd_sys::ZSTD_cParameter::ZSTD_c_enableLongDistanceMatching as i32; + const ZSTD_c_enableLongDistanceMatching: i32 = 160; #[pyattr] - const ZSTD_c_ldmHashLog: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_ldmHashLog as i32; + const ZSTD_c_ldmHashLog: i32 = 161; #[pyattr] - const ZSTD_c_ldmMinMatch: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_ldmMinMatch as i32; + const ZSTD_c_ldmMinMatch: i32 = 162; #[pyattr] - const ZSTD_c_ldmBucketSizeLog: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_ldmBucketSizeLog as i32; + const ZSTD_c_ldmBucketSizeLog: i32 = 163; #[pyattr] - const ZSTD_c_ldmHashRateLog: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_ldmHashRateLog as i32; + const ZSTD_c_ldmHashRateLog: i32 = 164; #[pyattr] - const ZSTD_c_contentSizeFlag: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_contentSizeFlag as i32; + const ZSTD_c_contentSizeFlag: i32 = 200; #[pyattr] - const ZSTD_c_checksumFlag: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_checksumFlag as i32; + const ZSTD_c_checksumFlag: i32 = 201; #[pyattr] - const ZSTD_c_dictIDFlag: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_dictIDFlag as i32; + const ZSTD_c_dictIDFlag: i32 = 202; #[pyattr] - const ZSTD_c_nbWorkers: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_nbWorkers as i32; + const ZSTD_c_nbWorkers: i32 = 400; #[pyattr] - const ZSTD_c_jobSize: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_jobSize as i32; + const ZSTD_c_jobSize: i32 = 401; #[pyattr] - const ZSTD_c_overlapLog: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_overlapLog as i32; + const ZSTD_c_overlapLog: i32 = 402; // Decompression parameter identifiers. libzstd only exposes one non- // experimental decompression parameter. #[pyattr] - const ZSTD_d_windowLogMax: i32 = zstd_sys::ZSTD_dParameter::ZSTD_d_windowLogMax as i32; + const ZSTD_d_windowLogMax: i32 = 100; // Strategy enum members ordered from fastest to strongest. These power // the `Strategy` IntEnum in `Lib/compression/zstd/__init__.py`. #[pyattr] - const ZSTD_fast: i32 = zstd_sys::ZSTD_strategy::ZSTD_fast as i32; + const ZSTD_fast: i32 = libzstd_rs_sys::lib::zstd::ZSTD_fast as i32; #[pyattr] - const ZSTD_dfast: i32 = zstd_sys::ZSTD_strategy::ZSTD_dfast as i32; + const ZSTD_dfast: i32 = libzstd_rs_sys::lib::zstd::ZSTD_dfast as i32; #[pyattr] - const ZSTD_greedy: i32 = zstd_sys::ZSTD_strategy::ZSTD_greedy as i32; + const ZSTD_greedy: i32 = libzstd_rs_sys::lib::zstd::ZSTD_greedy as i32; #[pyattr] - const ZSTD_lazy: i32 = zstd_sys::ZSTD_strategy::ZSTD_lazy as i32; + const ZSTD_lazy: i32 = libzstd_rs_sys::lib::zstd::ZSTD_lazy as i32; #[pyattr] - const ZSTD_lazy2: i32 = zstd_sys::ZSTD_strategy::ZSTD_lazy2 as i32; + const ZSTD_lazy2: i32 = libzstd_rs_sys::lib::zstd::ZSTD_lazy2 as i32; #[pyattr] - const ZSTD_btlazy2: i32 = zstd_sys::ZSTD_strategy::ZSTD_btlazy2 as i32; + const ZSTD_btlazy2: i32 = libzstd_rs_sys::lib::zstd::ZSTD_btlazy2 as i32; #[pyattr] - const ZSTD_btopt: i32 = zstd_sys::ZSTD_strategy::ZSTD_btopt as i32; + const ZSTD_btopt: i32 = libzstd_rs_sys::lib::zstd::ZSTD_btopt as i32; #[pyattr] - const ZSTD_btultra: i32 = zstd_sys::ZSTD_strategy::ZSTD_btultra as i32; + const ZSTD_btultra: i32 = libzstd_rs_sys::lib::zstd::ZSTD_btultra as i32; #[pyattr] - const ZSTD_btultra2: i32 = zstd_sys::ZSTD_strategy::ZSTD_btultra2 as i32; + const ZSTD_btultra2: i32 = libzstd_rs_sys::lib::zstd::ZSTD_btultra2 as i32; #[pyattr(once, name = "zstd_version")] fn zstd_version(_vm: &VirtualMachine) -> String { - zstd_safe::version_string().to_string() + // SAFETY: `ZSTD_versionString` returns a pointer to libzstd's static, + // NUL-terminated version string. + unsafe { CStr::from_ptr(ZSTD_versionString()) } + .to_string_lossy() + .into_owned() } #[pyattr(once, name = "zstd_version_number")] fn zstd_version_number(_vm: &VirtualMachine) -> u32 { - zstd_safe::version_number() + ZSTD_versionNumber() } #[pyattr(once, name = "ZSTD_DStreamOutSize")] fn zstd_dstream_out_size(_vm: &VirtualMachine) -> usize { - DCtx::out_size() + ZSTD_DStreamOutSize() } // Dictionary load type markers. The `ZstdDict.as_*` properties wrap the @@ -152,7 +171,10 @@ mod _zstd { /// functions when `ZSTD_isError(code)` is non-zero) into a `ZstdError` /// carrying the human-readable message from `ZSTD_getErrorName`. fn catch_zstd_error(code: usize, vm: &VirtualMachine) -> PyBaseExceptionRef { - new_zstd_error(zstd_safe::get_error_name(code).to_string(), vm) + // SAFETY: `ZSTD_getErrorName` returns a pointer to a static, + // NUL-terminated error string from libzstd's error table. + let name = unsafe { CStr::from_ptr(ZSTD_getErrorName(code)) }; + new_zstd_error(name.to_string_lossy(), vm) } /// Reject an options-dict `key` whose class is the parameter enum that is @@ -184,92 +206,80 @@ mod _zstd { Ok(()) } - /// Map a compression parameter id to its raw libzstd `ZSTD_cParameter` - /// enum value. Returns `None` for unknown ids so callers can surface a + /// Map a compression parameter id to its libzstd `ZSTD_cParameter` + /// constant. Returns `None` for unknown ids so callers can surface a /// targeted `ValueError`. Done with an explicit match rather than /// `mem::transmute` so passing junk like `ZSTD_cParameter(42)` cannot /// be triggered from Python. - fn c_param_enum(param: i32) -> Option { - use zstd_sys::ZSTD_cParameter as P; + fn c_param_enum(param: i32) -> Option { Some(match param { - ZSTD_c_compressionLevel => P::ZSTD_c_compressionLevel, - ZSTD_c_windowLog => P::ZSTD_c_windowLog, - ZSTD_c_hashLog => P::ZSTD_c_hashLog, - ZSTD_c_chainLog => P::ZSTD_c_chainLog, - ZSTD_c_searchLog => P::ZSTD_c_searchLog, - ZSTD_c_minMatch => P::ZSTD_c_minMatch, - ZSTD_c_targetLength => P::ZSTD_c_targetLength, - ZSTD_c_strategy => P::ZSTD_c_strategy, - ZSTD_c_enableLongDistanceMatching => P::ZSTD_c_enableLongDistanceMatching, - ZSTD_c_ldmHashLog => P::ZSTD_c_ldmHashLog, - ZSTD_c_ldmMinMatch => P::ZSTD_c_ldmMinMatch, - ZSTD_c_ldmBucketSizeLog => P::ZSTD_c_ldmBucketSizeLog, - ZSTD_c_ldmHashRateLog => P::ZSTD_c_ldmHashRateLog, - ZSTD_c_contentSizeFlag => P::ZSTD_c_contentSizeFlag, - ZSTD_c_checksumFlag => P::ZSTD_c_checksumFlag, - ZSTD_c_dictIDFlag => P::ZSTD_c_dictIDFlag, - ZSTD_c_nbWorkers => P::ZSTD_c_nbWorkers, - ZSTD_c_jobSize => P::ZSTD_c_jobSize, - ZSTD_c_overlapLog => P::ZSTD_c_overlapLog, + ZSTD_c_compressionLevel => ZSTD_cParameter::ZSTD_c_compressionLevel, + ZSTD_c_windowLog => ZSTD_cParameter::ZSTD_c_windowLog, + ZSTD_c_hashLog => ZSTD_cParameter::ZSTD_c_hashLog, + ZSTD_c_chainLog => ZSTD_cParameter::ZSTD_c_chainLog, + ZSTD_c_searchLog => ZSTD_cParameter::ZSTD_c_searchLog, + ZSTD_c_minMatch => ZSTD_cParameter::ZSTD_c_minMatch, + ZSTD_c_targetLength => ZSTD_cParameter::ZSTD_c_targetLength, + ZSTD_c_strategy => ZSTD_cParameter::ZSTD_c_strategy, + ZSTD_c_enableLongDistanceMatching => ZSTD_cParameter::ZSTD_c_enableLongDistanceMatching, + ZSTD_c_ldmHashLog => ZSTD_cParameter::ZSTD_c_ldmHashLog, + ZSTD_c_ldmMinMatch => ZSTD_cParameter::ZSTD_c_ldmMinMatch, + ZSTD_c_ldmBucketSizeLog => ZSTD_cParameter::ZSTD_c_ldmBucketSizeLog, + ZSTD_c_ldmHashRateLog => ZSTD_cParameter::ZSTD_c_ldmHashRateLog, + ZSTD_c_contentSizeFlag => ZSTD_cParameter::ZSTD_c_contentSizeFlag, + ZSTD_c_checksumFlag => ZSTD_cParameter::ZSTD_c_checksumFlag, + ZSTD_c_dictIDFlag => ZSTD_cParameter::ZSTD_c_dictIDFlag, + ZSTD_c_nbWorkers => ZSTD_cParameter::ZSTD_c_nbWorkers, + ZSTD_c_jobSize => ZSTD_cParameter::ZSTD_c_jobSize, + ZSTD_c_overlapLog => ZSTD_cParameter::ZSTD_c_overlapLog, _ => return None, }) } - /// Map a decompression parameter id to its raw libzstd `ZSTD_dParameter` - /// enum value. See [`c_param_enum`] for rationale. - fn d_param_enum(param: i32) -> Option { - use zstd_sys::ZSTD_dParameter as P; + /// Map a decompression parameter id to its libzstd `ZSTD_dParameter` + /// constant. See [`c_param_enum`] for rationale. + fn d_param_enum(param: i32) -> Option { match param { - ZSTD_d_windowLogMax => Some(P::ZSTD_d_windowLogMax), + ZSTD_d_windowLogMax => Some(ZSTD_dParameter::ZSTD_d_windowLogMax), _ => None, } } - /// Map a compression-parameter id and integer value to a `CParameter` - /// variant. Used by the compressor's `options=` constructor argument. - fn cparameter_from_int(param: i32, value: i32, vm: &VirtualMachine) -> PyResult { - let p = match param { - ZSTD_c_compressionLevel => CParameter::CompressionLevel(value), - ZSTD_c_windowLog => CParameter::WindowLog(value as u32), - ZSTD_c_hashLog => CParameter::HashLog(value as u32), - ZSTD_c_chainLog => CParameter::ChainLog(value as u32), - ZSTD_c_searchLog => CParameter::SearchLog(value as u32), - ZSTD_c_minMatch => CParameter::MinMatch(value as u32), - ZSTD_c_targetLength => CParameter::TargetLength(value as u32), - ZSTD_c_strategy => { - CParameter::Strategy(strategy_from_int(value).ok_or_else(|| { - new_zstd_error(format!("invalid strategy value: {value}"), vm) - })?) - } - ZSTD_c_enableLongDistanceMatching => CParameter::EnableLongDistanceMatching(value != 0), - ZSTD_c_ldmHashLog => CParameter::LdmHashLog(value as u32), - ZSTD_c_ldmMinMatch => CParameter::LdmMinMatch(value as u32), - ZSTD_c_ldmBucketSizeLog => CParameter::LdmBucketSizeLog(value as u32), - ZSTD_c_ldmHashRateLog => CParameter::LdmHashRateLog(value as u32), - ZSTD_c_contentSizeFlag => CParameter::ContentSizeFlag(value != 0), - ZSTD_c_checksumFlag => CParameter::ChecksumFlag(value != 0), - ZSTD_c_dictIDFlag => CParameter::DictIdFlag(value != 0), - ZSTD_c_nbWorkers => CParameter::NbWorkers(value as u32), - ZSTD_c_jobSize => CParameter::JobSize(value as u32), - ZSTD_c_overlapLog => CParameter::OverlapSizeLog(value as u32), - _ => { - return Err(vm.new_value_error(format!( - "invalid compression parameter 'unknown parameter (key {param})'" - ))); - } - }; + /// Validate a compression-parameter id and value pair, returning the + /// parameter constant on success. Used by the compressor's `options=` + /// constructor argument. The value is only validated here for the + /// strategy parameter (which libzstd does not bounds-check on its own in + /// a way CPython can rely on); everything else is guarded by the bounds + /// pre-validation in `apply_options` and passed through to libzstd + /// untouched, like CPython does. + fn cparameter_from_int( + param: i32, + value: i32, + vm: &VirtualMachine, + ) -> PyResult { + let p = c_param_enum(param).ok_or_else(|| { + vm.new_value_error(format!( + "invalid compression parameter 'unknown parameter (key {param})'" + )) + })?; + if param == ZSTD_c_strategy && strategy_from_int(value).is_none() { + return Err(new_zstd_error( + format!("invalid strategy value: {value}"), + vm, + )); + } Ok(p) } - /// Map a decompression-parameter id and integer value to a `DParameter`. - /// Used by the decompressor's `options=` constructor argument. - fn dparameter_from_int(param: i32, value: i32, vm: &VirtualMachine) -> PyResult { - match param { - ZSTD_d_windowLogMax => Ok(DParameter::WindowLogMax(value as u32)), - _ => Err(vm.new_value_error(format!( + /// Validate a decompression-parameter id, returning the parameter + /// constant on success. Used by the decompressor's `options=` constructor + /// argument. + fn dparameter_from_int(param: i32, vm: &VirtualMachine) -> PyResult { + d_param_enum(param).ok_or_else(|| { + vm.new_value_error(format!( "invalid decompression parameter 'unknown parameter (key {param})'" - ))), - } + )) + }) } /// Build the CPython-compatible " parameter 'name' received an @@ -312,21 +322,23 @@ mod _zstd { /// parameter id. Returns `None` if the id is not recognized (callers /// validate the id separately). fn lookup_param_bounds(param: i32, is_compress: bool) -> Option<(i32, i32)> { - // SAFETY: `ZSTD_*Param_getBounds` reads no memory beyond the enum - // discriminant; the helpers above validated that `param` maps to a - // real enum variant. - let bounds = if is_compress { + // The helpers above validated that `param` maps to a real parameter + // constant before `ZSTD_*Param_getBounds` inspects it. The two + // functions return distinct (but layout-identical) `ZSTD_bounds` + // types, so destructure into a tuple right away. + let (error, lo, hi) = if is_compress { let p = c_param_enum(param)?; - unsafe { zstd_sys::ZSTD_cParam_getBounds(p) } + let b = ZSTD_cParam_getBounds(p); + (b.error, b.lowerBound, b.upperBound) } else { let p = d_param_enum(param)?; - unsafe { zstd_sys::ZSTD_dParam_getBounds(p) } + let b = ZSTD_dParam_getBounds(p); + (b.error, b.lowerBound, b.upperBound) }; - // SAFETY: ZSTD_isError just inspects the bounds.error integer. - if unsafe { zstd_sys::ZSTD_isError(bounds.error) } != 0 { + if ZSTD_isError(error) != 0 { return None; } - Some((bounds.lowerBound, bounds.upperBound)) + Some((lo, hi)) } /// Map a parameter integer id back to the Python-visible enum member @@ -365,21 +377,21 @@ mod _zstd { } /// Map a strategy integer (as exposed via the `Strategy` IntEnum) back - /// to the underlying `ZSTD_strategy` C enum value. Done via an explicit + /// to the underlying `ZSTD_strategy` C value. Done via an explicit /// match for the same reason as [`c_param_enum`]: an untrusted int - /// might not correspond to any real enum variant. - fn strategy_from_int(v: i32) -> Option { - use zstd_sys::ZSTD_strategy as S; + /// might not correspond to any real strategy value. + fn strategy_from_int(v: i32) -> Option { + use libzstd_rs_sys::lib::zstd as s; Some(match v { - ZSTD_fast => S::ZSTD_fast, - ZSTD_dfast => S::ZSTD_dfast, - ZSTD_greedy => S::ZSTD_greedy, - ZSTD_lazy => S::ZSTD_lazy, - ZSTD_lazy2 => S::ZSTD_lazy2, - ZSTD_btlazy2 => S::ZSTD_btlazy2, - ZSTD_btopt => S::ZSTD_btopt, - ZSTD_btultra => S::ZSTD_btultra, - ZSTD_btultra2 => S::ZSTD_btultra2, + ZSTD_fast => s::ZSTD_fast, + ZSTD_dfast => s::ZSTD_dfast, + ZSTD_greedy => s::ZSTD_greedy, + ZSTD_lazy => s::ZSTD_lazy, + ZSTD_lazy2 => s::ZSTD_lazy2, + ZSTD_btlazy2 => s::ZSTD_btlazy2, + ZSTD_btopt => s::ZSTD_btopt, + ZSTD_btultra => s::ZSTD_btultra, + ZSTD_btultra2 => s::ZSTD_btultra2, _ => return None, }) } @@ -462,12 +474,16 @@ mod _zstd { fn py_new(_cls: &Py, args: Self::Args, vm: &VirtualMachine) -> PyResult { let dict_content = args.dict_content.with_ref(|b| b.to_vec()); - // libzstd's `get_dict_id_from_dict` returns `None` either when + // libzstd's `ZSTD_getDictID_fromDict` returns 0 either when // the content is too small to contain a valid header or when it // does not carry the dictionary magic. Both are runtime errors // when `is_raw=False`, matching CPython's behavior of raising // `ValueError` on a non-conformant dictionary. - let parsed_id = zstd_safe::get_dict_id_from_dict(&dict_content).map_or(0, |n| n.get()); + // SAFETY: `dict_content` is a live slice for the duration of the + // call. + let parsed_id = unsafe { + ZSTD_getDictID_fromDict(dict_content.as_ptr().cast(), dict_content.len()) + }; if !args.is_raw && parsed_id == 0 { return Err(vm.new_value_error( "ZSTD_DICT_MAGIC_NUMBER not found, dict_content cannot be a 'raw content' \ @@ -549,6 +565,150 @@ mod _zstd { const COMP_MODE_FLUSH_BLOCK: i32 = 1; const COMP_MODE_FLUSH_FRAME: i32 = 2; + /// Owning wrapper around a raw libzstd compression context. Frees the + /// context with `ZSTD_freeCCtx` on drop. + struct CCtx(*mut ZSTD_CCtx); + + // SAFETY: the wrapped context is only ever accessed while holding the + // owning pyclass's `PyMutex` (`CompressorState`), so no concurrent + // access can occur. + unsafe impl Send for CCtx {} + + impl CCtx { + fn create() -> Self { + // SAFETY: `ZSTD_createCCtx` has no preconditions; it returns NULL + // only on allocation failure. + let ptr = unsafe { ZSTD_createCCtx() }; + assert!(!ptr.is_null(), "ZSTD_createCCtx failed"); + Self(ptr) + } + + /// Forward a `(parameter, value)` pair to libzstd, passing the value + /// straight through like CPython does. The error code is returned on + /// failure so the caller can pick the right error mapping. + fn set_parameter(&mut self, param: ZSTD_cParameter, value: c_int) -> Result<(), usize> { + // SAFETY: `self.0` is a live context owned by this wrapper. + let code = unsafe { ZSTD_CCtx_setParameter(self.0, param, value) }; + if ZSTD_isError(code) != 0 { + Err(code) + } else { + Ok(()) + } + } + + fn set_pledged_src_size(&mut self, pledged: Option) -> Result<(), usize> { + // libzstd represents "unknown" as `ZSTD_CONTENTSIZE_UNKNOWN`. + let pledged = pledged.map_or(ZSTD_CONTENTSIZE_UNKNOWN, |v| v); + // SAFETY: `self.0` is a live context owned by this wrapper. + let code = unsafe { ZSTD_CCtx_setPledgedSrcSize(self.0, pledged) }; + if ZSTD_isError(code) != 0 { + Err(code) + } else { + Ok(()) + } + } + } + + impl Drop for CCtx { + fn drop(&mut self) { + // SAFETY: `self.0` was created by `ZSTD_createCCtx` and is freed + // exactly once, here. + unsafe { ZSTD_freeCCtx(self.0) }; + } + } + + /// Owning wrapper around a raw libzstd decompression context. Frees the + /// context with `ZSTD_freeDCtx` on drop. + struct DCtx(*mut ZSTD_DCtx); + + // SAFETY: same reasoning as `CCtx`: access is serialized by the owning + // pyclass's `PyMutex` (`DecompressorState`). + unsafe impl Send for DCtx {} + + impl DCtx { + fn create() -> Self { + // SAFETY: `ZSTD_createDCtx` has no preconditions; it returns NULL + // only on allocation failure. + let ptr = unsafe { ZSTD_createDCtx() }; + assert!(!ptr.is_null(), "ZSTD_createDCtx failed"); + Self(ptr) + } + + /// See [`CCtx::set_parameter`]. + fn set_parameter(&mut self, param: ZSTD_dParameter, value: c_int) -> Result<(), usize> { + // SAFETY: `self.0` is a live context owned by this wrapper. + let code = unsafe { ZSTD_DCtx_setParameter(self.0, param, value) }; + if ZSTD_isError(code) != 0 { + Err(code) + } else { + Ok(()) + } + } + } + + impl Drop for DCtx { + fn drop(&mut self) { + // SAFETY: `self.0` was created by `ZSTD_createDCtx` and is freed + // exactly once, here. + unsafe { ZSTD_freeDCtx(self.0) }; + } + } + + /// Owning wrapper around a digested compression dictionary. Frees the + /// dictionary with `ZSTD_freeCDict` on drop. + struct CDict(*mut ZSTD_CDict); + + // SAFETY: the dictionary is only referenced by the `CCtx` stored + // alongside it in `CompressorState`, which the pyclass's `PyMutex` + // serializes access to. + unsafe impl Send for CDict {} + + impl CDict { + /// Build a digested dictionary; libzstd copies the bytes into its own + /// storage. Returns `None` when libzstd rejects the content (NULL), + /// e.g. for a corrupted dictionary. + fn try_create(bytes: &[u8], level: c_int) -> Option { + // SAFETY: `bytes` is a live slice for the duration of the call, + // and `ZSTD_createCDict` copies what it needs. + let ptr = unsafe { ZSTD_createCDict(bytes.as_ptr().cast(), bytes.len(), level) }; + if ptr.is_null() { None } else { Some(Self(ptr)) } + } + } + + impl Drop for CDict { + fn drop(&mut self) { + // SAFETY: `self.0` was created by `ZSTD_createCDict` and is freed + // exactly once, here. + unsafe { ZSTD_freeCDict(self.0) }; + } + } + + /// Owning wrapper around a digested decompression dictionary. Frees the + /// dictionary with `ZSTD_freeDDict` on drop. + struct DDict(*mut ZSTD_DDict); + + // SAFETY: same reasoning as `CDict` (guarded by the decompressor's + // `PyMutex` via `DecompressorState`). + unsafe impl Send for DDict {} + + impl DDict { + /// See [`CDict::try_create`]. + fn try_create(bytes: &[u8]) -> Option { + // SAFETY: `bytes` is a live slice for the duration of the call, + // and `ZSTD_createDDict` copies what it needs. + let ptr = unsafe { ZSTD_createDDict(bytes.as_ptr().cast(), bytes.len()) }; + if ptr.is_null() { None } else { Some(Self(ptr)) } + } + } + + impl Drop for DDict { + fn drop(&mut self) { + // SAFETY: `self.0` was created by `ZSTD_createDDict` and is freed + // exactly once, here. + unsafe { ZSTD_freeDDict(self.0) }; + } + } + /// Internal state of a `ZstdCompressor`. Holds the libzstd context, the /// last mode used (for `last_mode` and `set_pledged_input_size` /// validation), and the dictionary handles that the context may reference @@ -557,12 +717,12 @@ mod _zstd { /// `PyRef` go away afterwards, which is the safe order for /// teardown. struct CompressorState { - cctx: CCtx<'static>, + cctx: CCtx, /// Cached digested dictionary. The CCtx references this via - /// `ref_cdict`, so it must outlive the CCtx (handled by Rust's - /// field drop order: `cctx` drops first). - _cdict: Option>, - /// Keeps the ZstdDict's bytes alive for `ref_prefix` mode. + /// `ZSTD_CCtx_refCDict`, so it must outlive the CCtx (handled by + /// Rust's field drop order: `cctx` drops first). + _cdict: Option, + /// Keeps the ZstdDict's bytes alive for `ZSTD_CCtx_refPrefix` mode. _dict: Option>, last_mode: i32, } @@ -592,14 +752,11 @@ mod _zstd { /// Translate the public `mode` int to the libzstd `ZSTD_EndDirective` /// the streaming API takes. - fn end_directive_from_mode( - mode: i32, - vm: &VirtualMachine, - ) -> PyResult { + fn end_directive_from_mode(mode: i32, vm: &VirtualMachine) -> PyResult { match mode { - COMP_MODE_CONTINUE => Ok(zstd_sys::ZSTD_EndDirective::ZSTD_e_continue), - COMP_MODE_FLUSH_BLOCK => Ok(zstd_sys::ZSTD_EndDirective::ZSTD_e_flush), - COMP_MODE_FLUSH_FRAME => Ok(zstd_sys::ZSTD_EndDirective::ZSTD_e_end), + COMP_MODE_CONTINUE => Ok(ZSTD_e_continue), + COMP_MODE_FLUSH_BLOCK => Ok(ZSTD_e_flush), + COMP_MODE_FLUSH_FRAME => Ok(ZSTD_e_end), _ => Err(vm.new_value_error(format!( "mode argument wrong value, it should be one of \ ZstdCompressor.CONTINUE ({COMP_MODE_CONTINUE}), \ @@ -621,11 +778,11 @@ mod _zstd { return Err(vm.new_type_error("Only one of level or options should be used.")); } - let mut cctx = CCtx::<'static>::create(); + let mut cctx = CCtx::create(); if let Some(level_obj) = level_opt { let level = parse_compression_level(&level_obj, vm)?; - cctx.set_parameter(CParameter::CompressionLevel(level)) + cctx.set_parameter(ZSTD_cParameter::ZSTD_c_compressionLevel, level) .map_err(|_| param_value_error_for(ZSTD_c_compressionLevel, level, true, vm))?; } @@ -708,25 +865,25 @@ mod _zstd { /// Trait wrapper over `CCtx::set_parameter` and `DCtx::set_parameter` so /// `apply_options` can drive either context without duplicated code. - /// `set_parameter` translates the (id, value) pair to the appropriate - /// `CParameter` / `DParameter` enum variant, then forwards to libzstd. + /// `apply` validates the (id, value) pair via `cparameter_from_int` / + /// `dparameter_from_int`, then forwards to libzstd. trait ParamSetter { fn apply(&mut self, param: i32, value: i32, vm: &VirtualMachine) -> PyResult<()>; } - impl ParamSetter for CCtx<'static> { + impl ParamSetter for CCtx { fn apply(&mut self, param: i32, value: i32, vm: &VirtualMachine) -> PyResult<()> { let p = cparameter_from_int(param, value, vm)?; - self.set_parameter(p) + self.set_parameter(p, value) .map_err(|_| param_value_error_for(param, value, true, vm))?; Ok(()) } } - impl ParamSetter for DCtx<'static> { + impl ParamSetter for DCtx { fn apply(&mut self, param: i32, value: i32, vm: &VirtualMachine) -> PyResult<()> { - let p = dparameter_from_int(param, value, vm)?; - self.set_parameter(p) + let p = dparameter_from_int(param, vm)?; + self.set_parameter(p, value) .map_err(|_| param_value_error_for(param, value, false, vm))?; Ok(()) } @@ -736,46 +893,95 @@ mod _zstd { /// and decompressor consume a dictionary: the name of the type that /// appears in error messages, the eager-validation constructor for the /// digested variant, and the three ways of attaching it to the context. - trait DictLoader<'a> { + /// The `Err` payload of each method is a libzstd error code. + trait DictLoader { type Digested; const KIND_NAME: &'static str; fn try_create_digested(bytes: &[u8]) -> Option; - fn ref_digested(&mut self, dict: &Self::Digested) -> zstd_safe::SafeResult; - fn load_undigested(&mut self, bytes: &[u8]) -> zstd_safe::SafeResult; - fn ref_prefix_static(&mut self, bytes: &'static [u8]) -> zstd_safe::SafeResult; + fn ref_digested(&mut self, dict: &Self::Digested) -> Result<(), usize>; + fn load_undigested(&mut self, bytes: &[u8]) -> Result<(), usize>; + fn ref_prefix(&mut self, bytes: &[u8]) -> Result<(), usize>; } - impl DictLoader<'static> for CCtx<'static> { - type Digested = zstd_safe::CDict<'static>; + impl DictLoader for CCtx { + type Digested = CDict; const KIND_NAME: &'static str = "ZSTD_CDict"; fn try_create_digested(bytes: &[u8]) -> Option { - zstd_safe::CDict::try_create(bytes, ZSTD_CLEVEL_DEFAULT) + CDict::try_create(bytes, libzstd_rs_sys::ZSTD_CLEVEL_DEFAULT) } - fn ref_digested(&mut self, dict: &Self::Digested) -> zstd_safe::SafeResult { - self.ref_cdict(dict) + fn ref_digested(&mut self, dict: &Self::Digested) -> Result<(), usize> { + // SAFETY: both handles are live; `load_dict`'s safety contract + // keeps `dict` alive at least as long as `self`. + let code = unsafe { ZSTD_CCtx_refCDict(self.0, dict.0.cast_const()) }; + if ZSTD_isError(code) != 0 { + Err(code) + } else { + Ok(()) + } } - fn load_undigested(&mut self, bytes: &[u8]) -> zstd_safe::SafeResult { - self.load_dictionary(bytes) + fn load_undigested(&mut self, bytes: &[u8]) -> Result<(), usize> { + // SAFETY: `bytes` is a live slice for the duration of the call; + // its contents are copied into the context. + let code = + unsafe { ZSTD_CCtx_loadDictionary(self.0, bytes.as_ptr().cast(), bytes.len()) }; + if ZSTD_isError(code) != 0 { + Err(code) + } else { + Ok(()) + } } - fn ref_prefix_static(&mut self, bytes: &'static [u8]) -> zstd_safe::SafeResult { - self.ref_prefix(bytes) + fn ref_prefix(&mut self, bytes: &[u8]) -> Result<(), usize> { + // SAFETY: `bytes` is a live slice for the duration of the call. + // libzstd keeps a raw pointer to the bytes afterwards; + // `load_dict`'s safety contract keeps them alive at least as + // long as `self`. + let code = unsafe { ZSTD_CCtx_refPrefix(self.0, bytes.as_ptr().cast(), bytes.len()) }; + if ZSTD_isError(code) != 0 { + Err(code) + } else { + Ok(()) + } } } - impl DictLoader<'static> for DCtx<'static> { - type Digested = zstd_safe::DDict<'static>; + impl DictLoader for DCtx { + type Digested = DDict; const KIND_NAME: &'static str = "ZSTD_DDict"; fn try_create_digested(bytes: &[u8]) -> Option { - zstd_safe::DDict::try_create(bytes) + DDict::try_create(bytes) } - fn ref_digested(&mut self, dict: &Self::Digested) -> zstd_safe::SafeResult { - self.ref_ddict(dict) + fn ref_digested(&mut self, dict: &Self::Digested) -> Result<(), usize> { + // SAFETY: both handles are live; `load_dict`'s safety contract + // keeps `dict` alive at least as long as `self`. + let code = unsafe { ZSTD_DCtx_refDDict(self.0, dict.0.cast_const()) }; + if ZSTD_isError(code) != 0 { + Err(code) + } else { + Ok(()) + } } - fn load_undigested(&mut self, bytes: &[u8]) -> zstd_safe::SafeResult { - self.load_dictionary(bytes) + fn load_undigested(&mut self, bytes: &[u8]) -> Result<(), usize> { + // SAFETY: `bytes` is a live slice for the duration of the call; + // its contents are copied into the context. + let code = + unsafe { ZSTD_DCtx_loadDictionary(self.0, bytes.as_ptr().cast(), bytes.len()) }; + if ZSTD_isError(code) != 0 { + Err(code) + } else { + Ok(()) + } } - fn ref_prefix_static(&mut self, bytes: &'static [u8]) -> zstd_safe::SafeResult { - self.ref_prefix(bytes) + fn ref_prefix(&mut self, bytes: &[u8]) -> Result<(), usize> { + // SAFETY: `bytes` is a live slice for the duration of the call. + // libzstd keeps a raw pointer to the bytes afterwards; + // `load_dict`'s safety contract keeps them alive at least as + // long as `self`. + let code = unsafe { ZSTD_DCtx_refPrefix(self.0, bytes.as_ptr().cast(), bytes.len()) }; + if ZSTD_isError(code) != 0 { + Err(code) + } else { + Ok(()) + } } } @@ -803,7 +1009,7 @@ mod _zstd { /// In `undigested` mode the bytes are copied into `ctx`, so neither /// return value carries a safety obligation — but the caller cannot tell /// the modes apart, so it must keep both alive regardless. - unsafe fn load_dict>( + unsafe fn load_dict( ctx: &mut L, dict_obj: Option, vm: &VirtualMachine, @@ -825,17 +1031,7 @@ mod _zstd { let mut digested = None; match marker { DICT_TYPE_PREFIX => { - // SAFETY: we extend `dict_bytes`' lifetime to `'static` only - // to thread it through `ref_prefix_static`'s signature. - // `load_dict`'s own safety contract requires the caller to - // keep the returned `PyRef` (which owns - // `dict_bytes`) alive at least as long as `ctx`, so the - // raw pointer libzstd stores stays valid for the required - // window. - let static_bytes: &'static [u8] = - unsafe { core::slice::from_raw_parts(dict_bytes.as_ptr(), dict_bytes.len()) }; - ctx.ref_prefix_static(static_bytes) - .map_err(|_| bad_dict_err())?; + ctx.ref_prefix(dict_bytes).map_err(|_| bad_dict_err())?; } DICT_TYPE_DIGESTED => { // Build the digested dict eagerly so a corrupted dictionary @@ -861,7 +1057,7 @@ mod _zstd { /// here, both invariants `load_dict` documents become structural and a /// safe-Rust caller cannot split the pieces apart. fn build_compressor_state( - mut cctx: CCtx<'static>, + mut cctx: CCtx, dict_obj: Option, last_mode: i32, vm: &VirtualMachine, @@ -874,7 +1070,7 @@ mod _zstd { // and `_dict` (prefix mode) are freed. `CompressorState` is private // to this module and is never destructured, so no safe caller can // reorder the drops. - let (cdict, dict) = unsafe { load_dict::>(&mut cctx, dict_obj, vm) }?; + let (cdict, dict) = unsafe { load_dict::(&mut cctx, dict_obj, vm) }?; Ok(CompressorState { cctx, _cdict: cdict, @@ -887,12 +1083,12 @@ mod _zstd { /// [`build_compressor_state`] for the safety reasoning; /// `DecompressorState`'s field order plays the same role here. fn build_decompressor_state( - mut dctx: DCtx<'static>, + mut dctx: DCtx, dict_obj: Option, vm: &VirtualMachine, ) -> PyResult { // SAFETY: see [`build_compressor_state`]. - let (ddict, dict) = unsafe { load_dict::>(&mut dctx, dict_obj, vm) }?; + let (ddict, dict) = unsafe { load_dict::(&mut dctx, dict_obj, vm) }?; Ok(DecompressorState { dctx, _ddict: ddict, @@ -904,13 +1100,14 @@ mod _zstd { }) } - /// Drive `compress_stream2` until the input is fully consumed and, for - /// flush/end directives, the internal buffers report zero remaining bytes. - /// Grows the output `Vec` by `CStreamOutSize` chunks as needed. + /// Drive `ZSTD_compressStream2` until the input is fully consumed and, + /// for flush/end directives, the internal buffers report zero remaining + /// bytes. Grows the output `Vec` by `ZSTD_CStreamOutSize` chunks as + /// needed. fn do_compress( state: &mut CompressorState, data: &[u8], - end_op: zstd_sys::ZSTD_EndDirective, + end_op: ZSTD_EndDirective, vm: &VirtualMachine, ) -> PyResult> { // Release the GIL for the duration of the compression loop. Safety: @@ -921,21 +1118,41 @@ mod _zstd { // and the output `Vec` is local to this function. No Python object // access happens inside the closure — error codes are surfaced as // `usize` and converted into exceptions after re-attaching. - let is_end = end_op != zstd_sys::ZSTD_EndDirective::ZSTD_e_continue; - let chunk_size = CCtx::out_size().max(1); + let is_end = end_op != ZSTD_e_continue; + let chunk_size = ZSTD_CStreamOutSize().max(1); let result: Result, usize> = vm.allow_threads(|| { - let mut output = Vec::new(); - let mut input = InBuffer::around(data); + let mut output: Vec = Vec::new(); + let mut input = ZSTD_inBuffer { + src: data.as_ptr().cast(), + size: data.len(), + pos: 0, + }; loop { let prev_len = output.len(); output.reserve(chunk_size); let remaining = { - let mut out_buf = OutBuffer::around_pos(&mut output, prev_len); - state - .cctx - .compress_stream2(&mut out_buf, &mut input, end_op) - }?; - let consumed_all = input.pos == input.src.len(); + let spare = output.spare_capacity_mut(); + let mut out_buf = ZSTD_outBuffer { + dst: spare.as_mut_ptr().cast(), + size: spare.len(), + pos: 0, + }; + // SAFETY: `state.cctx` is a live context; `out_buf` + // covers the Vec's spare capacity, which is valid for + // `spare.len()` bytes; `input` points to `data`, which + // outlives the call. + let code = unsafe { + ZSTD_compressStream2(state.cctx.0, &mut out_buf, &mut input, end_op) + }; + if ZSTD_isError(code) != 0 { + return Err(code); + } + // SAFETY: libzstd wrote `out_buf.pos` bytes + // (`out_buf.pos <= spare.len()`) into the spare capacity. + unsafe { output.set_len(prev_len + out_buf.pos) }; + code + }; + let consumed_all = input.pos == input.size; // Stop when input is fully consumed and, for flush/end // directives, libzstd reports that all internal buffers have // been drained (remaining == 0). Otherwise loop; the next @@ -986,8 +1203,8 @@ mod _zstd { // likewise converts the argument before touching the compressor. // // Python passes `None` to mean "unknown"; libzstd represents that - // internally as `ZSTD_CONTENTSIZE_UNKNOWN` (`u64::MAX`), and - // `zstd_safe` translates `None` accordingly. libzstd also reserves + // internally as `ZSTD_CONTENTSIZE_UNKNOWN` (`u64::MAX`), which + // `set_pledged_src_size` maps `None` to. libzstd also reserves // `ZSTD_CONTENTSIZE_ERROR` (`u64::MAX - 1`), so a concrete size must // be strictly less than that; reject anything else up front so // callers see the documented `ValueError`, not a libzstd-level error. @@ -1072,9 +1289,10 @@ mod _zstd { /// the `dctx` is freed first and must give up its internal pointers /// before any referenced `DDict`/`PyRef` is dropped. struct DecompressorState { - dctx: DCtx<'static>, - /// Cached decompression dictionary referenced by the DCtx. - _ddict: Option>, + dctx: DCtx, + /// Cached decompression dictionary referenced by the DCtx via + /// `ZSTD_DCtx_refDDict`. + _ddict: Option, _dict: Option>, eof: bool, needs_input: bool, @@ -1113,7 +1331,7 @@ mod _zstd { let dict_opt = args.zstd_dict.flatten(); let options_opt = args.options.flatten(); - let mut dctx = DCtx::<'static>::create(); + let mut dctx = DCtx::create(); if let Some(options_obj) = options_opt { apply_options(&mut dctx, options_obj, false, vm)?; @@ -1162,18 +1380,22 @@ mod _zstd { alloc::borrow::Cow::Owned(combined) }; - let chunk_size = DCtx::out_size().max(1); + let chunk_size = ZSTD_DStreamOutSize().max(1); // Release the GIL for the streaming loop. Safety: see `do_compress`; // the closure captures only Rust-owned buffers and `&mut state` (held // under the decompressor's `PyMutex`), and surfaces error codes as // `usize` so we can build the exception after re-attaching. let loop_result: Result<(Vec, bool, usize), usize> = vm.allow_threads(|| { - let mut input = InBuffer::around(&work_data); + let mut input = ZSTD_inBuffer { + src: work_data.as_ptr().cast(), + size: work_data.len(), + pos: 0, + }; let mut output: Vec = Vec::new(); // Reusable scratch buffer for each decompress_stream call. We need // an exact-size output buffer because `Vec::reserve` may - // over-allocate; `OutBuffer` reports the full Vec capacity to - // libzstd, which would then happily write past `max_length`. + // over-allocate; reporting the full Vec capacity to libzstd would + // let it write past `max_length`. let mut scratch: Vec = vec![0u8; chunk_size]; let mut hit_max = false; let mut iteration = 0usize; @@ -1195,41 +1417,46 @@ mod _zstd { Some(maxl) => (maxl - output.len()).min(chunk_size), None => chunk_size, }; - let result; + let code; let written; { let slot = &mut scratch[..grow]; - let mut out_buf = OutBuffer::around(slot as &mut [u8]); - result = state.dctx.decompress_stream(&mut out_buf, &mut input); - written = out_buf.pos(); + let mut out_buf = ZSTD_outBuffer { + dst: slot.as_mut_ptr().cast(), + size: slot.len(), + pos: 0, + }; + // SAFETY: `state.dctx` is a live context; `out_buf` + // covers `slot`; `input` points into `work_data`, which + // outlives the call. + code = unsafe { ZSTD_decompressStream(state.dctx.0, &mut out_buf, &mut input) }; + written = out_buf.pos; } output.extend_from_slice(&scratch[..written]); - match result { - Ok(0) => { - // Frame fully decompressed; the decompressor is at EOF. - state.eof = true; - break Ok(()); - } - Ok(_) => { - let output_was_full = written == grow; - let input_consumed = input.pos == input.src.len(); - - if let Some(maxl) = max_length - && output.len() >= maxl - && iteration > 1 - { - hit_max = true; - break Ok(()); - } - - // Input is gone and libzstd had room to write but did - // not, which means the frame is incomplete and the - // caller has to supply more input. - if input_consumed && !output_was_full { - break Ok(()); - } - } - Err(code) => break Err(code), + if ZSTD_isError(code) != 0 { + break Err(code); + } + if code == 0 { + // Frame fully decompressed; the decompressor is at EOF. + state.eof = true; + break Ok(()); + } + let output_was_full = written == grow; + let input_consumed = input.pos == input.size; + + if let Some(maxl) = max_length + && output.len() >= maxl + && iteration > 1 + { + hit_max = true; + break Ok(()); + } + + // Input is gone and libzstd had room to write but did + // not, which means the frame is incomplete and the + // caller has to supply more input. + if input_consumed && !output_was_full { + break Ok(()); } }; outcome.map(|()| (output, hit_max, input.pos)) @@ -1318,14 +1545,17 @@ mod _zstd { #[pyfunction] fn get_frame_size(frame_buffer: ArgBytesLike, vm: &VirtualMachine) -> PyResult { let buf = frame_buffer.with_ref(|b| b.to_vec()); - zstd_safe::find_frame_compressed_size(&buf).map_err(|_| { - new_zstd_error( + // SAFETY: `buf` is a live slice for the duration of the call. + let size = unsafe { ZSTD_findFrameCompressedSize(buf.as_ptr().cast(), buf.len()) }; + if ZSTD_isError(size) != 0 { + return Err(new_zstd_error( "Error when finding the compressed size of a Zstandard frame. \ Ensure the frame_buffer argument starts from the beginning of a frame, \ and its length not less than this complete frame.", vm, - ) - }) + )); + } + Ok(size) } #[pyfunction] @@ -1334,19 +1564,23 @@ mod _zstd { vm: &VirtualMachine, ) -> PyResult<(PyObjectRef, u32)> { let buf = frame_buffer.with_ref(|b| b.to_vec()); - let content_size = zstd_safe::get_frame_content_size(&buf).map_err(|_| { - new_zstd_error( + // SAFETY: `buf` is a live slice for the duration of the call. + let content_size = unsafe { ZSTD_getFrameContentSize(buf.as_ptr().cast(), buf.len()) }; + if content_size == ZSTD_CONTENTSIZE_ERROR { + return Err(new_zstd_error( "Error when getting information from the header of a Zstandard frame. \ Ensure the frame_buffer argument starts from the beginning of a frame, \ and its length not less than the frame header (6~18 bytes).", vm, - ) - })?; - let content_size_obj: PyObjectRef = match content_size { - Some(n) => vm.ctx.new_int(n).into(), - None => vm.ctx.none(), + )); + } + let content_size_obj: PyObjectRef = if content_size == ZSTD_CONTENTSIZE_UNKNOWN { + vm.ctx.none() + } else { + vm.ctx.new_int(content_size).into() }; - let dict_id = zstd_safe::get_dict_id_from_frame(&buf).map_or(0, |n| n.get()); + // SAFETY: `buf` is a live slice for the duration of the call. + let dict_id = unsafe { ZSTD_getDictID_fromFrame(buf.as_ptr().cast(), buf.len()) }; Ok((content_size_obj, dict_id)) } @@ -1405,9 +1639,9 @@ mod _zstd { } /// Sum the per-sample sizes and check they exactly cover `expected_total`, - /// rejecting overflow. A safe-fn interface that calls into libzstd with - /// these sizes must not let a wrapping sum sneak past the equality check, - /// since libzstd would then read past the samples buffer. + /// rejecting overflow. Any call into libzstd with these sizes must not + /// let a wrapping sum sneak past the equality check, since libzstd would + /// then read past the samples buffer. fn check_sample_sizes_match( sizes: &[usize], expected_total: usize, @@ -1433,8 +1667,27 @@ mod _zstd { let sizes = parse_sample_sizes(args.samples_sizes, vm)?; check_sample_sizes_match(&sizes, samples_buffer.len(), vm)?; let mut dict_buffer: Vec = Vec::with_capacity(dict_size); - zstd_safe::train_from_buffer(&mut dict_buffer, &samples_buffer, &sizes) - .map_err(|c| catch_zstd_error(c, vm))?; + // SAFETY: `dict_buffer`'s spare capacity is valid for writes of up to + // `dict_size` bytes; the samples and sizes buffers are live slices + // for the duration of the call. + let written = unsafe { + ZDICT_trainFromBuffer( + dict_buffer.as_mut_ptr().cast(), + dict_size, + samples_buffer.as_ptr().cast(), + sizes.as_ptr(), + sizes.len() as u32, + ) + }; + if ZDICT_isError(written) != 0 { + // SAFETY: `ZDICT_getErrorName` returns a pointer to a static, + // NUL-terminated error string from libzstd's error table. + let name = unsafe { CStr::from_ptr(ZDICT_getErrorName(written)) }; + return Err(new_zstd_error(name.to_string_lossy(), vm)); + } + // SAFETY: `ZDICT_trainFromBuffer` wrote `written` (<= dict_size) + // bytes into the buffer's spare capacity. + unsafe { dict_buffer.set_len(written) }; Ok(dict_buffer) } @@ -1459,8 +1712,6 @@ mod _zstd { compression_level: PyObjectRef, } - // `zstd_safe` does not wrap `ZDICT_finalizeDictionary`, so we drop down to - // raw `zstd_sys` FFI for it. #[pyfunction] fn finalize_dict(args: FinalizeDictArgs, vm: &VirtualMachine) -> PyResult> { let dict_size = parse_positive_dict_size(&args.dict_size, vm)?; @@ -1471,7 +1722,7 @@ mod _zstd { check_sample_sizes_match(&sizes, samples_buffer.len(), vm)?; let mut dict_buffer: Vec = vec![0u8; dict_size]; - let params = zstd_sys::ZDICT_params_t { + let params = ZDICT_params_t { compressionLevel: compression_level, notificationLevel: 0, dictID: 0, @@ -1481,7 +1732,7 @@ mod _zstd { // that outlive the FFI call. ZDICT_finalizeDictionary just reads from // the sample/dict buffers and writes into `dict_buffer`. let written = unsafe { - zstd_sys::ZDICT_finalizeDictionary( + ZDICT_finalizeDictionary( dict_buffer.as_mut_ptr() as *mut _, dict_buffer.len(), custom_dict.as_ptr() as *const _, @@ -1492,15 +1743,14 @@ mod _zstd { params, ) }; - // SAFETY: ZDICT_isError just inspects the integer return code. - if unsafe { zstd_sys::ZDICT_isError(written) } != 0 { - // SAFETY: ZDICT_getErrorName returns a static NUL-terminated - // C string from libzstd's internal error table. - let err_ptr = unsafe { zstd_sys::ZDICT_getErrorName(written) }; + if ZDICT_isError(written) != 0 { + let err_ptr = ZDICT_getErrorName(written); let msg = if err_ptr.is_null() { "zstd dictionary finalization failed".to_string() } else { - unsafe { core::ffi::CStr::from_ptr(err_ptr) } + // SAFETY: `ZDICT_getErrorName` returns a pointer to a static, + // NUL-terminated error string from libzstd's error table. + unsafe { CStr::from_ptr(err_ptr) } .to_string_lossy() .into_owned() }; @@ -1533,22 +1783,22 @@ mod _zstd { }; // Validate the id via the same safe enum-lookup helpers used in // `lookup_param_bounds`, then call libzstd directly so we can - // distinguish a libzstd-reported error from our own "unknown". - let bounds = if args.is_compress { + // distinguish a libzstd-reported error from our own "unknown". The + // two `getBounds` functions return distinct (but layout-identical) + // `ZSTD_bounds` types, so destructure into a tuple right away. + let (error, lo, hi) = if args.is_compress { let p = c_param_enum(args.parameter).ok_or_else(unknown)?; - // SAFETY: `c_param_enum` returned `Some`, so `p` is a real - // `ZSTD_cParameter` discriminant. - unsafe { zstd_sys::ZSTD_cParam_getBounds(p) } + let b = ZSTD_cParam_getBounds(p); + (b.error, b.lowerBound, b.upperBound) } else { let p = d_param_enum(args.parameter).ok_or_else(unknown)?; - // SAFETY: same as above. - unsafe { zstd_sys::ZSTD_dParam_getBounds(p) } + let b = ZSTD_dParam_getBounds(p); + (b.error, b.lowerBound, b.upperBound) }; - // SAFETY: ZSTD_isError just inspects the integer error code. - if unsafe { zstd_sys::ZSTD_isError(bounds.error) } != 0 { - return Err(catch_zstd_error(bounds.error, vm)); + if ZSTD_isError(error) != 0 { + return Err(catch_zstd_error(error, vm)); } - Ok((bounds.lowerBound, bounds.upperBound)) + Ok((lo, hi)) } // Register the `CompressionParameter` / `DecompressionParameter` enum From 8fe59b8e930a524231d3260909469e029866868d Mon Sep 17 00:00:00 2001 From: James David Clarke Date: Sun, 9 Aug 2026 17:52:07 +0100 Subject: [PATCH 14/20] Fix data loss in ZstdDecompressor.decompress(max_length=0) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The zero-cap probe iteration handed libzstd a 1-byte buffer and the post-loop truncation discarded whatever it emitted, while the input counter had already advanced past the bytes that produced it — every decompress(..., max_length=0) call with pending output silently dropped one byte. The suite never caught it because max_length=0 was only exercised against zero-output (skippable) frames. Match CPython's mechanism instead: probe with a zero-size output buffer, so libzstd consumes input without emitting and never reports frame completion with un-emitted content. Add a regression snippet covering lossless zero-cap probes (content, truncated, and skippable frames), verified byte-exact against CPython. --- crates/stdlib/src/zstd.rs | 31 +++++++----------- extra_tests/snippets/stdlib_zstd.py | 51 +++++++++++++++++++++++++++++ 2 files changed, 63 insertions(+), 19 deletions(-) create mode 100644 extra_tests/snippets/stdlib_zstd.py diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs index 8d9e194bfd6..e7d3613fef0 100644 --- a/crates/stdlib/src/zstd.rs +++ b/crates/stdlib/src/zstd.rs @@ -1403,17 +1403,18 @@ mod _zstd { let outcome = loop { iteration += 1; // Honor `max_length`: stop growing the output buffer once - // we have produced enough. Special-case the first iteration - // when the cap is zero so a zero-output frame (skippable - // frame, empty content frame) can still complete; we hand - // libzstd a 1-byte scratch and discard the byte if it ends - // up writing one. + // we have produced enough. When the cap is zero, hand + // libzstd a zero-size buffer instead: it then consumes + // input without emitting, so zero-output frames (skippable + // frame, empty content frame) still complete while real + // output stays inside libzstd until a later call has room. + // CPython's decompressor uses the same zero-size mechanism. let grow = match max_length { Some(maxl) if output.len() >= maxl && iteration > 1 => { hit_max = true; break Ok(()); } - Some(maxl) if output.len() >= maxl => 1, + Some(maxl) if output.len() >= maxl => 0, Some(maxl) => (maxl - output.len()).min(chunk_size), None => chunk_size, }; @@ -1441,7 +1442,10 @@ mod _zstd { state.eof = true; break Ok(()); } - let output_was_full = written == grow; + // Only meaningful with a non-zero buffer: libzstd had more + // to emit but ran out of room. A zero-size call (the + // max_length == 0 probe) is not "full". + let output_was_full = grow > 0 && written == grow; let input_consumed = input.pos == input.size; if let Some(maxl) = max_length @@ -1462,18 +1466,7 @@ mod _zstd { outcome.map(|()| (output, hit_max, input.pos)) }); - let (mut output, mut hit_max, consumed) = - loop_result.map_err(|c| catch_zstd_error(c, vm))?; - - // If `max_length == 0` opened a courtesy iteration that produced more - // bytes than the caller asked for, truncate. Should not happen with - // the scratch slicing above, but keep the safety net. - if let Some(maxl) = max_length - && output.len() > maxl - { - output.truncate(maxl); - hit_max = true; - } + let (output, hit_max, consumed) = loop_result.map_err(|c| catch_zstd_error(c, vm))?; let remaining = &work_data[consumed..]; diff --git a/extra_tests/snippets/stdlib_zstd.py b/extra_tests/snippets/stdlib_zstd.py new file mode 100644 index 00000000000..58fa742a3d7 --- /dev/null +++ b/extra_tests/snippets/stdlib_zstd.py @@ -0,0 +1,51 @@ +from compression.zstd import ZstdCompressor, ZstdDecompressor + +# basic roundtrip +payload = b"hello world" * 100 +c = ZstdCompressor() +frame = c.compress(payload) + c.flush() +d = ZstdDecompressor() +out = d.decompress(frame) +assert out == payload +assert d.eof +assert not d.needs_input + +# max_length=0 must not drop output bytes: probing with a zero cap has to +# leave every byte available to later drain calls (regression: the probe +# used to emit one byte per call into the void). +d = ZstdDecompressor() +out = d.decompress(frame, 0) +assert out == b"" +assert not d.eof +assert not d.needs_input +rest = d.decompress(b"", len(payload)) +assert rest == payload +assert d.eof + +# two zero-cap probes in a row, then drain; still lossless +d = ZstdDecompressor() +assert d.decompress(frame, 0) == b"" +assert d.decompress(b"", 0) == b"" +assert d.decompress(b"", len(payload)) == payload +assert d.eof + +# zero-cap probe on a truncated frame: no output, not at end +d = ZstdDecompressor() +out = d.decompress(frame[:-5], 0) +assert out == b"" +assert not d.eof +assert not d.needs_input + +# zero-cap probe on a skippable frame completes it +skippable = b"\x50\x2a\x4d\x18\x04\x00\x00\x00abcd" +d = ZstdDecompressor() +assert d.decompress(skippable, 0) == b"" +assert d.eof +assert not d.needs_input +assert d.unused_data == b"" + +# skippable frame followed by trailing bytes lands in unused_data +d = ZstdDecompressor() +assert d.decompress(skippable + b"trail", 0) == b"" +assert d.eof +assert d.unused_data == b"trail" From 3d4825e9f0f4fa8d7cb3a846cd3ffb33cd0b0d81 Mon Sep 17 00:00:00 2001 From: James David Clarke Date: Sun, 9 Aug 2026 18:03:50 +0100 Subject: [PATCH 15/20] Drop duplicate spell-checker line in zstd.rs --- crates/stdlib/src/zstd.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs index e7d3613fef0..69e952e73ce 100644 --- a/crates/stdlib/src/zstd.rs +++ b/crates/stdlib/src/zstd.rs @@ -4,7 +4,6 @@ // spell-checker:ignore srcsize zdict refprefix refcdict refddict pledgedsrcsize getframecontentsize // spell-checker:ignore Zstd Zstandard pylib RFC // spell-checker:ignore CLEVEL zstdmodule cparameter dparameter maxl c2rust -// spell-checker:ignore cctx dctx CCTX DCTX ldm cdict ddict windowlog hashlog chainlog searchlog CLEVEL //! The `_zstd` extension module. Backs the pure-Python `compression.zstd` //! package by exposing the same classes, functions and constants that From 52bea595fcf8a8ac42b692bf999b49e4e1ef20b2 Mon Sep 17 00:00:00 2001 From: James David Clarke Date: Wed, 19 Aug 2026 09:58:51 +0100 Subject: [PATCH 16/20] Address review on _zstd constants, version attrs and dict markers - Derive the compression/decompression parameter ids from libzstd's own `ZSTD_cParameter`/`ZSTD_dParameter` constants instead of restating them as literals. The crate declares both as `#[repr(transparent)]` newtypes over a private `u32` with no accessor, `From` impl or `Deref`, so the id is read back through that guaranteed layout. - Drop the redundant `name = ` on `#[pyattr(once)] fn zstd_version` and `fn zstd_version_number`; the function names already match. - Use `CStr::to_str` for the version string. libzstd returns a fixed ASCII `MAJOR.MINOR.RELEASE` literal, so there is nothing for a lossy conversion to repair. - Replace the `DICT_TYPE_*` int constants with a `DictType` enum carrying the pinned discriminants, decoded at the Python boundary by `DictType::from_marker`. - Explain the non-trivial assertions: why `level_bounds`' `expect` cannot fire, and why a NULL context allocation panics rather than raising `MemoryError`. Co-Authored-By: Claude Opus 5 --- crates/stdlib/src/zstd.rs | 184 ++++++++++++++++++++++++++------------ 1 file changed, 127 insertions(+), 57 deletions(-) diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs index 69e952e73ce..833e4da5cea 100644 --- a/crates/stdlib/src/zstd.rs +++ b/crates/stdlib/src/zstd.rs @@ -54,55 +54,79 @@ mod _zstd { #[pyattr] const ZSTD_CLEVEL_DEFAULT: i32 = libzstd_rs_sys::ZSTD_CLEVEL_DEFAULT; - // Compression parameter identifiers. Values match the `ZSTD_cParameter` - // enum in libzstd, which is what the public `CompressionParameter` IntEnum - // in `Lib/compression/zstd/__init__.py` derives its members from. - // libzstd-rs-sys models the enum as a newtype whose inner value is - // crate-private, so the ids — frozen C ABI values from zstd.h — are - // spelled out as literals here. + /// Read the numeric id back out of a libzstd compression-parameter + /// constant. + /// + /// The ids are needed as plain ints because the module constants below are + /// what `Lib/compression/zstd/__init__.py` turns into the public + /// `CompressionParameter` IntEnum. libzstd-rs-sys declares + /// `ZSTD_cParameter` as `#[repr(transparent)] pub struct + /// ZSTD_cParameter(pub(crate) u32)` and gives it no accessor, `From` impl + /// or `Deref`, so the id is unreachable through its public API. Reading + /// the wrapped `u32` back keeps libzstd the single source of truth for + /// these values instead of restating them here as literals. + const fn c_param_id(param: ZSTD_cParameter) -> i32 { + // SAFETY: `#[repr(transparent)]` guarantees `ZSTD_cParameter` has the + // same size, alignment and validity invariants as the `u32` it wraps, + // so reading one back as a `u32` is always well-defined. + unsafe { core::mem::transmute::(param) as i32 } + } + + /// `ZSTD_dParameter` counterpart of [`c_param_id`]; that type is likewise + /// a `#[repr(transparent)]` newtype over a private `u32`. + const fn d_param_id(param: ZSTD_dParameter) -> i32 { + // SAFETY: as in `c_param_id`. + unsafe { core::mem::transmute::(param) as i32 } + } + + // Compression parameter identifiers, taken straight from libzstd's + // `ZSTD_cParameter` constants. These are the ids the public + // `CompressionParameter` IntEnum in `Lib/compression/zstd/__init__.py` + // derives its members from. #[pyattr] - const ZSTD_c_compressionLevel: i32 = 100; + const ZSTD_c_compressionLevel: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_compressionLevel); #[pyattr] - const ZSTD_c_windowLog: i32 = 101; + const ZSTD_c_windowLog: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_windowLog); #[pyattr] - const ZSTD_c_hashLog: i32 = 102; + const ZSTD_c_hashLog: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_hashLog); #[pyattr] - const ZSTD_c_chainLog: i32 = 103; + const ZSTD_c_chainLog: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_chainLog); #[pyattr] - const ZSTD_c_searchLog: i32 = 104; + const ZSTD_c_searchLog: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_searchLog); #[pyattr] - const ZSTD_c_minMatch: i32 = 105; + const ZSTD_c_minMatch: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_minMatch); #[pyattr] - const ZSTD_c_targetLength: i32 = 106; + const ZSTD_c_targetLength: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_targetLength); #[pyattr] - const ZSTD_c_strategy: i32 = 107; + const ZSTD_c_strategy: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_strategy); #[pyattr] - const ZSTD_c_enableLongDistanceMatching: i32 = 160; + const ZSTD_c_enableLongDistanceMatching: i32 = + c_param_id(ZSTD_cParameter::ZSTD_c_enableLongDistanceMatching); #[pyattr] - const ZSTD_c_ldmHashLog: i32 = 161; + const ZSTD_c_ldmHashLog: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_ldmHashLog); #[pyattr] - const ZSTD_c_ldmMinMatch: i32 = 162; + const ZSTD_c_ldmMinMatch: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_ldmMinMatch); #[pyattr] - const ZSTD_c_ldmBucketSizeLog: i32 = 163; + const ZSTD_c_ldmBucketSizeLog: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_ldmBucketSizeLog); #[pyattr] - const ZSTD_c_ldmHashRateLog: i32 = 164; + const ZSTD_c_ldmHashRateLog: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_ldmHashRateLog); #[pyattr] - const ZSTD_c_contentSizeFlag: i32 = 200; + const ZSTD_c_contentSizeFlag: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_contentSizeFlag); #[pyattr] - const ZSTD_c_checksumFlag: i32 = 201; + const ZSTD_c_checksumFlag: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_checksumFlag); #[pyattr] - const ZSTD_c_dictIDFlag: i32 = 202; + const ZSTD_c_dictIDFlag: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_dictIDFlag); #[pyattr] - const ZSTD_c_nbWorkers: i32 = 400; + const ZSTD_c_nbWorkers: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_nbWorkers); #[pyattr] - const ZSTD_c_jobSize: i32 = 401; + const ZSTD_c_jobSize: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_jobSize); #[pyattr] - const ZSTD_c_overlapLog: i32 = 402; + const ZSTD_c_overlapLog: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_overlapLog); // Decompression parameter identifiers. libzstd only exposes one non- // experimental decompression parameter. #[pyattr] - const ZSTD_d_windowLogMax: i32 = 100; + const ZSTD_d_windowLogMax: i32 = d_param_id(ZSTD_dParameter::ZSTD_d_windowLogMax); // Strategy enum members ordered from fastest to strongest. These power // the `Strategy` IntEnum in `Lib/compression/zstd/__init__.py`. @@ -125,16 +149,21 @@ mod _zstd { #[pyattr] const ZSTD_btultra2: i32 = libzstd_rs_sys::lib::zstd::ZSTD_btultra2 as i32; - #[pyattr(once, name = "zstd_version")] + #[pyattr(once)] fn zstd_version(_vm: &VirtualMachine) -> String { // SAFETY: `ZSTD_versionString` returns a pointer to libzstd's static, // NUL-terminated version string. - unsafe { CStr::from_ptr(ZSTD_versionString()) } - .to_string_lossy() - .into_owned() + let version = unsafe { CStr::from_ptr(ZSTD_versionString()) }; + // The string is a fixed `MAJOR.MINOR.RELEASE` ASCII literal baked into + // libzstd, so it is always valid UTF-8 and there is nothing for a + // lossy conversion to repair. + version + .to_str() + .expect("libzstd version string is ASCII") + .to_owned() } - #[pyattr(once, name = "zstd_version_number")] + #[pyattr(once)] fn zstd_version_number(_vm: &VirtualMachine) -> u32 { ZSTD_versionNumber() } @@ -144,13 +173,35 @@ mod _zstd { ZSTD_DStreamOutSize() } - // Dictionary load type markers. The `ZstdDict.as_*` properties wrap the - // dictionary in a `(zdict, marker)` tuple so the compressor or decompressor - // constructor knows which load mode to apply. Numbering matches CPython's - // `Modules/_zstd/_zstdmodule.h::DictType`. - const DICT_TYPE_DIGESTED: i32 = 0; - const DICT_TYPE_UNDIGESTED: i32 = 1; - const DICT_TYPE_PREFIX: i32 = 2; + /// Which of the three load modes a `zstd_dict=` argument asks for. The + /// `ZstdDict.as_*` properties pair the dictionary with the matching + /// discriminant in a `(zdict, marker)` tuple, and the compressor and + /// decompressor constructors turn that marker back into this enum. + /// + /// The discriminants are pinned because they cross the Python boundary as + /// plain ints: they are the values CPython's + /// `Modules/_zstd/_zstdmodule.h::DictType` uses, and `test_zstd` builds + /// the tuples by hand with those literals. + #[derive(Copy, Clone)] + enum DictType { + Digested = 0, + Undigested = 1, + Prefix = 2, + } + + impl DictType { + /// Decode the int half of a `(zdict, marker)` tuple, which may come + /// straight from user code. `None` for anything outside the three + /// documented markers. + fn from_marker(marker: i32) -> Option { + Some(match marker { + m if m == Self::Digested as i32 => Self::Digested, + m if m == Self::Undigested as i32 => Self::Undigested, + m if m == Self::Prefix as i32 => Self::Prefix, + _ => return None, + }) + } + } #[pyattr(once, name = "ZstdError")] fn zstd_error(vm: &VirtualMachine) -> PyTypeRef { @@ -313,6 +364,11 @@ mod _zstd { /// libzstd silently clamps out-of-range values rather than surfacing /// them as errors. fn level_bounds() -> (i32, i32) { + // `lookup_param_bounds` only returns `None` for an unrecognized id or + // when libzstd flags the bounds as an error. `ZSTD_c_compressionLevel` + // is recognized, and `ZSTD_cParam_getBounds` answers it from + // `ZSTD_minCLevel()`/`ZSTD_maxCLevel()` without ever setting the error + // field, so `None` is unreachable here. lookup_param_bounds(ZSTD_c_compressionLevel, true) .expect("compressionLevel always has valid bounds") } @@ -401,12 +457,12 @@ mod _zstd { fn parse_zstd_dict_arg( obj: PyObjectRef, vm: &VirtualMachine, - ) -> PyResult<(PyRef, i32)> { + ) -> PyResult<(PyRef, DictType)> { // The first downcast clones `obj` because we fall through to the // tuple branch if it fails. The second downcast (the tuple one) is // the last use of `obj`, so we let it move directly. if let Ok(d) = obj.clone().downcast::() { - return Ok((d, DICT_TYPE_DIGESTED)); + return Ok((d, DictType::Digested)); } if let Ok(tuple) = obj.downcast::() { let items = tuple.as_slice(); @@ -433,10 +489,10 @@ mod _zstd { vm.new_type_error("zstd_dict argument should be a ZstdDict object") } })?; - if !(DICT_TYPE_DIGESTED..=DICT_TYPE_PREFIX).contains(&marker) { + let Some(dict_type) = DictType::from_marker(marker) else { return Err(vm.new_type_error("zstd_dict argument should be a ZstdDict object")); - } - return Ok((d, marker)); + }; + return Ok((d, dict_type)); } Err(vm.new_type_error("zstd_dict argument should be a ZstdDict object")) } @@ -538,22 +594,26 @@ mod _zstd { #[pygetset] fn as_digested_dict(zelf: PyRef, vm: &VirtualMachine) -> PyTupleRef { - vm.ctx - .new_tuple(vec![zelf.into(), vm.ctx.new_int(DICT_TYPE_DIGESTED).into()]) + vm.ctx.new_tuple(vec![ + zelf.into(), + vm.ctx.new_int(DictType::Digested as i32).into(), + ]) } #[pygetset] fn as_undigested_dict(zelf: PyRef, vm: &VirtualMachine) -> PyTupleRef { vm.ctx.new_tuple(vec![ zelf.into(), - vm.ctx.new_int(DICT_TYPE_UNDIGESTED).into(), + vm.ctx.new_int(DictType::Undigested as i32).into(), ]) } #[pygetset] fn as_prefix(zelf: PyRef, vm: &VirtualMachine) -> PyTupleRef { - vm.ctx - .new_tuple(vec![zelf.into(), vm.ctx.new_int(DICT_TYPE_PREFIX).into()]) + vm.ctx.new_tuple(vec![ + zelf.into(), + vm.ctx.new_int(DictType::Prefix as i32).into(), + ]) } } @@ -576,8 +636,13 @@ mod _zstd { impl CCtx { fn create() -> Self { // SAFETY: `ZSTD_createCCtx` has no preconditions; it returns NULL - // only on allocation failure. + // only when its `malloc` of the context struct fails. let ptr = unsafe { ZSTD_createCCtx() }; + // A NULL return therefore means the allocator is out of memory, + // which is the one case Rust already handles by aborting, so + // panicking matches the rest of the VM rather than threading a + // `vm` reference through every construction site to raise + // `MemoryError` the way CPython's `_zstd` does. assert!(!ptr.is_null(), "ZSTD_createCCtx failed"); Self(ptr) } @@ -627,8 +692,13 @@ mod _zstd { impl DCtx { fn create() -> Self { // SAFETY: `ZSTD_createDCtx` has no preconditions; it returns NULL - // only on allocation failure. + // only when its `malloc` of the context struct fails. let ptr = unsafe { ZSTD_createDCtx() }; + // A NULL return therefore means the allocator is out of memory, + // which is the one case Rust already handles by aborting, so + // panicking matches the rest of the VM rather than threading a + // `vm` reference through every construction site to raise + // `MemoryError` the way CPython's `_zstd` does. assert!(!ptr.is_null(), "ZSTD_createDCtx failed"); Self(ptr) } @@ -1016,7 +1086,7 @@ mod _zstd { let Some(dict_obj) = dict_obj else { return Ok((None, None)); }; - let (zdict, marker) = parse_zstd_dict_arg(dict_obj, vm)?; + let (zdict, dict_type) = parse_zstd_dict_arg(dict_obj, vm)?; let bad_dict_err = || -> PyBaseExceptionRef { new_zstd_error( format!( @@ -1028,11 +1098,11 @@ mod _zstd { }; let dict_bytes = zdict.dict_content.as_bytes(); let mut digested = None; - match marker { - DICT_TYPE_PREFIX => { + match dict_type { + DictType::Prefix => { ctx.ref_prefix(dict_bytes).map_err(|_| bad_dict_err())?; } - DICT_TYPE_DIGESTED => { + DictType::Digested => { // Build the digested dict eagerly so a corrupted dictionary // surfaces as a `ZstdError` at construction time, not when // the first compress/decompress call runs. @@ -1040,9 +1110,9 @@ mod _zstd { ctx.ref_digested(&d).map_err(|_| bad_dict_err())?; digested = Some(d); } - _ => { - // Undigested: copy the bytes into the context. Validation - // happens lazily at the first stream call in this mode. + DictType::Undigested => { + // Copy the bytes into the context. Validation happens lazily + // at the first stream call in this mode. ctx.load_undigested(dict_bytes) .map_err(|_| bad_dict_err())?; } From 13351483c9c5b82508899c265978d8fed7858eaa Mon Sep 17 00:00:00 2001 From: James David Clarke Date: Wed, 19 Aug 2026 10:06:33 +0100 Subject: [PATCH 17/20] Replace the _zstd flush-mode constants with a CompressMode enum `COMP_MODE_CONTINUE`/`FLUSH_BLOCK`/`FLUSH_FRAME` become a `CompressMode` enum with pinned discriminants, decoded at the Python boundary by `CompressMode::from_int` and mapped to libzstd by `CompressMode::end_directive`. `CompressorState::last_mode` now holds the enum, so the states `set_pledged_input_size` accepts are checked by the type system rather than by comparing ints. `flush()` keeps rejecting `CONTINUE` and out-of-range modes with its own message rather than the one `compress()` uses, so the Python-visible errors are unchanged. Co-Authored-By: Claude Opus 5 --- crates/stdlib/src/zstd.rs | 121 ++++++++++++++++++++++++++------------ 1 file changed, 83 insertions(+), 38 deletions(-) diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs index 833e4da5cea..63624ba6d2f 100644 --- a/crates/stdlib/src/zstd.rs +++ b/crates/stdlib/src/zstd.rs @@ -617,12 +617,42 @@ mod _zstd { } } - // The three flush modes for `ZstdCompressor.compress()`, mirrored as - // class attributes via `extend_class` below. Values are positional and - // chosen to match what CPython exposes. - const COMP_MODE_CONTINUE: i32 = 0; - const COMP_MODE_FLUSH_BLOCK: i32 = 1; - const COMP_MODE_FLUSH_FRAME: i32 = 2; + /// The flush modes `ZstdCompressor.compress()` and + /// `ZstdCompressor.flush()` accept, mirrored as the class attributes + /// `CONTINUE`, `FLUSH_BLOCK` and `FLUSH_FRAME` by `extend_class` below. + /// + /// The discriminants are pinned because they cross the Python boundary as + /// plain ints, both as those class attributes and as the value the + /// `last_mode` getset hands back. They match what CPython's `_zstd` + /// exposes. + #[derive(Copy, Clone, PartialEq, Eq)] + enum CompressMode { + Continue = 0, + FlushBlock = 1, + FlushFrame = 2, + } + + impl CompressMode { + /// Decode a user-supplied `mode=` argument. `None` for anything + /// outside the three documented modes. + fn from_int(mode: i32) -> Option { + Some(match mode { + m if m == Self::Continue as i32 => Self::Continue, + m if m == Self::FlushBlock as i32 => Self::FlushBlock, + m if m == Self::FlushFrame as i32 => Self::FlushFrame, + _ => return None, + }) + } + + /// The libzstd end directive this mode drives the streaming API with. + fn end_directive(self) -> ZSTD_EndDirective { + match self { + Self::Continue => ZSTD_e_continue, + Self::FlushBlock => ZSTD_e_flush, + Self::FlushFrame => ZSTD_e_end, + } + } + } /// Owning wrapper around a raw libzstd compression context. Frees the /// context with `ZSTD_freeCCtx` on drop. @@ -793,7 +823,7 @@ mod _zstd { _cdict: Option, /// Keeps the ZstdDict's bytes alive for `ZSTD_CCtx_refPrefix` mode. _dict: Option>, - last_mode: i32, + last_mode: CompressMode, } #[pyattr] @@ -819,20 +849,22 @@ mod _zstd { zstd_dict: OptionalOption, } - /// Translate the public `mode` int to the libzstd `ZSTD_EndDirective` - /// the streaming API takes. - fn end_directive_from_mode(mode: i32, vm: &VirtualMachine) -> PyResult { - match mode { - COMP_MODE_CONTINUE => Ok(ZSTD_e_continue), - COMP_MODE_FLUSH_BLOCK => Ok(ZSTD_e_flush), - COMP_MODE_FLUSH_FRAME => Ok(ZSTD_e_end), - _ => Err(vm.new_value_error(format!( + /// Decode the public `mode` int accepted by `ZstdCompressor.compress()`, + /// which takes all three modes. + fn compress_mode_from_int(mode: i32, vm: &VirtualMachine) -> PyResult { + CompressMode::from_int(mode).ok_or_else(|| { + let (cont, block, frame) = ( + CompressMode::Continue as i32, + CompressMode::FlushBlock as i32, + CompressMode::FlushFrame as i32, + ); + vm.new_value_error(format!( "mode argument wrong value, it should be one of \ - ZstdCompressor.CONTINUE ({COMP_MODE_CONTINUE}), \ - ZstdCompressor.FLUSH_BLOCK ({COMP_MODE_FLUSH_BLOCK}), or \ - ZstdCompressor.FLUSH_FRAME ({COMP_MODE_FLUSH_FRAME})" - ))), - } + ZstdCompressor.CONTINUE ({cont}), \ + ZstdCompressor.FLUSH_BLOCK ({block}), or \ + ZstdCompressor.FLUSH_FRAME ({frame})" + )) + }) } impl Constructor for ZstdCompressor { @@ -859,7 +891,7 @@ mod _zstd { apply_options(&mut cctx, options_obj, true, vm)?; } - let state = build_compressor_state(cctx, dict_opt, COMP_MODE_FLUSH_FRAME, vm)?; + let state = build_compressor_state(cctx, dict_opt, CompressMode::FlushFrame, vm)?; Ok(Self { state: PyMutex::new(state), }) @@ -1128,7 +1160,7 @@ mod _zstd { fn build_compressor_state( mut cctx: CCtx, dict_obj: Option, - last_mode: i32, + last_mode: CompressMode, vm: &VirtualMachine, ) -> PyResult { // SAFETY: `load_dict` requires its two return values to outlive `ctx`. @@ -1238,8 +1270,11 @@ mod _zstd { impl ZstdCompressor { #[pymethod] fn compress(&self, args: CompressMethodArgs, vm: &VirtualMachine) -> PyResult> { - let mode = args.mode.unwrap_or(COMP_MODE_CONTINUE); - let end_op = end_directive_from_mode(mode, vm)?; + let mode = match args.mode { + Some(mode) => compress_mode_from_int(mode, vm)?, + None => CompressMode::Continue, + }; + let end_op = mode.end_directive(); let data = args.data.with_ref(|b| b.to_vec()); let mut state = self.state.lock(); let out = do_compress(&mut state, &data, end_op, vm)?; @@ -1249,15 +1284,25 @@ mod _zstd { #[pymethod] fn flush(&self, args: FlushMethodArgs, vm: &VirtualMachine) -> PyResult> { - let mode = args.mode.unwrap_or(COMP_MODE_FLUSH_FRAME); - if mode != COMP_MODE_FLUSH_BLOCK && mode != COMP_MODE_FLUSH_FRAME { - return Err(vm.new_value_error(format!( - "mode argument wrong value, it should be \ - ZstdCompressor.FLUSH_FRAME ({COMP_MODE_FLUSH_FRAME}) or \ - ZstdCompressor.FLUSH_BLOCK ({COMP_MODE_FLUSH_BLOCK})" - ))); - } - let end_op = end_directive_from_mode(mode, vm)?; + // `flush()` takes only the two flushing modes, and rejects + // everything else — `CONTINUE` included — with its own message + // rather than the one `compress()` uses. + let mode = match args.mode.map(CompressMode::from_int) { + None => CompressMode::FlushFrame, + Some(Some(mode @ (CompressMode::FlushBlock | CompressMode::FlushFrame))) => mode, + Some(_) => { + let (block, frame) = ( + CompressMode::FlushBlock as i32, + CompressMode::FlushFrame as i32, + ); + return Err(vm.new_value_error(format!( + "mode argument wrong value, it should be \ + ZstdCompressor.FLUSH_FRAME ({frame}) or \ + ZstdCompressor.FLUSH_BLOCK ({block})" + ))); + } + }; + let end_op = mode.end_directive(); let mut state = self.state.lock(); let out = do_compress(&mut state, &[], end_op, vm)?; state.last_mode = mode; @@ -1299,7 +1344,7 @@ mod _zstd { Some(v) }; let mut state = self.state.lock(); - if state.last_mode != COMP_MODE_FLUSH_FRAME { + if state.last_mode != CompressMode::FlushFrame { return Err(vm.new_value_error( "set_pledged_input_size() method must be called when last_mode == FLUSH_FRAME", )); @@ -1313,7 +1358,7 @@ mod _zstd { #[pygetset] fn last_mode(&self) -> i32 { - self.state.lock().last_mode + self.state.lock().last_mode as i32 } // Install class-level constants `CONTINUE`, `FLUSH_BLOCK`, and @@ -1324,15 +1369,15 @@ mod _zstd { fn extend_class(ctx: &Context, class: &'static Py) { class.set_attr( ctx.intern_str("CONTINUE"), - ctx.new_int(COMP_MODE_CONTINUE).into(), + ctx.new_int(CompressMode::Continue as i32).into(), ); class.set_attr( ctx.intern_str("FLUSH_BLOCK"), - ctx.new_int(COMP_MODE_FLUSH_BLOCK).into(), + ctx.new_int(CompressMode::FlushBlock as i32).into(), ); class.set_attr( ctx.intern_str("FLUSH_FRAME"), - ctx.new_int(COMP_MODE_FLUSH_FRAME).into(), + ctx.new_int(CompressMode::FlushFrame as i32).into(), ); } } From 68bed0441efe2c26ba77ef205bf629832d5316bb Mon Sep 17 00:00:00 2001 From: James David Clarke Date: Wed, 19 Aug 2026 10:54:25 +0100 Subject: [PATCH 18/20] Correct the _zstd explanatory comments and raise MemoryError on OOM MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Verified each claim the previous two commits added, against the pinned libzstd-rs-sys rev, CPython's `Modules/_zstd/` and `Lib/test/test_zstd.py`, and corrected the ones that did not hold: - The `DictType` discriminants come from CPython's `dictionary_type` enum, not a type named `DictType`, and `test_zstd` never hand-builds a valid `(zdict, marker)` tuple — it pins the accepted range from the outside with `(zd, -1)` and `(zd, 3)` rejection cases. - `#[repr(transparent)]` guarantees layout, not validity invariants. The transmute is sound because every bit pattern is a valid `u32`, which holds only in the enum-to-int direction; say so, and note that this is why `c_param_enum` decodes untrusted ids with an explicit match. - The id is unreachable through the crate's public API only in the sense that no accessor exists; the derived `Debug` does print it. - Undigested dictionary loading is not uniformly lazy: `ZSTD_DCtx_loadDictionary` digests eagerly and rejects corrupted content at construction time, while `ZSTD_CCtx_loadDictionary` accepts it. - `CompressMode`'s discriminants are libzstd's `ZSTD_e_continue`/ `ZSTD_e_flush`/`ZSTD_e_end` values. `CCtx::create`/`DCtx::create` now return `PyResult` and raise `MemoryError` like CPython's `_zstd` does, rather than panicking. The old justification for the panic claimed a `vm` reference would have to be threaded through every construction site; there are two, and both already had one. Co-Authored-By: Claude Opus 5 --- crates/stdlib/src/zstd.rs | 75 +++++++++++++++++++++------------------ 1 file changed, 40 insertions(+), 35 deletions(-) diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs index 63624ba6d2f..59247cb2fed 100644 --- a/crates/stdlib/src/zstd.rs +++ b/crates/stdlib/src/zstd.rs @@ -62,13 +62,17 @@ mod _zstd { /// `CompressionParameter` IntEnum. libzstd-rs-sys declares /// `ZSTD_cParameter` as `#[repr(transparent)] pub struct /// ZSTD_cParameter(pub(crate) u32)` and gives it no accessor, `From` impl - /// or `Deref`, so the id is unreachable through its public API. Reading - /// the wrapped `u32` back keeps libzstd the single source of truth for - /// these values instead of restating them here as literals. + /// or `Deref`, so short of parsing its derived `Debug` output there is no + /// way to ask it for its id. Reading the wrapped `u32` back keeps libzstd + /// the single source of truth for these values instead of restating them + /// here as literals. const fn c_param_id(param: ZSTD_cParameter) -> i32 { - // SAFETY: `#[repr(transparent)]` guarantees `ZSTD_cParameter` has the - // same size, alignment and validity invariants as the `u32` it wraps, - // so reading one back as a `u32` is always well-defined. + // SAFETY: `#[repr(transparent)]` gives `ZSTD_cParameter` the layout of + // the single `u32` it wraps, and every bit pattern is a valid `u32`, + // so reading one back out is well-defined. Note this holds only in + // this direction: `u32 -> ZSTD_cParameter` would be a widening into a + // type whose valid values libzstd defines, which is why `c_param_enum` + // below decodes untrusted ids with an explicit match instead. unsafe { core::mem::transmute::(param) as i32 } } @@ -179,9 +183,10 @@ mod _zstd { /// decompressor constructors turn that marker back into this enum. /// /// The discriminants are pinned because they cross the Python boundary as - /// plain ints: they are the values CPython's - /// `Modules/_zstd/_zstdmodule.h::DictType` uses, and `test_zstd` builds - /// the tuples by hand with those literals. + /// plain ints: they are the values of CPython's `dictionary_type` enum in + /// `Modules/_zstd/_zstdmodule.h`. `test_zstd` pins the accepted range from + /// the outside, hand-building `(zd, -1)` and `(zd, 3)` tuples and + /// requiring both to raise. #[derive(Copy, Clone)] enum DictType { Digested = 0, @@ -621,10 +626,11 @@ mod _zstd { /// `ZstdCompressor.flush()` accept, mirrored as the class attributes /// `CONTINUE`, `FLUSH_BLOCK` and `FLUSH_FRAME` by `extend_class` below. /// - /// The discriminants are pinned because they cross the Python boundary as - /// plain ints, both as those class attributes and as the value the - /// `last_mode` getset hands back. They match what CPython's `_zstd` - /// exposes. + /// The discriminants are libzstd's own `ZSTD_e_continue`, `ZSTD_e_flush` + /// and `ZSTD_e_end` values, which is what CPython's `_zstd` exposes under + /// these attribute names. They are pinned here because they cross the + /// Python boundary as plain ints, both as those class attributes and as + /// the value the `last_mode` getset hands back. #[derive(Copy, Clone, PartialEq, Eq)] enum CompressMode { Continue = 0, @@ -664,17 +670,15 @@ mod _zstd { unsafe impl Send for CCtx {} impl CCtx { - fn create() -> Self { + fn create(vm: &VirtualMachine) -> PyResult { // SAFETY: `ZSTD_createCCtx` has no preconditions; it returns NULL - // only when its `malloc` of the context struct fails. + // only when its allocation of the context struct fails, so that is + // the sole case to handle. let ptr = unsafe { ZSTD_createCCtx() }; - // A NULL return therefore means the allocator is out of memory, - // which is the one case Rust already handles by aborting, so - // panicking matches the rest of the VM rather than threading a - // `vm` reference through every construction site to raise - // `MemoryError` the way CPython's `_zstd` does. - assert!(!ptr.is_null(), "ZSTD_createCCtx failed"); - Self(ptr) + if ptr.is_null() { + return Err(vm.new_memory_error("failed to allocate a zstd context")); + } + Ok(Self(ptr)) } /// Forward a `(parameter, value)` pair to libzstd, passing the value @@ -720,17 +724,15 @@ mod _zstd { unsafe impl Send for DCtx {} impl DCtx { - fn create() -> Self { + fn create(vm: &VirtualMachine) -> PyResult { // SAFETY: `ZSTD_createDCtx` has no preconditions; it returns NULL - // only when its `malloc` of the context struct fails. + // only when its allocation of the context struct fails, so that is + // the sole case to handle. let ptr = unsafe { ZSTD_createDCtx() }; - // A NULL return therefore means the allocator is out of memory, - // which is the one case Rust already handles by aborting, so - // panicking matches the rest of the VM rather than threading a - // `vm` reference through every construction site to raise - // `MemoryError` the way CPython's `_zstd` does. - assert!(!ptr.is_null(), "ZSTD_createDCtx failed"); - Self(ptr) + if ptr.is_null() { + return Err(vm.new_memory_error("failed to allocate a zstd context")); + } + Ok(Self(ptr)) } /// See [`CCtx::set_parameter`]. @@ -879,7 +881,7 @@ mod _zstd { return Err(vm.new_type_error("Only one of level or options should be used.")); } - let mut cctx = CCtx::create(); + let mut cctx = CCtx::create(vm)?; if let Some(level_obj) = level_opt { let level = parse_compression_level(&level_obj, vm)?; @@ -1143,8 +1145,11 @@ mod _zstd { digested = Some(d); } DictType::Undigested => { - // Copy the bytes into the context. Validation happens lazily - // at the first stream call in this mode. + // Copy the bytes into the context. Whether a corrupted + // dictionary is caught here is up to libzstd and differs by + // direction: `ZSTD_DCtx_loadDictionary` digests the content + // eagerly and fails now, while `ZSTD_CCtx_loadDictionary` + // defers that work and accepts it. ctx.load_undigested(dict_bytes) .map_err(|_| bad_dict_err())?; } @@ -1445,7 +1450,7 @@ mod _zstd { let dict_opt = args.zstd_dict.flatten(); let options_opt = args.options.flatten(); - let mut dctx = DCtx::create(); + let mut dctx = DCtx::create(vm)?; if let Some(options_obj) = options_opt { apply_options(&mut dctx, options_obj, false, vm)?; From e1635610b92c0336b6ad54c39814da477fadce1c Mon Sep 17 00:00:00 2001 From: James David Clarke Date: Wed, 19 Aug 2026 12:44:11 +0100 Subject: [PATCH 19/20] Fix CPython divergences and a crash in the _zstd module MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Found by reviewing the module against CPython's `Modules/_zstd/` and the pinned libzstd-rs-sys sources; each item below was reproduced before the fix and re-checked after. - `train_dict`/`finalize_dict` sized their output buffer with an infallible `Vec::with_capacity`/`vec![]` from a caller-supplied `dict_size`, so `train_dict([b'x'], 2**62)` aborted the interpreter instead of raising. Reserve fallibly and raise `MemoryError`, as CPython does. - A digested dictionary was always built at `ZSTD_CLEVEL_DEFAULT`. Since `ZSTD_CCtx_refCDict` takes its parameters from the CDict, that silently discarded the requested level: `level=19` with `as_digested_dict` compressed a test corpus to 39817 bytes instead of 31951. Build the CDict at the compressor's effective level, tracked through `level=` and the `options` dict like CPython's `self->compression_level`. - A bare `ZstdDict` was loaded as digested in both directions; CPython compresses with an undigested dictionary by default and only decompresses with a digested one. Give `DictLoader` a per-direction default. - Neither compressor nor decompressor reset its session after an error, leaving the object permanently broken — every later call failed with "Operation not authorized at current processing stage" or re-reported the original corruption on valid input. Reset as CPython's error paths do, including restoring `last_mode` to `FLUSH_FRAME`. - `apply_options` range-checked every parameter, rejecting values libzstd and CPython accept: 0 ("use the default") for `window_log`, `strategy`, `window_log_max` and friends, any non-zero value for the boolean flags, and the clamped `nb_workers`/`job_size`/`overlap_log`. Only `compression_level` needs the check — libzstd clamps that one silently, which is why `test_compress_parameters` requires a `ValueError` — so everything else now goes to libzstd, whose error code already maps to the same message. - `ZstdDict` accepted content shorter than 8 bytes; CPython rejects it up front regardless of `is_raw`. - `compress`/`decompress` refused `data` as a keyword and `get_param_bounds` required `is_compress` as one, both contrary to CPython's clinic signatures. - The decompressor allocated a full 128 KiB scratch buffer even when `max_length` capped output far below that. `test_zstd` (119 tests) and the snippet test pass; clippy is clean. Co-Authored-By: Claude Opus 5 --- crates/stdlib/src/zstd.rs | 230 ++++++++++++++++++++++++++++++-------- 1 file changed, 182 insertions(+), 48 deletions(-) diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs index 59247cb2fed..c873fda3cbd 100644 --- a/crates/stdlib/src/zstd.rs +++ b/crates/stdlib/src/zstd.rs @@ -29,17 +29,17 @@ mod _zstd { use libzstd_rs_sys::lib::zdict::ZDICT_finalizeDictionary; use libzstd_rs_sys::{ ZDICT_getErrorName, ZDICT_isError, ZDICT_params_t, ZDICT_trainFromBuffer, ZSTD_CCtx, - ZSTD_CCtx_loadDictionary, ZSTD_CCtx_refCDict, ZSTD_CCtx_refPrefix, ZSTD_CCtx_setParameter, - ZSTD_CCtx_setPledgedSrcSize, ZSTD_CDict, ZSTD_CONTENTSIZE_ERROR, ZSTD_CONTENTSIZE_UNKNOWN, - ZSTD_CStreamOutSize, ZSTD_DCtx, ZSTD_DCtx_loadDictionary, ZSTD_DCtx_refDDict, - ZSTD_DCtx_refPrefix, ZSTD_DCtx_setParameter, ZSTD_DDict, ZSTD_DStreamOutSize, - ZSTD_EndDirective, ZSTD_cParam_getBounds, ZSTD_cParameter, ZSTD_compressStream2, - ZSTD_createCCtx, ZSTD_createCDict, ZSTD_createDCtx, ZSTD_createDDict, - ZSTD_dParam_getBounds, ZSTD_dParameter, ZSTD_decompressStream, - ZSTD_findFrameCompressedSize, ZSTD_freeCCtx, ZSTD_freeCDict, ZSTD_freeDCtx, ZSTD_freeDDict, - ZSTD_getDictID_fromDict, ZSTD_getDictID_fromFrame, ZSTD_getErrorName, - ZSTD_getFrameContentSize, ZSTD_inBuffer, ZSTD_isError, ZSTD_outBuffer, ZSTD_strategy, - ZSTD_versionNumber, ZSTD_versionString, + ZSTD_CCtx_loadDictionary, ZSTD_CCtx_refCDict, ZSTD_CCtx_refPrefix, ZSTD_CCtx_reset, + ZSTD_CCtx_setParameter, ZSTD_CCtx_setPledgedSrcSize, ZSTD_CDict, ZSTD_CONTENTSIZE_ERROR, + ZSTD_CONTENTSIZE_UNKNOWN, ZSTD_CStreamOutSize, ZSTD_DCtx, ZSTD_DCtx_loadDictionary, + ZSTD_DCtx_refDDict, ZSTD_DCtx_refPrefix, ZSTD_DCtx_reset, ZSTD_DCtx_setParameter, + ZSTD_DDict, ZSTD_DStreamOutSize, ZSTD_EndDirective, ZSTD_ResetDirective, + ZSTD_cParam_getBounds, ZSTD_cParameter, ZSTD_compressStream2, ZSTD_createCCtx, + ZSTD_createCDict, ZSTD_createDCtx, ZSTD_createDDict, ZSTD_dParam_getBounds, + ZSTD_dParameter, ZSTD_decompressStream, ZSTD_findFrameCompressedSize, ZSTD_freeCCtx, + ZSTD_freeCDict, ZSTD_freeDCtx, ZSTD_freeDDict, ZSTD_getDictID_fromDict, + ZSTD_getDictID_fromFrame, ZSTD_getErrorName, ZSTD_getFrameContentSize, ZSTD_inBuffer, + ZSTD_isError, ZSTD_outBuffer, ZSTD_strategy, ZSTD_versionNumber, ZSTD_versionString, }; use rustpython_common::lock::PyMutex; use rustpython_vm::builtins::{ @@ -317,7 +317,9 @@ mod _zstd { "invalid compression parameter 'unknown parameter (key {param})'" )) })?; - if param == ZSTD_c_strategy && strategy_from_int(value).is_none() { + // Strategy 0 is libzstd's "use the default" sentinel, like the other + // parameters; only a non-zero value has to name a real strategy. + if param == ZSTD_c_strategy && value != 0 && strategy_from_int(value).is_none() { return Err(new_zstd_error( format!("invalid strategy value: {value}"), vm, @@ -457,17 +459,18 @@ mod _zstd { } /// Decode the `zstd_dict=` constructor argument. Accepts either a - /// `ZstdDict` instance (treated as the default digested form) or a - /// `(ZstdDict, marker)` tuple produced by one of `ZstdDict.as_*`. + /// `ZstdDict` instance, which takes the caller's `default` load mode, or + /// a `(ZstdDict, marker)` tuple produced by one of `ZstdDict.as_*`. fn parse_zstd_dict_arg( obj: PyObjectRef, + default: DictType, vm: &VirtualMachine, ) -> PyResult<(PyRef, DictType)> { // The first downcast clones `obj` because we fall through to the // tuple branch if it fails. The second downcast (the tuple one) is // the last use of `obj`, so we let it move directly. if let Ok(d) = obj.clone().downcast::() { - return Ok((d, DictType::Digested)); + return Ok((d, default)); } if let Ok(tuple) = obj.downcast::() { let items = tuple.as_slice(); @@ -534,6 +537,15 @@ mod _zstd { fn py_new(_cls: &Py, args: Self::Args, vm: &VirtualMachine) -> PyResult { let dict_content = args.dict_content.with_ref(|b| b.to_vec()); + // CPython checks this first, before and independently of the + // `is_raw` handling below, so a too-short buffer never reaches + // libzstd. + if dict_content.len() < 8 { + return Err(vm.new_value_error( + "Zstandard dictionary content too short \ + (must have at least eight bytes)", + )); + } // libzstd's `ZSTD_getDictID_fromDict` returns 0 either when // the content is too small to contain a valid header or when it // does not carry the dictionary magic. Both are runtime errors @@ -694,6 +706,14 @@ mod _zstd { } } + /// Return the context to a frame boundary, discarding a partially + /// written frame but keeping parameters and dictionary. libzstd + /// documents a session-only reset as infallible. + fn reset_session(&mut self) { + // SAFETY: `self.0` is a live context owned by this wrapper. + unsafe { ZSTD_CCtx_reset(self.0, ZSTD_ResetDirective::ZSTD_reset_session_only) }; + } + fn set_pledged_src_size(&mut self, pledged: Option) -> Result<(), usize> { // libzstd represents "unknown" as `ZSTD_CONTENTSIZE_UNKNOWN`. let pledged = pledged.map_or(ZSTD_CONTENTSIZE_UNKNOWN, |v| v); @@ -735,6 +755,12 @@ mod _zstd { Ok(Self(ptr)) } + /// See [`CCtx::reset_session`]. + fn reset_session(&mut self) { + // SAFETY: `self.0` is a live context owned by this wrapper. + unsafe { ZSTD_DCtx_reset(self.0, ZSTD_ResetDirective::ZSTD_reset_session_only) }; + } + /// See [`CCtx::set_parameter`]. fn set_parameter(&mut self, param: ZSTD_dParameter, value: c_int) -> Result<(), usize> { // SAFETY: `self.0` is a live context owned by this wrapper. @@ -828,6 +854,17 @@ mod _zstd { last_mode: CompressMode, } + impl CompressorState { + /// Recover from a failed `compress()`/`flush()`, as CPython does: + /// drop the half-written frame and report `last_mode == FLUSH_FRAME`. + /// Without this the context stays mid-frame and every later call + /// fails with "Operation not authorized at current processing stage". + fn reset_session(&mut self) { + self.last_mode = CompressMode::FlushFrame; + self.cctx.reset_session(); + } + } + #[pyattr] #[pyclass(name = "ZstdCompressor")] #[derive(PyPayload)] @@ -883,17 +920,25 @@ mod _zstd { let mut cctx = CCtx::create(vm)?; + // Mirror CPython's `self->compression_level`: the level libzstd + // will actually compress at, which a digested dictionary has to + // be built for. + let mut level = ZSTD_CLEVEL_DEFAULT; + if let Some(level_obj) = level_opt { - let level = parse_compression_level(&level_obj, vm)?; + level = parse_compression_level(&level_obj, vm)?; cctx.set_parameter(ZSTD_cParameter::ZSTD_c_compressionLevel, level) .map_err(|_| param_value_error_for(ZSTD_c_compressionLevel, level, true, vm))?; } - if let Some(options_obj) = options_opt { - apply_options(&mut cctx, options_obj, true, vm)?; + if let Some(options_obj) = options_opt + && let Some(from_options) = apply_options(&mut cctx, options_obj, true, vm)? + { + level = from_options; } - let state = build_compressor_state(cctx, dict_opt, CompressMode::FlushFrame, vm)?; + let state = + build_compressor_state(cctx, dict_opt, level, CompressMode::FlushFrame, vm)?; Ok(Self { state: PyMutex::new(state), }) @@ -932,7 +977,8 @@ mod _zstd { options_obj: PyObjectRef, is_compress: bool, vm: &VirtualMachine, - ) -> PyResult<()> { + ) -> PyResult> { + let mut level = None; let dict = options_obj .downcast::() .map_err(|_| vm.new_type_error("options must be a dict"))?; @@ -953,17 +999,34 @@ mod _zstd { check_wrong_param_kind(&k, wrong_kind.as_ref(), kind, vm)?; let key_int: i32 = k.try_to_value(vm)?; let val_int: i32 = v.try_to_value(vm)?; - // libzstd silently clamps out-of-range values for some - // parameters (notably compression_level) rather than rejecting - // them, so validate against the documented bounds upfront. - if let Some((lo, hi)) = lookup_param_bounds(key_int, is_compress) + // Only the compression level is range-checked here. libzstd + // silently clamps it (`ZSTD_cParam_clampBounds`) instead of + // reporting an error, so without this check the `ValueError` + // that `test_compress_parameters` requires would never be + // raised. + // + // Every other parameter is left to libzstd, which does its own + // bounds checking and whose error code becomes the same + // `ValueError` via `ParamSetter::apply`. A blanket range test + // here would reject values CPython accepts: 0 means "use the + // default" for most parameters, the boolean flags take any + // non-zero value, and `nb_workers`/`job_size`/`overlap_log` are + // clamped rather than rejected. + if is_compress + && key_int == ZSTD_c_compressionLevel + && let Some((lo, hi)) = lookup_param_bounds(key_int, is_compress) && (val_int < lo || val_int > hi) { return Err(param_value_error_for(key_int, val_int, is_compress, vm)); } ctx.apply(key_int, val_int, vm)?; + // Track the level the compressor ends up with: a digested + // dictionary has to be built at that level (see `load_dict`). + if is_compress && key_int == ZSTD_c_compressionLevel { + level = Some(val_int); + } } - Ok(()) + Ok(level) } /// Trait wrapper over `CCtx::set_parameter` and `DCtx::set_parameter` so @@ -1000,7 +1063,13 @@ mod _zstd { trait DictLoader { type Digested; const KIND_NAME: &'static str; - fn try_create_digested(bytes: &[u8]) -> Option; + /// Load mode applied when the caller passes a bare `ZstdDict` instead + /// of one of the `as_*` tuples. CPython compresses with an undigested + /// dictionary by default and decompresses with a digested one. + const DEFAULT_DICT_TYPE: DictType; + /// `level` is the compression level the digested dictionary must be + /// built for; the decompressor has no level and ignores it. + fn try_create_digested(bytes: &[u8], level: c_int) -> Option; fn ref_digested(&mut self, dict: &Self::Digested) -> Result<(), usize>; fn load_undigested(&mut self, bytes: &[u8]) -> Result<(), usize>; fn ref_prefix(&mut self, bytes: &[u8]) -> Result<(), usize>; @@ -1009,8 +1078,12 @@ mod _zstd { impl DictLoader for CCtx { type Digested = CDict; const KIND_NAME: &'static str = "ZSTD_CDict"; - fn try_create_digested(bytes: &[u8]) -> Option { - CDict::try_create(bytes, libzstd_rs_sys::ZSTD_CLEVEL_DEFAULT) + const DEFAULT_DICT_TYPE: DictType = DictType::Undigested; + fn try_create_digested(bytes: &[u8], level: c_int) -> Option { + // The level matters: `ZSTD_CCtx_refCDict` takes its compression + // parameters from the CDict, so a CDict built at the default + // level would override whatever the caller asked for. + CDict::try_create(bytes, level) } fn ref_digested(&mut self, dict: &Self::Digested) -> Result<(), usize> { // SAFETY: both handles are live; `load_dict`'s safety contract @@ -1050,7 +1123,8 @@ mod _zstd { impl DictLoader for DCtx { type Digested = DDict; const KIND_NAME: &'static str = "ZSTD_DDict"; - fn try_create_digested(bytes: &[u8]) -> Option { + const DEFAULT_DICT_TYPE: DictType = DictType::Digested; + fn try_create_digested(bytes: &[u8], _level: c_int) -> Option { DDict::try_create(bytes) } fn ref_digested(&mut self, dict: &Self::Digested) -> Result<(), usize> { @@ -1115,12 +1189,13 @@ mod _zstd { unsafe fn load_dict( ctx: &mut L, dict_obj: Option, + level: c_int, vm: &VirtualMachine, ) -> DictLoadResult { let Some(dict_obj) = dict_obj else { return Ok((None, None)); }; - let (zdict, dict_type) = parse_zstd_dict_arg(dict_obj, vm)?; + let (zdict, dict_type) = parse_zstd_dict_arg(dict_obj, L::DEFAULT_DICT_TYPE, vm)?; let bad_dict_err = || -> PyBaseExceptionRef { new_zstd_error( format!( @@ -1140,7 +1215,7 @@ mod _zstd { // Build the digested dict eagerly so a corrupted dictionary // surfaces as a `ZstdError` at construction time, not when // the first compress/decompress call runs. - let d = L::try_create_digested(dict_bytes).ok_or_else(bad_dict_err)?; + let d = L::try_create_digested(dict_bytes, level).ok_or_else(bad_dict_err)?; ctx.ref_digested(&d).map_err(|_| bad_dict_err())?; digested = Some(d); } @@ -1165,6 +1240,7 @@ mod _zstd { fn build_compressor_state( mut cctx: CCtx, dict_obj: Option, + level: c_int, last_mode: CompressMode, vm: &VirtualMachine, ) -> PyResult { @@ -1176,7 +1252,7 @@ mod _zstd { // and `_dict` (prefix mode) are freed. `CompressorState` is private // to this module and is never destructured, so no safe caller can // reorder the drops. - let (cdict, dict) = unsafe { load_dict::(&mut cctx, dict_obj, vm) }?; + let (cdict, dict) = unsafe { load_dict::(&mut cctx, dict_obj, level, vm) }?; Ok(CompressorState { cctx, _cdict: cdict, @@ -1194,7 +1270,9 @@ mod _zstd { vm: &VirtualMachine, ) -> PyResult { // SAFETY: see [`build_compressor_state`]. - let (ddict, dict) = unsafe { load_dict::(&mut dctx, dict_obj, vm) }?; + // The decompressor has no compression level; `DCtx` ignores it. + let (ddict, dict) = + unsafe { load_dict::(&mut dctx, dict_obj, ZSTD_CLEVEL_DEFAULT, vm) }?; Ok(DecompressorState { dctx, _ddict: ddict, @@ -1282,9 +1360,16 @@ mod _zstd { let end_op = mode.end_directive(); let data = args.data.with_ref(|b| b.to_vec()); let mut state = self.state.lock(); - let out = do_compress(&mut state, &data, end_op, vm)?; - state.last_mode = mode; - Ok(out) + match do_compress(&mut state, &data, end_op, vm) { + Ok(out) => { + state.last_mode = mode; + Ok(out) + } + Err(e) => { + state.reset_session(); + Err(e) + } + } } #[pymethod] @@ -1309,9 +1394,16 @@ mod _zstd { }; let end_op = mode.end_directive(); let mut state = self.state.lock(); - let out = do_compress(&mut state, &[], end_op, vm)?; - state.last_mode = mode; - Ok(out) + match do_compress(&mut state, &[], end_op, vm) { + Ok(out) => { + state.last_mode = mode; + Ok(out) + } + Err(e) => { + state.reset_session(); + Err(e) + } + } } #[pymethod] @@ -1389,7 +1481,9 @@ mod _zstd { #[derive(FromArgs)] pub(super) struct CompressMethodArgs { - #[pyarg(positional)] + /// Positional-or-keyword, as in CPython's clinic signature for + /// `_zstd.ZstdCompressor.compress` (no `/` marker). + #[pyarg(any)] data: ArgBytesLike, #[pyarg(any, optional)] mode: Option, @@ -1422,6 +1516,20 @@ mod _zstd { input_buffer: Vec, } + impl DecompressorState { + /// Recover from a failed `decompress()`, mirroring CPython's + /// `decompressor_reset_session_lock_held`. Without it the DCtx stays + /// wedged on the corrupt frame and every later call re-reports the + /// same error, even for valid input. + fn reset_session(&mut self, vm: &VirtualMachine) { + self.input_buffer.clear(); + self.unused_data = vm.ctx.empty_bytes.clone(); + self.needs_input = true; + self.eof = false; + self.dctx.reset_session(); + } + } + #[pyattr] #[pyclass(name = "ZstdDecompressor")] #[derive(PyPayload)] @@ -1465,7 +1573,9 @@ mod _zstd { #[derive(FromArgs)] pub(super) struct DecompressMethodArgs { - #[pyarg(positional)] + /// Positional-or-keyword, as in CPython's clinic signature for + /// `_zstd.ZstdDecompressor.decompress`. + #[pyarg(any)] data: ArgBytesLike, #[pyarg(any, default = -1)] max_length: isize, @@ -1514,8 +1624,13 @@ mod _zstd { // Reusable scratch buffer for each decompress_stream call. We need // an exact-size output buffer because `Vec::reserve` may // over-allocate; reporting the full Vec capacity to libzstd would - // let it write past `max_length`. - let mut scratch: Vec = vec![0u8; chunk_size]; + // let it write past `max_length`. A cap below one chunk bounds + // how much of it can ever be used, so do not allocate more. + let scratch_size = match max_length { + Some(maxl) => maxl.min(chunk_size), + None => chunk_size, + }; + let mut scratch: Vec = vec![0u8; scratch_size]; let mut hit_max = false; let mut iteration = 0usize; @@ -1585,7 +1700,13 @@ mod _zstd { outcome.map(|()| (output, hit_max, input.pos)) }); - let (output, hit_max, consumed) = loop_result.map_err(|c| catch_zstd_error(c, vm))?; + let (output, hit_max, consumed) = match loop_result { + Ok(v) => v, + Err(code) => { + state.reset_session(vm); + return Err(catch_zstd_error(code, vm)); + } + }; let remaining = &work_data[consumed..]; @@ -1778,7 +1899,13 @@ mod _zstd { let samples_buffer = args.samples_bytes.as_bytes().to_vec(); let sizes = parse_sample_sizes(args.samples_sizes, vm)?; check_sample_sizes_match(&sizes, samples_buffer.len(), vm)?; - let mut dict_buffer: Vec = Vec::with_capacity(dict_size); + // `dict_size` is caller-controlled, so reserve fallibly: an + // infallible `Vec::with_capacity` aborts the whole interpreter on a + // failed allocation, where CPython raises `MemoryError`. + let mut dict_buffer: Vec = Vec::new(); + dict_buffer + .try_reserve_exact(dict_size) + .map_err(|_| vm.new_memory_error("failed to allocate the dictionary buffer"))?; // SAFETY: `dict_buffer`'s spare capacity is valid for writes of up to // `dict_size` bytes; the samples and sizes buffers are live slices // for the duration of the call. @@ -1833,7 +1960,12 @@ mod _zstd { let sizes = parse_sample_sizes(args.samples_sizes, vm)?; check_sample_sizes_match(&sizes, samples_buffer.len(), vm)?; - let mut dict_buffer: Vec = vec![0u8; dict_size]; + // Reserve fallibly before filling: see `train_dict`. + let mut dict_buffer: Vec = Vec::new(); + dict_buffer + .try_reserve_exact(dict_size) + .map_err(|_| vm.new_memory_error("failed to allocate the dictionary buffer"))?; + dict_buffer.resize(dict_size, 0); let params = ZDICT_params_t { compressionLevel: compression_level, notificationLevel: 0, @@ -1874,9 +2006,11 @@ mod _zstd { #[derive(FromArgs)] pub(super) struct ParamBoundsArgs { - #[pyarg(positional)] + #[pyarg(any)] parameter: i32, - #[pyarg(named)] + /// Positional-or-keyword, matching CPython's clinic signature + /// `get_param_bounds(parameter, is_compress)`. + #[pyarg(any)] is_compress: bool, } From 041f631a5e7df668d3a90bbc020a52d04ee85f35 Mon Sep 17 00:00:00 2001 From: James David Clarke Date: Wed, 19 Aug 2026 13:40:51 +0100 Subject: [PATCH 20/20] Restore Lib/test/test_zstd.py to its vendored CPython state MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The branch dropped the `skipIf(not SUPPORT_MULTITHREADING)` guard from `test_zstd_multithread_compress`. That guard can never fire against this backend: libzstd-rs-sys compiles multi-threaded compression in unconditionally — its Cargo.toml has no `zstdmt`-style feature — so `CompressionParameter.nb_workers.bounds()` is never `(0, 0)` and `SUPPORT_MULTITHREADING` is always true. Removing the decorator therefore changed nothing except to put a modification of a vendored CPython test in the diff, and it would turn into a hard failure rather than a skip if a future backend ever lacked multi-threading. Restore the file so the PR touches no CPython test. Co-Authored-By: Claude Opus 5 --- Lib/test/test_zstd.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/Lib/test/test_zstd.py b/Lib/test/test_zstd.py index 44a5bbb4b19..cf618534add 100644 --- a/Lib/test/test_zstd.py +++ b/Lib/test/test_zstd.py @@ -321,6 +321,8 @@ def test_unknown_compression_parameter(self): with self.assertRaisesRegex(ValueError, pattern): ZstdCompressor(options=option) + @unittest.skipIf(not SUPPORT_MULTITHREADING, + "zstd build doesn't support multi-threaded compression") def test_zstd_multithread_compress(self): size = 40*_1M b = THIS_FILE_BYTES * (size // len(THIS_FILE_BYTES))