From 77e99b1f0f44958f5fbed87ac05f9c8298822b4b Mon Sep 17 00:00:00 2001 From: Ray Walker Date: Sun, 1 Nov 2026 03:23:39 +0000 Subject: [PATCH 1/9] fix(rust): decode StandardSerializer MessagePack in safe Rust --- rust/src/msgpack_decode.rs | 345 +++++++++++++---------------- tests/unit/test_standard_decode.py | 34 ++- 2 files changed, 183 insertions(+), 196 deletions(-) diff --git a/rust/src/msgpack_decode.rs b/rust/src/msgpack_decode.rs index 79d5a776..d32ba5c4 100644 --- a/rust/src/msgpack_decode.rs +++ b/rust/src/msgpack_decode.rs @@ -12,16 +12,16 @@ //! //! The hook runs only on a map holding a `str` key equal to a marker key. Any other map it would //! return unchanged, because each marker test reads its key with `dict.get`. - -use std::os::raw::{c_char, c_long}; -use std::ptr; +//! +//! Safe Rust only (`SECURITY.md`): every object is a PyO3 `Bound`, so each reference is released +//! by its owner on any exit, and an array's items wait on a stack until it closes, so no +//! partly-filled list is ever built. use cachekit_core::check_msgpack_structure; use pyo3::exceptions::PyValueError; -use pyo3::ffi; use pyo3::prelude::*; use pyo3::sync::PyOnceLock; -use pyo3::types::PyBytes; +use pyo3::types::{PyBool, PyBytes, PyDict, PyFloat, PyList, PyString}; /// The keys `_standard_object_hook` tests (`standard_serializer.py`). const MARKER_KEYS: [&[u8]; 3] = [b"__datetime__", b"__date__", b"__time__"]; @@ -33,64 +33,52 @@ const VALUE_CACHE_MAX_LEN: usize = 32; const CACHE_BITS: u32 = 8; const CACHE_SLOTS: usize = 1 << CACHE_BITS; -/// A Python error is already set (`PyErr::fetch` takes it), or the input broke a bound. +/// A Python error is set (`restore` put it there, `PyErr::fetch` takes it), or the input broke a +/// bound. A unit enum, so a decode step's result stays two words. enum Fail { Py, Bound, } -/// An open array or map, holding a strong reference to its container. -struct Frame { - obj: *mut ffi::PyObject, +/// An open array or map. +struct Frame<'py> { + /// Map: the dict, filled as each value arrives. Array: None. + dict: Option>, + /// Map: the key read, awaiting its value. + key: Option>, + /// Array: where its items start on the decoder's `items` stack. + base: usize, /// Elements still to read: items for an array, keys plus values for a map. owed: u64, - /// Array: next index. Map: unused. - index: ffi::Py_ssize_t, - /// Map: the key read, awaiting its value (strong reference), or null. - key: *mut ffi::PyObject, - is_map: bool, /// Map: a `str` key equal to a marker key was read. marker: bool, } -/// One cached string: its bytes' position in the input, and the `str` (a strong reference). -#[derive(Clone, Copy)] -struct StrSlot { +/// One cached string: its bytes' position in the input, and the `str`. +struct StrSlot<'py> { start: usize, len: usize, - obj: *mut ffi::PyObject, + obj: Bound<'py, PyAny>, } -struct Decoder<'a> { +struct Decoder<'a, 'py> { + py: Python<'py>, data: &'a [u8], pos: usize, max_depth: usize, /// Elements owed across every open container, the root included (the walk's `pending`). pending: u64, - stack: Vec, + stack: Vec>, + /// The items of every open array, innermost last; an array's become its list when it closes. + /// Only decoded elements land here, so it never outgrows the input. + items: Vec>, /// Cached strings; a table slot holds an index into it plus one, or 0 when empty. A slot holds /// one string at a time, a colliding one replacing it, so this never exceeds `2 * CACHE_SLOTS`. - cached: Vec, + cached: Vec>, keys: [u16; CACHE_SLOTS], values: [u16; CACHE_SLOTS], } -impl Drop for Decoder<'_> { - fn drop(&mut self) { - // SAFETY: every pointer here is a strong reference this decoder owns, or null; the - // thread is attached (the decoder never outlives the pyfunction call). - unsafe { - for f in self.stack.drain(..) { - ffi::Py_XDECREF(f.key); - ffi::Py_DECREF(f.obj); - } - for slot in &self.cached { - ffi::Py_DECREF(slot.obj); - } - } - } -} - /// The cache slot for `bytes`: its first and last eight bytes, multiplied and shifted. #[inline] fn slot_of(bytes: &[u8]) -> usize { @@ -106,23 +94,16 @@ fn slot_of(bytes: &[u8]) -> usize { (h >> (64 - CACHE_BITS)) as usize } -#[inline] -unsafe fn decode_utf8(bytes: &[u8]) -> *mut ffi::PyObject { - ffi::PyUnicode_DecodeUTF8( - bytes.as_ptr().cast::(), - bytes.len() as ffi::Py_ssize_t, - ptr::null(), - ) -} - -impl<'a> Decoder<'a> { - fn new(data: &'a [u8], max_depth: usize) -> Self { +impl<'a, 'py> Decoder<'a, 'py> { + fn new(py: Python<'py>, data: &'a [u8], max_depth: usize) -> Self { Self { + py, data, pos: 0, max_depth, pending: 1, stack: Vec::new(), + items: Vec::new(), cached: Vec::new(), keys: [0; CACHE_SLOTS], values: [0; CACHE_SLOTS], @@ -156,14 +137,11 @@ impl<'a> Decoder<'a> { /// Whether the next element read is a map key. #[inline] fn at_key(&self) -> bool { - matches!(self.stack.last(), Some(f) if f.is_map && f.key.is_null()) + matches!(self.stack.last(), Some(f) if f.dict.is_some() && f.key.is_none()) } - fn decode( - &mut self, - py: Python<'_>, - hook: &Bound<'_, PyAny>, - ) -> Result<*mut ffi::PyObject, Fail> { + fn decode(&mut self, hook: &Bound<'py, PyAny>) -> Result, Fail> { + let py = self.py; loop { let marker = *self.data.get(self.pos).ok_or(Fail::Bound)?; self.pos += 1; @@ -171,82 +149,71 @@ impl<'a> Decoder<'a> { if let Some(f) = self.stack.last_mut() { f.owed -= 1; } - // SAFETY: the thread is attached; every constructor returns a new reference or null - // with an exception set, which `check` turns into Fail::Py. - let mut value = unsafe { - match marker { - 0x00..=0x7f => self.scalar(ffi::PyLong_FromLong(c_long::from(marker)))?, - 0xe0..=0xff => self.scalar(ffi::PyLong_FromLong(c_long::from(marker as i8)))?, - 0xc0 => self.scalar(ffi::Py_NewRef(ffi::Py_None()))?, - 0xc2 => self.scalar(ffi::Py_NewRef(ffi::Py_False()))?, - 0xc3 => self.scalar(ffi::Py_NewRef(ffi::Py_True()))?, - 0xa0..=0xbf => self.string(u64::from(marker & 0x1f))?, - 0xd9 => self.prefixed(1).and_then(|n| self.string(n))?, - 0xda => self.prefixed(2).and_then(|n| self.string(n))?, - 0xdb => self.prefixed(4).and_then(|n| self.string(n))?, - 0xc4 => self.prefixed(1).and_then(|n| self.binary(n))?, - 0xc5 => self.prefixed(2).and_then(|n| self.binary(n))?, - 0xc6 => self.prefixed(4).and_then(|n| self.binary(n))?, - 0xc7 => self.prefixed(1).and_then(|n| self.ext(py, n))?, - 0xc8 => self.prefixed(2).and_then(|n| self.ext(py, n))?, - 0xc9 => self.prefixed(4).and_then(|n| self.ext(py, n))?, - 0xd4..=0xd8 => self.ext(py, 1u64 << (marker - 0xd4))?, - 0xca..=0xd3 => self.number(marker)?, - 0x90..=0x9f => self.open(false, u64::from(marker & 0x0f))?, - 0xdc => self.prefixed(2).and_then(|n| self.open(false, n))?, - 0xdd => self.prefixed(4).and_then(|n| self.open(false, n))?, - 0x80..=0x8f => self.open(true, u64::from(marker & 0x0f))?, - 0xde => self.prefixed(2).and_then(|n| self.open(true, n))?, - 0xdf => self.prefixed(4).and_then(|n| self.open(true, n))?, - 0xc1 => return Err(Fail::Bound), - } + let value = match marker { + 0x00..=0x7f => self.scalar(int(py, i64::from(marker)))?, + 0xe0..=0xff => self.scalar(int(py, i64::from(marker as i8)))?, + 0xc0 => self.scalar(py.None().into_bound(py))?, + 0xc2 => self.scalar(PyBool::new(py, false).to_owned().into_any())?, + 0xc3 => self.scalar(PyBool::new(py, true).to_owned().into_any())?, + 0xa0..=0xbf => self.string(u64::from(marker & 0x1f))?, + 0xd9 => self.prefixed(1).and_then(|n| self.string(n))?, + 0xda => self.prefixed(2).and_then(|n| self.string(n))?, + 0xdb => self.prefixed(4).and_then(|n| self.string(n))?, + 0xc4 => self.prefixed(1).and_then(|n| self.binary(n))?, + 0xc5 => self.prefixed(2).and_then(|n| self.binary(n))?, + 0xc6 => self.prefixed(4).and_then(|n| self.binary(n))?, + 0xc7 => self.prefixed(1).and_then(|n| self.ext(n))?, + 0xc8 => self.prefixed(2).and_then(|n| self.ext(n))?, + 0xc9 => self.prefixed(4).and_then(|n| self.ext(n))?, + 0xd4..=0xd8 => self.ext(1u64 << (marker - 0xd4))?, + 0xca..=0xd3 => self.number(marker)?, + 0x90..=0x9f => self.open(false, u64::from(marker & 0x0f))?, + 0xdc => self.prefixed(2).and_then(|n| self.open(false, n))?, + 0xdd => self.prefixed(4).and_then(|n| self.open(false, n))?, + 0x80..=0x8f => self.open(true, u64::from(marker & 0x0f))?, + 0xde => self.prefixed(2).and_then(|n| self.open(true, n))?, + 0xdf => self.prefixed(4).and_then(|n| self.open(true, n))?, + 0xc1 => return Err(Fail::Bound), }; - // A non-empty container was pushed: read its first element. - if value.is_null() { + // None: a non-empty container was pushed, so read its first element. + let Some(mut value) = value else { continue; - } + }; // Hand the value to its container, closing every container it completes. loop { let Some(f) = self.stack.last_mut() else { return Ok(value); }; - // SAFETY: `f.obj` is a strong reference to a list or dict this decoder built; - // `value` and `f.key` are strong references the set calls consume. - unsafe { - if !f.is_map { - if ffi::PyList_SetItem(f.obj, f.index, value) != 0 { - return Err(Fail::Py); + match &f.dict { + None => self.items.push(value), + Some(dict) => match f.key.take() { + None => { + f.key = Some(value); + break; } - f.index += 1; - } else if f.key.is_null() { - f.key = value; - break; - } else { - let key = std::mem::replace(&mut f.key, ptr::null_mut()); - let ok = key_admitted(key); - let rc = if ok { - ffi::PyDict_SetItem(f.obj, key, value) - } else { - 0 - }; - if !ok { - set_key_error(py, key); + Some(key) => { + // strict_map_key=True: a map key must be exactly `str` or `bytes`. + if !(key.is_exact_instance_of::() + || key.is_exact_instance_of::()) + { + return Err(restore(py, key_error(&key))); + } + dict.set_item(key, value).map_err(|e| restore(py, e))?; } - ffi::Py_DECREF(key); - ffi::Py_DECREF(value); - if !ok || rc != 0 { - return Err(Fail::Py); - } - } - if f.owed > 0 { - break; - } - let done = self.stack.pop().expect("a frame was just read"); - value = done.obj; - if done.marker { - let obj = Bound::from_owned_ptr(py, value); - value = hook.call1((obj,)).map_err(|e| restore(py, e))?.into_ptr(); - } + }, + } + if f.owed > 0 { + break; + } + let done = self.stack.pop().expect("a frame was just read"); + value = match done.dict { + Some(dict) => dict.into_any(), + None => PyList::new(py, self.items.drain(done.base..)) + .map_err(|e| restore(py, e))? + .into_any(), + }; + if done.marker { + value = hook.call1((value,)).map_err(|e| restore(py, e))?; } } } @@ -265,19 +232,14 @@ impl<'a> Decoder<'a> { Ok(n) } - /// A one-byte element, already built as `obj`. + /// A one-byte element, already built as `obj` (dropped on a bound failure). #[inline] - fn scalar(&mut self, obj: *mut ffi::PyObject) -> Result<*mut ffi::PyObject, Fail> { - // Accounting first: on a bound failure the object is dropped, not leaked. - if let Err(e) = self.account(0, 0) { - // SAFETY: obj is a new reference or null. - unsafe { ffi::Py_XDECREF(obj) }; - return Err(e); - } - check(obj) + fn scalar(&mut self, obj: Bound<'py, PyAny>) -> Result>, Fail> { + self.account(0, 0)?; + Ok(Some(obj)) } - unsafe fn string(&mut self, len: u64) -> Result<*mut ffi::PyObject, Fail> { + fn string(&mut self, len: u64) -> Result>, Fail> { self.account(len, 0)?; let at_key = self.at_key(); let start = self.pos; @@ -293,22 +255,22 @@ impl<'a> Decoder<'a> { VALUE_CACHE_MAX_LEN }; if bytes.len() > max { - return check(decode_utf8(bytes)); + return self.decode_utf8(bytes).map(Some); } let i = slot_of(bytes); let held = if at_key { self.keys[i] } else { self.values[i] }; - if let Some(s) = (held as usize).checked_sub(1).map(|j| self.cached[j]) { + if let Some(s) = (held as usize).checked_sub(1).map(|j| &self.cached[j]) { if s.len == bytes.len() && &self.data[s.start..s.start + s.len] == bytes { - return Ok(ffi::Py_NewRef(s.obj)); + return Ok(Some(s.obj.clone())); } } // Not interned, unlike msgpack-python's keys: equal either way, and interning a distinct key // costs more than every later lookup of it saves. - let obj = check(decode_utf8(bytes))?; + let obj = self.decode_utf8(bytes)?; let slot = StrSlot { start, len: bytes.len(), - obj: ffi::Py_NewRef(obj), + obj: obj.clone(), }; if held == 0 { self.cached.push(slot); @@ -319,41 +281,50 @@ impl<'a> Decoder<'a> { &mut self.values[i] }) = self.cached.len() as u16; } else { - let old = std::mem::replace(&mut self.cached[held as usize - 1], slot); - ffi::Py_DECREF(old.obj); + self.cached[held as usize - 1] = slot; } - Ok(obj) + Ok(Some(obj)) + } + + /// `bytes` as a `str`, or the UnicodeDecodeError of the strict UTF-8 decode msgpack-python runs. + #[inline] + fn decode_utf8(&self, bytes: &[u8]) -> Result, Fail> { + PyString::from_bytes(self.py, bytes) + .map(Bound::into_any) + .map_err(|e| restore(self.py, e)) } - unsafe fn binary(&mut self, len: u64) -> Result<*mut ffi::PyObject, Fail> { + fn binary(&mut self, len: u64) -> Result>, Fail> { self.account(len, 0)?; let bytes = self.take(len); - check(ffi::PyBytes_FromStringAndSize( - bytes.as_ptr().cast::(), - bytes.len() as ffi::Py_ssize_t, - )) + Ok(Some(PyBytes::new(self.py, bytes).into_any())) } - unsafe fn number(&mut self, marker: u8) -> Result<*mut ffi::PyObject, Fail> { + fn number(&mut self, marker: u8) -> Result>, Fail> { + let py = self.py; let width = 1u64 << (marker & 0x03); self.account(width, 0)?; let raw = big_endian(self.take(width)); - check(match marker { - 0xca => ffi::PyFloat_FromDouble(f64::from(f32::from_bits(raw as u32))), - 0xcb => ffi::PyFloat_FromDouble(f64::from_bits(raw)), - 0xcc..=0xce => ffi::PyLong_FromLongLong(raw as i64), - 0xcf => ffi::PyLong_FromUnsignedLongLong(raw), - 0xd0 => ffi::PyLong_FromLong(c_long::from(raw as u8 as i8)), - 0xd1 => ffi::PyLong_FromLong(c_long::from(raw as u16 as i16)), - 0xd2 => ffi::PyLong_FromLong(c_long::from(raw as u32 as i32)), - _ => ffi::PyLong_FromLongLong(raw as i64), - }) + Ok(Some(match marker { + 0xca => PyFloat::new(py, f64::from(f32::from_bits(raw as u32))).into_any(), + 0xcb => PyFloat::new(py, f64::from_bits(raw)).into_any(), + 0xcf => match raw.into_pyobject(py) { + Ok(n) => n.into_any(), + Err(never) => match never {}, + }, + 0xd0 => int(py, i64::from(raw as u8 as i8)), + 0xd1 => int(py, i64::from(raw as u16 as i16)), + 0xd2 => int(py, i64::from(raw as u32 as i32)), + // 0xcc..=0xce (unsigned, at most 32 bits) and 0xd3 (signed 64 bits). + _ => int(py, raw as i64), + })) } /// An ext of `len` data bytes after its type byte, decoded as msgpack-python's defaults do: /// type -1 to `msgpack.Timestamp` (through its own `from_bytes`), any other to /// `msgpack.ExtType(code, data)`. cachekit never writes an ext, so this path is not hot. - unsafe fn ext(&mut self, py: Python<'_>, len: u64) -> Result<*mut ffi::PyObject, Fail> { + fn ext(&mut self, len: u64) -> Result>, Fail> { + let py = self.py; self.account(len + 1, 0)?; let code = self.take(1)[0] as i8; let data = self.take(len); @@ -363,66 +334,59 @@ impl<'a> Decoder<'a> { } else { msgpack_attr(py, "ExtType").and_then(|t| t.call1((code, PyBytes::new(py, data)))) }; - made.map(Bound::into_ptr).map_err(|e| restore(py, e)) + made.map(Some).map_err(|e| restore(py, e)) } - /// Open a container of `n` elements (`2n` slots for a map). Returns null after pushing a + /// Open a container of `n` elements (`2n` slots for a map). Returns None after pushing a /// non-empty one, or the empty container. - unsafe fn open(&mut self, is_map: bool, n: u64) -> Result<*mut ffi::PyObject, Fail> { + fn open(&mut self, is_map: bool, n: u64) -> Result>, Fail> { // An empty container is still a level, so depth is checked before the n > 0 push. if self.stack.len() >= self.max_depth { return Err(Fail::Bound); } let owed = if is_map { 2 * n } else { n }; self.account(0, owed)?; - // n <= the input length now, so it fits Py_ssize_t. - let obj = check(if is_map { - ffi::PyDict_New() - } else { - ffi::PyList_New(n as ffi::Py_ssize_t) - })?; if n == 0 { - return Ok(obj); + return Ok(Some(if is_map { + PyDict::new(self.py).into_any() + } else { + PyList::empty(self.py).into_any() + })); } self.stack.push(Frame { - obj, + dict: is_map.then(|| PyDict::new(self.py)), + key: None, + base: self.items.len(), owed, - index: 0, - key: ptr::null_mut(), - is_map, marker: false, }); - Ok(ptr::null_mut()) + Ok(None) } } +/// A Python `int`. #[inline] -fn big_endian(bytes: &[u8]) -> u64 { - bytes.iter().fold(0u64, |acc, b| (acc << 8) | u64::from(*b)) +fn int(py: Python<'_>, n: i64) -> Bound<'_, PyAny> { + match n.into_pyobject(py) { + Ok(n) => n.into_any(), + Err(never) => match never {}, + } } #[inline] -fn check(obj: *mut ffi::PyObject) -> Result<*mut ffi::PyObject, Fail> { - if obj.is_null() { - Err(Fail::Py) - } else { - Ok(obj) - } +fn big_endian(bytes: &[u8]) -> u64 { + bytes.iter().fold(0u64, |acc, b| (acc << 8) | u64::from(*b)) } -/// Put a fetched error back so the `Fail::Py` path finds it. +/// Set `e` as the current Python error, for the `Fail::Py` path to fetch. fn restore(py: Python<'_>, e: PyErr) -> Fail { e.restore(py); Fail::Py } -/// `strict_map_key=True`: a map key must be exactly `str` or `bytes`. -unsafe fn key_admitted(key: *mut ffi::PyObject) -> bool { - ffi::PyUnicode_CheckExact(key) != 0 || ffi::PyBytes_CheckExact(key) != 0 -} - -unsafe fn set_key_error(py: Python<'_>, key: *mut ffi::PyObject) { - let name = Bound::from_borrowed_ptr(py, key) +/// msgpack-python's `strict_map_key=True` refusal of `key`. +fn key_error(key: &Bound<'_, PyAny>) -> PyErr { + let name = key .get_type() .name() .map(|n| n.to_string()) @@ -430,7 +394,6 @@ unsafe fn set_key_error(py: Python<'_>, key: *mut ffi::PyObject) { PyValueError::new_err(format!( "{name:.100} is not allowed for map key when strict_map_key=True" )) - .restore(py); } fn msgpack_attr<'py>(py: Python<'py>, name: &str) -> PyResult> { @@ -449,14 +412,12 @@ pub fn unpackb_standard( max_depth: usize, hook: &Bound<'_, PyAny>, ) -> PyResult> { - let mut decoder = Decoder::new(data, max_depth); - let outcome = decoder.decode(py, hook); + let mut decoder = Decoder::new(py, data, max_depth); + let outcome = decoder.decode(hook); let consumed = decoder.pos; drop(decoder); let failure = match outcome { Ok(value) => { - // SAFETY: a strong reference the decode returned. - let value = unsafe { Bound::from_owned_ptr(py, value) }; if consumed == data.len() { return Ok(value.unbind()); } diff --git a/tests/unit/test_standard_decode.py b/tests/unit/test_standard_decode.py index bdeacaaa..81927a73 100644 --- a/tests/unit/test_standard_decode.py +++ b/tests/unit/test_standard_decode.py @@ -11,7 +11,8 @@ read paths' error handling see what they saw before; - the hook runs only on maps holding a marker key (the rule the speed rests on), and a heap amplification payload costs no more than it did; -- no path leaks a reference, and the decode is safe to run from many threads. +- no path leaks an object it allocates, nor a reference to its input, the input's buffer or the hook, + and the decode is safe to run from many threads. The inputs are values packed as the serializer packs them, the same with ``bytes`` packed as str (invalid UTF-8), and both after random byte edits, so the bounds, truncation, reserved-marker and @@ -21,6 +22,7 @@ from __future__ import annotations import gc +import sys import threading from collections.abc import Callable from datetime import UTC, datetime, timedelta, timezone @@ -326,8 +328,10 @@ def every_path() -> None: return sizes -# A leak on any path keeps at least one object per call, and the smallest object (an int) is 28 bytes, so it adds at -# least 500 * 28 = 13.7 KiB to EVERY round. The interpreter's own growth (caches, free lists, a dict resizing) comes +# This gate catches per-call allocations only: a path that keeps an object it allocated keeps at least one per call, +# and the smallest object (an int) is 28 bytes, so it adds at least 500 * 28 = 13.7 KiB to EVERY round. A leaked +# reference to an object that already exists (the input, its buffer, the hook) allocates nothing, so tracemalloc +# cannot see it: test_no_decode_keeps_a_reference_to_its_input counts those. The interpreter's own growth (caches, free lists, a dict resizing) comes # in a few one-off steps, and one such step can exceed any fixed budget on the total. So the gate bounds the quietest # of five rounds, which grows by under 1 KiB when nothing leaks. LEAK_ROUNDS = 5 @@ -336,12 +340,34 @@ def every_path() -> None: def test_no_path_leaks_a_reference() -> None: - """A leaked reference keeps its object alive, so traced memory would climb in every round.""" + """An object a decode allocates and never frees climbs traced memory every round (per-call allocations only).""" sizes = measure_in_subprocess(_measure_leaks) rounds = [after - before for before, after in zip(sizes, sizes[1:], strict=False)] assert min(rounds) < LEAK_ROUND_BUDGET, sizes +_MULTI_BYTE = {name: data for name, data in CASES.items() if len(data) > 1} + + +@pytest.mark.parametrize("data", _MULTI_BYTE.values(), ids=_MULTI_BYTE.keys()) +def test_no_decode_keeps_a_reference_to_its_input(data: bytes) -> None: + """Every path, from every buffer type, gives back each reference it took to the input, its buffer and the hook.""" + owned = bytes(bytearray(data)) # a fresh object: a code constant may be immortal, and its count never moves + view, buffer = memoryview(owned), bytearray(data) + assert view.obj is owned + watched = (owned, view, buffer, _standard_object_hook) + before = [sys.getrefcount(o) for o in watched] + for source in (owned, view, buffer): + for _ in range(100): + try: + rust(source) + except (ValueError, SerializationError, TypeError): # noqa: S110 - every refusal CASES pins; any other error fails + pass + del source # still bound to the last input + assert [sys.getrefcount(o) for o in watched] == before + buffer.append(0) # a buffer export never released refuses the resize with BufferError + + def test_threads_decode_one_document_concurrently() -> None: doc = _pack({"users": [{"id": i, "name": f"user {i}", "at": datetime(2024, 1, 15)} for i in range(50)]}, True) expected, failures, barrier = reference(doc), [], threading.Barrier(8) From f5bb32d039b1d2654b7612ccc72dfc3c8ed39ded Mon Sep 17 00:00:00 2001 From: Ray Walker Date: Sun, 1 Nov 2026 03:47:34 +0000 Subject: [PATCH 2/9] fix(rust): preallocate lists and raise MemoryError on input-sized allocations --- rust/src/msgpack_decode.rs | 86 ++++++++++++++++++++---------- tests/unit/test_standard_decode.py | 30 +++++++++++ 2 files changed, 89 insertions(+), 27 deletions(-) diff --git a/rust/src/msgpack_decode.rs b/rust/src/msgpack_decode.rs index d32ba5c4..8fe6b75b 100644 --- a/rust/src/msgpack_decode.rs +++ b/rust/src/msgpack_decode.rs @@ -14,14 +14,15 @@ //! return unchanged, because each marker test reads its key with `dict.get`. //! //! Safe Rust only (`SECURITY.md`): every object is a PyO3 `Bound`, so each reference is released -//! by its owner on any exit, and an array's items wait on a stack until it closes, so no -//! partly-filled list is ever built. +//! by its owner on any exit. PyO3's plain constructors abort the process when CPython is out of +//! memory (the release profile sets `panic = "abort"`), so the two allocations the input sizes, an +//! array's list and a bin's bytes, go through ones that raise MemoryError, as the old decode did. use cachekit_core::check_msgpack_structure; use pyo3::exceptions::PyValueError; use pyo3::prelude::*; use pyo3::sync::PyOnceLock; -use pyo3::types::{PyBool, PyBytes, PyDict, PyFloat, PyList, PyString}; +use pyo3::types::{PyBool, PyBytes, PyDict, PyFloat, PyList, PySequence, PyString}; /// The keys `_standard_object_hook` tests (`standard_serializer.py`). const MARKER_KEYS: [&[u8]; 3] = [b"__datetime__", b"__date__", b"__time__"]; @@ -40,14 +41,17 @@ enum Fail { Bound, } +/// An open container, filled as each element arrives. +enum Open<'py> { + /// An array: its list, made with `None` in every slot, and the next index to fill. + List(Bound<'py, PyList>, usize), + /// A map: its dict, and the key read, awaiting its value. + Map(Bound<'py, PyDict>, Option>), +} + /// An open array or map. struct Frame<'py> { - /// Map: the dict, filled as each value arrives. Array: None. - dict: Option>, - /// Map: the key read, awaiting its value. - key: Option>, - /// Array: where its items start on the decoder's `items` stack. - base: usize, + open: Open<'py>, /// Elements still to read: items for an array, keys plus values for a map. owed: u64, /// Map: a `str` key equal to a marker key was read. @@ -69,9 +73,8 @@ struct Decoder<'a, 'py> { /// Elements owed across every open container, the root included (the walk's `pending`). pending: u64, stack: Vec>, - /// The items of every open array, innermost last; an array's become its list when it closes. - /// Only decoded elements land here, so it never outgrows the input. - items: Vec>, + /// `[None]`, made at the first non-empty array, for `nones`. + seed: Option>, /// Cached strings; a table slot holds an index into it plus one, or 0 when empty. A slot holds /// one string at a time, a colliding one replacing it, so this never exceeds `2 * CACHE_SLOTS`. cached: Vec>, @@ -103,7 +106,7 @@ impl<'a, 'py> Decoder<'a, 'py> { max_depth, pending: 1, stack: Vec::new(), - items: Vec::new(), + seed: None, cached: Vec::new(), keys: [0; CACHE_SLOTS], values: [0; CACHE_SLOTS], @@ -137,7 +140,13 @@ impl<'a, 'py> Decoder<'a, 'py> { /// Whether the next element read is a map key. #[inline] fn at_key(&self) -> bool { - matches!(self.stack.last(), Some(f) if f.dict.is_some() && f.key.is_none()) + matches!( + self.stack.last(), + Some(Frame { + open: Open::Map(_, None), + .. + }) + ) } fn decode(&mut self, hook: &Bound<'py, PyAny>) -> Result, Fail> { @@ -184,11 +193,14 @@ impl<'a, 'py> Decoder<'a, 'py> { let Some(f) = self.stack.last_mut() else { return Ok(value); }; - match &f.dict { - None => self.items.push(value), - Some(dict) => match f.key.take() { + match &mut f.open { + Open::List(list, index) => { + list.set_item(*index, value).map_err(|e| restore(py, e))?; + *index += 1; + } + Open::Map(dict, pending) => match pending.take() { None => { - f.key = Some(value); + *pending = Some(value); break; } Some(key) => { @@ -206,11 +218,9 @@ impl<'a, 'py> Decoder<'a, 'py> { break; } let done = self.stack.pop().expect("a frame was just read"); - value = match done.dict { - Some(dict) => dict.into_any(), - None => PyList::new(py, self.items.drain(done.base..)) - .map_err(|e| restore(py, e))? - .into_any(), + value = match done.open { + Open::List(list, _) => list.into_any(), + Open::Map(dict, _) => dict.into_any(), }; if done.marker { value = hook.call1((value,)).map_err(|e| restore(py, e))?; @@ -297,7 +307,13 @@ impl<'a, 'py> Decoder<'a, 'py> { fn binary(&mut self, len: u64) -> Result>, Fail> { self.account(len, 0)?; let bytes = self.take(len); - Ok(Some(PyBytes::new(self.py, bytes).into_any())) + // new_with, not new: an input-sized allocation that fails raises MemoryError, not abort. + PyBytes::new_with(self.py, bytes.len(), |b| { + b.copy_from_slice(bytes); + Ok(()) + }) + .map(|b| Some(b.into_any())) + .map_err(|e| restore(self.py, e)) } fn number(&mut self, marker: u8) -> Result>, Fail> { @@ -353,15 +369,31 @@ impl<'a, 'py> Decoder<'a, 'py> { PyList::empty(self.py).into_any() })); } + let open = if is_map { + Open::Map(PyDict::new(self.py), None) + } else { + Open::List(self.nones(n).map_err(|e| restore(self.py, e))?, 0) + }; self.stack.push(Frame { - dict: is_map.then(|| PyDict::new(self.py)), - key: None, - base: self.items.len(), + open, owed, marker: false, }); Ok(None) } + + /// A list of `n` `None`s, which the decode replaces in order: `[None] * n`, which raises + /// MemoryError where PyO3's list constructors would abort. `n` is at most the input length. + fn nones(&mut self, n: u64) -> PyResult> { + let py = self.py; + let seed = match self.seed.take() { + Some(seed) => seed, + None => PyList::new(py, [py.None()])?.into_sequence(), + }; + let list = seed.repeat(n as usize); + self.seed = Some(seed); + Ok(list?.cast_into::()?) + } } /// A Python `int`. diff --git a/tests/unit/test_standard_decode.py b/tests/unit/test_standard_decode.py index 81927a73..6f46343e 100644 --- a/tests/unit/test_standard_decode.py +++ b/tests/unit/test_standard_decode.py @@ -22,6 +22,7 @@ from __future__ import annotations import gc +import subprocess import sys import threading from collections.abc import Callable @@ -304,6 +305,35 @@ def test_an_amplifying_payload_costs_no_more_heap_than_before() -> None: assert now <= before, f"{name}: {now} bytes at peak, {before} before" +_OUT_OF_MEMORY = """ +import resource, sys +from cachekit._rust_serializer import unpackb_standard +n = 1 << 27 +data = {shape} +with open("/proc/self/statm") as f: + mapped = int(f.read().split()[0]) * resource.getpagesize() +resource.setrlimit(resource.RLIMIT_AS, (mapped + (64 << 20), resource.RLIM_INFINITY)) +try: + unpackb_standard(data, 1024, None) +except MemoryError: + sys.exit(0) +sys.exit(3) +""" + + +@pytest.mark.skipif(not sys.platform.startswith("linux"), reason="RLIMIT_AS is enforced on Linux") +@pytest.mark.parametrize( + "shape", + ['b"\\xc6" + n.to_bytes(4, "big") + bytes(n)', 'b"\\xdd" + (n >> 3).to_bytes(4, "big") + b"\\xc0" * (n >> 3)'], + ids=["128 MiB bin", "16M-item array"], +) +def test_an_allocation_the_input_sizes_raises_memory_error(shape: str) -> None: + """With panic = "abort", a PyO3 constructor that meets a NULL would kill the process, not raise.""" + script = _OUT_OF_MEMORY.format(shape=shape) + run = subprocess.run([sys.executable, "-c", script], capture_output=True, check=False) # noqa: S603 - this module's own script + assert run.returncode == 0, run.stderr.decode()[-2000:] + + def _measure_leaks() -> list[int]: import tracemalloc From c51cb78590bb46dc94dd125749f357cebca81371 Mon Sep 17 00:00:00 2001 From: Ray Walker Date: Sun, 1 Nov 2026 04:22:45 +0000 Subject: [PATCH 3/9] test: skip the out-of-memory decode check on free-threaded builds; refresh decode figures --- docs/performance.md | 2 +- tests/unit/test_standard_decode.py | 5 +++++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/docs/performance.md b/docs/performance.md index 9c0d0685..ffd66372 100644 --- a/docs/performance.md +++ b/docs/performance.md @@ -73,7 +73,7 @@ With `backend=None`, cachekit keeps the returned object itself in memory and han With an L2 backend configured, L1 holds the serialized bytes instead, and a hit deserializes them, so that hit costs more as the payload grows. For the 100-user dict it took 284 ± 2μs on the same runner, measured before the decode changes described below, about 910 times an L1-only hit of the same dict. Nearly all of the difference is deserializing 23.5KB of MessagePack back into about 400 nested dicts. It is still about half of a Redis L2 hit there (539 ± 3μs), which deserializes the same bytes after the round trip. -Every read with the default serializer decodes its MessagePack in the Rust extension, in one pass that also enforces the bounds on untrusted MessagePack: nesting at most 1,024 levels deep, and no header declaring more elements or bytes than the rest of the input can back. That covers L2 reads, L1 hits and the decrypted plaintext of encrypted reads. The decode builds the Python objects directly. It shares most repeated short strings, map keys up to 64 bytes and values up to 32, as one object within a value (a small table per decode; two strings that land in one slot replace each other). It runs the Python hook that restores datetimes only on maps holding a `__datetime__`, `__date__` or `__time__` key, because the hook hands back any other map unchanged. It returns the same values and refuses the same inputs with the same exception classes as the decode it replaced: cachekit-core's structure walk, then msgpack-python's C unpacker calling the hook once per map. On the 100-user dict (23,547 bytes of MessagePack, about 3,400 objects), the decode fell from 2,580,104 to 1,238,204 instructions on CPython 3.14 (−52.0%, callgrind main-thread instructions per call, the [Instruction Budgets](#instruction-budgets) method), level with `pickle.loads` of the same value (1,203,437). The 1,000-user dict (239 KB) fell 51.0% and a 1,090-byte one 45.3%. A list of 200 records each holding a datetime fell 33.2%, since every datetime still runs the hook. On the abi3t extension that CPython 3.15 and later load, the 100-user dict fell 50.5%. +Every read with the default serializer decodes its MessagePack in the Rust extension, in one pass that also enforces the bounds on untrusted MessagePack: nesting at most 1,024 levels deep, and no header declaring more elements or bytes than the rest of the input can back. That covers L2 reads, L1 hits and the decrypted plaintext of encrypted reads. The decode builds the Python objects directly. It shares most repeated short strings, map keys up to 64 bytes and values up to 32, as one object within a value (a small table per decode; two strings that land in one slot replace each other). It runs the Python hook that restores datetimes only on maps holding a `__datetime__`, `__date__` or `__time__` key, because the hook hands back any other map unchanged. It returns the same values and refuses the same inputs with the same exception classes as the decode it replaced: cachekit-core's structure walk, then msgpack-python's C unpacker calling the hook once per map. On the 100-user dict (23,547 bytes of MessagePack, about 3,400 objects), the decode fell from 2,584,067 to 1,230,047 instructions on CPython 3.14 (−52.4%, callgrind main-thread instructions per call, the [Instruction Budgets](#instruction-budgets) method), level with `pickle.loads` of the same value (1,197,777). The 1,000-user dict (239 KB) fell 51.3% and a 1,090-byte one 45.1%. A list of 200 records each holding a datetime fell 33.8%, since every datetime still runs the hook. On the abi3t extension that CPython 3.15 and later load, the 100-user dict fell 50.8%. For a plaintext cache with the default serializer and integrity checking on, an entry's first L1 hit decodes it as every read does, and then, when its verified MessagePack is no longer than the envelope, which is true of small values and incompressible ones, L1 holds that MessagePack in the envelope's place. Later hits on it skip the envelope's LZ4 decompress and xxHash3-64 checksum. Every hit still runs the bounded decode, which costs no more than the unbounded one did. The 100-user dict compresses to a 3,171-byte envelope from 23.5KB of MessagePack, so L1 keeps its envelope; no budget gates this large-value hit, and `l1_hit_backed` in [Instruction Budgets](#instruction-budgets) gates the small-value hit on held MessagePack. Encrypted caches always keep the ciphertext envelope, and every hit on one decrypts and decodes. diff --git a/tests/unit/test_standard_decode.py b/tests/unit/test_standard_decode.py index 6f46343e..ad8d9fc6 100644 --- a/tests/unit/test_standard_decode.py +++ b/tests/unit/test_standard_decode.py @@ -24,6 +24,7 @@ import gc import subprocess import sys +import sysconfig import threading from collections.abc import Callable from datetime import UTC, datetime, timedelta, timezone @@ -322,6 +323,10 @@ def test_an_amplifying_payload_costs_no_more_heap_than_before() -> None: @pytest.mark.skipif(not sys.platform.startswith("linux"), reason="RLIMIT_AS is enforced on Linux") +@pytest.mark.skipif( + bool(sysconfig.get_config_var("Py_GIL_DISABLED")), + reason="mimalloc's reserved arenas hold the allocation, so RLIMIT_AS never refuses it; the GIL legs run the same Rust", +) @pytest.mark.parametrize( "shape", ['b"\\xc6" + n.to_bytes(4, "big") + bytes(n)', 'b"\\xdd" + (n >> 3).to_bytes(4, "big") + b"\\xc0" * (n >> 3)'], From c16b826833072eef32708d56ef8a2284fd91106c Mon Sep 17 00:00:00 2001 From: Ray Walker Date: Sun, 1 Nov 2026 04:24:14 +0000 Subject: [PATCH 4/9] fix(rust): forbid unsafe code in the MessagePack decoder module --- rust/src/msgpack_decode.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/rust/src/msgpack_decode.rs b/rust/src/msgpack_decode.rs index 8fe6b75b..859b6a25 100644 --- a/rust/src/msgpack_decode.rs +++ b/rust/src/msgpack_decode.rs @@ -18,6 +18,9 @@ //! memory (the release profile sets `panic = "abort"`), so the two allocations the input sizes, an //! array's list and a bin's bytes, go through ones that raise MemoryError, as the old decode did. +// cargo-geiger checks the crate only after a merge; this fails any build, a PR's included. +#![forbid(unsafe_code)] + use cachekit_core::check_msgpack_structure; use pyo3::exceptions::PyValueError; use pyo3::prelude::*; From ca227c4a601995edf92d8ab6cbd25e724ee87518 Mon Sep 17 00:00:00 2001 From: Ray Walker Date: Sun, 1 Nov 2026 04:38:00 +0000 Subject: [PATCH 5/9] perf(rust): make lists of up to 64 items with the plain constructor --- rust/src/msgpack_decode.rs | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/rust/src/msgpack_decode.rs b/rust/src/msgpack_decode.rs index 859b6a25..999ac2fe 100644 --- a/rust/src/msgpack_decode.rs +++ b/rust/src/msgpack_decode.rs @@ -37,6 +37,9 @@ const VALUE_CACHE_MAX_LEN: usize = 32; const CACHE_BITS: u32 = 8; const CACHE_SLOTS: usize = 1 << CACHE_BITS; +/// Lists at most this long (512 bytes of slots) are made by PyO3's plain constructor (`nones`). +const SHORT_LIST: u64 = 64; + /// A Python error is set (`restore` put it there, `PyErr::fetch` takes it), or the input broke a /// bound. A unit enum, so a decode step's result stays two words. enum Fail { @@ -385,10 +388,18 @@ impl<'a, 'py> Decoder<'a, 'py> { Ok(None) } - /// A list of `n` `None`s, which the decode replaces in order: `[None] * n`, which raises - /// MemoryError where PyO3's list constructors would abort. `n` is at most the input length. + /// A list of `n` `None`s, which the decode replaces in order. `n` is at most the input length, + /// so a long list is made as `[None] * n`, which raises MemoryError where PyO3's list + /// constructors would abort. A short one, no bigger than an int, comes from PyO3's: the + /// repeat costs most small documents about 1% of their decode. fn nones(&mut self, n: u64) -> PyResult> { let py = self.py; + if n <= SHORT_LIST { + return PyList::new( + py, + std::iter::repeat_n(py.None().into_bound(py), n as usize), + ); + } let seed = match self.seed.take() { Some(seed) => seed, None => PyList::new(py, [py.None()])?.into_sequence(), From 33823bc1f936bad3e3d63d54dc5e1dc2ca1ff9b5 Mon Sep 17 00:00:00 2001 From: Ray Walker Date: Sun, 1 Nov 2026 04:50:26 +0000 Subject: [PATCH 6/9] fix(rust): raise MemoryError for ext and trailing-byte copies too; drop the seed list --- rust/src/msgpack_decode.rs | 90 +++++++++++++++--------------- tests/unit/test_standard_decode.py | 26 +++++---- 2 files changed, 60 insertions(+), 56 deletions(-) diff --git a/rust/src/msgpack_decode.rs b/rust/src/msgpack_decode.rs index 999ac2fe..bd3a96d9 100644 --- a/rust/src/msgpack_decode.rs +++ b/rust/src/msgpack_decode.rs @@ -15,17 +15,20 @@ //! //! Safe Rust only (`SECURITY.md`): every object is a PyO3 `Bound`, so each reference is released //! by its owner on any exit. PyO3's plain constructors abort the process when CPython is out of -//! memory (the release profile sets `panic = "abort"`), so the two allocations the input sizes, an -//! array's list and a bin's bytes, go through ones that raise MemoryError, as the old decode did. +//! memory (the release profile sets `panic = "abort"`), so every allocation the input sizes, a +//! long array's list and the bytes of a bin, an ext or `ExtraData`, goes through a constructor +//! that raises MemoryError instead (`nones`, `bytes_of`). // cargo-geiger checks the crate only after a merge; this fails any build, a PR's included. #![forbid(unsafe_code)] +use std::convert::Infallible; + use cachekit_core::check_msgpack_structure; use pyo3::exceptions::PyValueError; use pyo3::prelude::*; use pyo3::sync::PyOnceLock; -use pyo3::types::{PyBool, PyBytes, PyDict, PyFloat, PyList, PySequence, PyString}; +use pyo3::types::{PyBool, PyBytes, PyDict, PyFloat, PyInt, PyList, PyString}; /// The keys `_standard_object_hook` tests (`standard_serializer.py`). const MARKER_KEYS: [&[u8]; 3] = [b"__datetime__", b"__date__", b"__time__"]; @@ -41,7 +44,7 @@ const CACHE_SLOTS: usize = 1 << CACHE_BITS; const SHORT_LIST: u64 = 64; /// A Python error is set (`restore` put it there, `PyErr::fetch` takes it), or the input broke a -/// bound. A unit enum, so a decode step's result stays two words. +/// bound. A fieldless enum, so a decode step's result stays two words. enum Fail { Py, Bound, @@ -79,8 +82,6 @@ struct Decoder<'a, 'py> { /// Elements owed across every open container, the root included (the walk's `pending`). pending: u64, stack: Vec>, - /// `[None]`, made at the first non-empty array, for `nones`. - seed: Option>, /// Cached strings; a table slot holds an index into it plus one, or 0 when empty. A slot holds /// one string at a time, a colliding one replacing it, so this never exceeds `2 * CACHE_SLOTS`. cached: Vec>, @@ -112,7 +113,6 @@ impl<'a, 'py> Decoder<'a, 'py> { max_depth, pending: 1, stack: Vec::new(), - seed: None, cached: Vec::new(), keys: [0; CACHE_SLOTS], values: [0; CACHE_SLOTS], @@ -313,13 +313,9 @@ impl<'a, 'py> Decoder<'a, 'py> { fn binary(&mut self, len: u64) -> Result>, Fail> { self.account(len, 0)?; let bytes = self.take(len); - // new_with, not new: an input-sized allocation that fails raises MemoryError, not abort. - PyBytes::new_with(self.py, bytes.len(), |b| { - b.copy_from_slice(bytes); - Ok(()) - }) - .map(|b| Some(b.into_any())) - .map_err(|e| restore(self.py, e)) + bytes_of(self.py, bytes) + .map(|b| Some(b.into_any())) + .map_err(|e| restore(self.py, e)) } fn number(&mut self, marker: u8) -> Result>, Fail> { @@ -330,10 +326,7 @@ impl<'a, 'py> Decoder<'a, 'py> { Ok(Some(match marker { 0xca => PyFloat::new(py, f64::from(f32::from_bits(raw as u32))).into_any(), 0xcb => PyFloat::new(py, f64::from_bits(raw)).into_any(), - 0xcf => match raw.into_pyobject(py) { - Ok(n) => n.into_any(), - Err(never) => match never {}, - }, + 0xcf => int(py, raw), 0xd0 => int(py, i64::from(raw as u8 as i8)), 0xd1 => int(py, i64::from(raw as u16 as i16)), 0xd2 => int(py, i64::from(raw as u32 as i32)), @@ -350,12 +343,13 @@ impl<'a, 'py> Decoder<'a, 'py> { self.account(len + 1, 0)?; let code = self.take(1)[0] as i8; let data = self.take(len); - let made = if code == -1 { - msgpack_attr(py, "Timestamp") - .and_then(|t| t.call_method1("from_bytes", (PyBytes::new(py, data),))) - } else { - msgpack_attr(py, "ExtType").and_then(|t| t.call1((code, PyBytes::new(py, data)))) - }; + let made = bytes_of(py, data).and_then(|data| { + if code == -1 { + msgpack_attr(py, "Timestamp").and_then(|t| t.call_method1("from_bytes", (data,))) + } else { + msgpack_attr(py, "ExtType").and_then(|t| t.call1((code, data))) + } + }); made.map(Some).map_err(|e| restore(py, e)) } @@ -389,34 +383,40 @@ impl<'a, 'py> Decoder<'a, 'py> { } /// A list of `n` `None`s, which the decode replaces in order. `n` is at most the input length, - /// so a long list is made as `[None] * n`, which raises MemoryError where PyO3's list - /// constructors would abort. A short one, no bigger than an int, comes from PyO3's: the - /// repeat costs most small documents about 1% of their decode. - fn nones(&mut self, n: u64) -> PyResult> { + /// so a list longer than `SHORT_LIST` is made as `[None] * n`, which raises MemoryError where + /// PyO3's list constructor would abort. A shorter one, at most 512 bytes of slots, comes from + /// that constructor: making it by repeat would add about 1% to a small document's decode. + fn nones(&self, n: u64) -> PyResult> { let py = self.py; + let none = py.None().into_bound(py); if n <= SHORT_LIST { - return PyList::new( - py, - std::iter::repeat_n(py.None().into_bound(py), n as usize), - ); + return PyList::new(py, std::iter::repeat_n(none, n as usize)); } - let seed = match self.seed.take() { - Some(seed) => seed, - None => PyList::new(py, [py.None()])?.into_sequence(), - }; - let list = seed.repeat(n as usize); - self.seed = Some(seed); - Ok(list?.cast_into::()?) + let seed = PyList::new(py, [none])?; + Ok(seed + .as_sequence() + .repeat(n as usize)? + .cast_into::()?) } } /// A Python `int`. #[inline] -fn int(py: Python<'_>, n: i64) -> Bound<'_, PyAny> { - match n.into_pyobject(py) { - Ok(n) => n.into_any(), - Err(never) => match never {}, - } +fn int<'py>( + py: Python<'py>, + n: impl IntoPyObject<'py, Output = Bound<'py, PyInt>, Error = Infallible>, +) -> Bound<'py, PyAny> { + let Ok(n) = n.into_pyobject(py); + n.into_any() +} + +/// `bytes` holding `data`. `new_with`, not `new`: it raises MemoryError where `new` would abort, +/// at the cost of zeroing the buffer before the copy. +fn bytes_of<'py>(py: Python<'py>, data: &[u8]) -> PyResult> { + PyBytes::new_with(py, data.len(), |b| { + b.copy_from_slice(data); + Ok(()) + }) } #[inline] @@ -469,7 +469,7 @@ pub fn unpackb_standard( } let extra = msgpack_attr(py, "ExtraData")?; return Err(PyErr::from_value( - extra.call1((value, PyBytes::new(py, &data[consumed..])))?, + extra.call1((value, bytes_of(py, &data[consumed..])?))?, )); } Err(Fail::Py) => PyErr::fetch(py), diff --git a/tests/unit/test_standard_decode.py b/tests/unit/test_standard_decode.py index ad8d9fc6..f6922f7e 100644 --- a/tests/unit/test_standard_decode.py +++ b/tests/unit/test_standard_decode.py @@ -313,13 +313,20 @@ def test_an_amplifying_payload_costs_no_more_heap_than_before() -> None: data = {shape} with open("/proc/self/statm") as f: mapped = int(f.read().split()[0]) * resource.getpagesize() -resource.setrlimit(resource.RLIMIT_AS, (mapped + (64 << 20), resource.RLIM_INFINITY)) +resource.setrlimit(resource.RLIMIT_AS, (mapped + (64 << 20), resource.getrlimit(resource.RLIMIT_AS)[1])) try: unpackb_standard(data, 1024, None) except MemoryError: sys.exit(0) sys.exit(3) """ +# Every allocation the input sizes (n is 128 MiB), as the `data` expression the script decodes. +_OUT_OF_MEMORY_SHAPES = { + "bin": 'b"\\xc6" + n.to_bytes(4, "big") + bytes(n)', + "ext": 'b"\\xc9" + n.to_bytes(4, "big") + b"\\x05" + bytes(n)', + "trailing bytes": 'b"\\xc0" + bytes(n)', + "array of n / 8 items": 'b"\\xdd" + (n >> 3).to_bytes(4, "big") + b"\\xc0" * (n >> 3)', +} @pytest.mark.skipif(not sys.platform.startswith("linux"), reason="RLIMIT_AS is enforced on Linux") @@ -327,11 +334,7 @@ def test_an_amplifying_payload_costs_no_more_heap_than_before() -> None: bool(sysconfig.get_config_var("Py_GIL_DISABLED")), reason="mimalloc's reserved arenas hold the allocation, so RLIMIT_AS never refuses it; the GIL legs run the same Rust", ) -@pytest.mark.parametrize( - "shape", - ['b"\\xc6" + n.to_bytes(4, "big") + bytes(n)', 'b"\\xdd" + (n >> 3).to_bytes(4, "big") + b"\\xc0" * (n >> 3)'], - ids=["128 MiB bin", "16M-item array"], -) +@pytest.mark.parametrize("shape", _OUT_OF_MEMORY_SHAPES.values(), ids=_OUT_OF_MEMORY_SHAPES.keys()) def test_an_allocation_the_input_sizes_raises_memory_error(shape: str) -> None: """With panic = "abort", a PyO3 constructor that meets a NULL would kill the process, not raise.""" script = _OUT_OF_MEMORY.format(shape=shape) @@ -366,28 +369,29 @@ def every_path() -> None: # This gate catches per-call allocations only: a path that keeps an object it allocated keeps at least one per call, # and the smallest object (an int) is 28 bytes, so it adds at least 500 * 28 = 13.7 KiB to EVERY round. A leaked # reference to an object that already exists (the input, its buffer, the hook) allocates nothing, so tracemalloc -# cannot see it: test_no_decode_keeps_a_reference_to_its_input counts those. The interpreter's own growth (caches, free lists, a dict resizing) comes -# in a few one-off steps, and one such step can exceed any fixed budget on the total. So the gate bounds the quietest -# of five rounds, which grows by under 1 KiB when nothing leaks. +# cannot see it: test_no_decode_keeps_a_reference_to_its_input counts those. The interpreter's own growth (caches, +# free lists, a dict resizing) comes in a few one-off steps, and one such step can exceed any fixed budget on the +# total. So the gate bounds the quietest of five rounds, which grows by under 1 KiB when nothing leaks. LEAK_ROUNDS = 5 LEAK_CALLS_PER_ROUND = 500 LEAK_ROUND_BUDGET = 8 * 1024 -def test_no_path_leaks_a_reference() -> None: +def test_no_path_leaks_an_object() -> None: """An object a decode allocates and never frees climbs traced memory every round (per-call allocations only).""" sizes = measure_in_subprocess(_measure_leaks) rounds = [after - before for before, after in zip(sizes, sizes[1:], strict=False)] assert min(rounds) < LEAK_ROUND_BUDGET, sizes +# Empty and one-byte bytes are immortal singletons, whose count never moves. _MULTI_BYTE = {name: data for name, data in CASES.items() if len(data) > 1} @pytest.mark.parametrize("data", _MULTI_BYTE.values(), ids=_MULTI_BYTE.keys()) def test_no_decode_keeps_a_reference_to_its_input(data: bytes) -> None: """Every path, from every buffer type, gives back each reference it took to the input, its buffer and the hook.""" - owned = bytes(bytearray(data)) # a fresh object: a code constant may be immortal, and its count never moves + owned = bytes(bytearray(data)) # a fresh object: a code constant may be immortal too view, buffer = memoryview(owned), bytearray(data) assert view.obj is owned watched = (owned, view, buffer, _standard_object_hook) From 71ef971bbb4c0622527321ec728a867928db8bf4 Mon Sep 17 00:00:00 2001 From: Ray Walker Date: Sun, 1 Nov 2026 05:14:50 +0000 Subject: [PATCH 7/9] docs(performance): decode figures for the safe decoder --- docs/performance.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/performance.md b/docs/performance.md index ffd66372..dbe18141 100644 --- a/docs/performance.md +++ b/docs/performance.md @@ -73,7 +73,7 @@ With `backend=None`, cachekit keeps the returned object itself in memory and han With an L2 backend configured, L1 holds the serialized bytes instead, and a hit deserializes them, so that hit costs more as the payload grows. For the 100-user dict it took 284 ± 2μs on the same runner, measured before the decode changes described below, about 910 times an L1-only hit of the same dict. Nearly all of the difference is deserializing 23.5KB of MessagePack back into about 400 nested dicts. It is still about half of a Redis L2 hit there (539 ± 3μs), which deserializes the same bytes after the round trip. -Every read with the default serializer decodes its MessagePack in the Rust extension, in one pass that also enforces the bounds on untrusted MessagePack: nesting at most 1,024 levels deep, and no header declaring more elements or bytes than the rest of the input can back. That covers L2 reads, L1 hits and the decrypted plaintext of encrypted reads. The decode builds the Python objects directly. It shares most repeated short strings, map keys up to 64 bytes and values up to 32, as one object within a value (a small table per decode; two strings that land in one slot replace each other). It runs the Python hook that restores datetimes only on maps holding a `__datetime__`, `__date__` or `__time__` key, because the hook hands back any other map unchanged. It returns the same values and refuses the same inputs with the same exception classes as the decode it replaced: cachekit-core's structure walk, then msgpack-python's C unpacker calling the hook once per map. On the 100-user dict (23,547 bytes of MessagePack, about 3,400 objects), the decode fell from 2,584,067 to 1,230,047 instructions on CPython 3.14 (−52.4%, callgrind main-thread instructions per call, the [Instruction Budgets](#instruction-budgets) method), level with `pickle.loads` of the same value (1,197,777). The 1,000-user dict (239 KB) fell 51.3% and a 1,090-byte one 45.1%. A list of 200 records each holding a datetime fell 33.8%, since every datetime still runs the hook. On the abi3t extension that CPython 3.15 and later load, the 100-user dict fell 50.8%. +Every read with the default serializer decodes its MessagePack in the Rust extension, in one pass that also enforces the bounds on untrusted MessagePack: nesting at most 1,024 levels deep, and no header declaring more elements or bytes than the rest of the input can back. That covers L2 reads, L1 hits and the decrypted plaintext of encrypted reads. The decode builds the Python objects directly. It shares most repeated short strings, map keys up to 64 bytes and values up to 32, as one object within a value (a small table per decode; two strings that land in one slot replace each other). It runs the Python hook that restores datetimes only on maps holding a `__datetime__`, `__date__` or `__time__` key, because the hook hands back any other map unchanged. It returns the same values and refuses the same inputs with the same exception classes as the decode it replaced: cachekit-core's structure walk, then msgpack-python's C unpacker calling the hook once per map. On the 100-user dict (23,547 bytes of MessagePack, about 3,400 objects), the decode fell from 2,577,967 to 1,220,541 instructions on CPython 3.14 (−52.7%, callgrind main-thread instructions per call, the [Instruction Budgets](#instruction-budgets) method), level with `pickle.loads` of the same value (1,199,580). The 1,000-user dict (239 KB) fell 51.5% and a 1,090-byte one 45.8%. A list of 200 records each holding a datetime fell 33.9%, since every datetime still runs the hook. On the abi3t extension that CPython 3.15 and later load, the 100-user dict fell 51.0%. For a plaintext cache with the default serializer and integrity checking on, an entry's first L1 hit decodes it as every read does, and then, when its verified MessagePack is no longer than the envelope, which is true of small values and incompressible ones, L1 holds that MessagePack in the envelope's place. Later hits on it skip the envelope's LZ4 decompress and xxHash3-64 checksum. Every hit still runs the bounded decode, which costs no more than the unbounded one did. The 100-user dict compresses to a 3,171-byte envelope from 23.5KB of MessagePack, so L1 keeps its envelope; no budget gates this large-value hit, and `l1_hit_backed` in [Instruction Budgets](#instruction-budgets) gates the small-value hit on held MessagePack. Encrypted caches always keep the ciphertext envelope, and every hit on one decrypts and decodes. From 12ab5e19b200f25606dd1aa434ecfbfd95e88746 Mon Sep 17 00:00:00 2001 From: Ray Walker Date: Sun, 1 Nov 2026 06:01:08 +0000 Subject: [PATCH 8/9] docs: state which out-of-memory paths abort; review cleanups --- docs/performance.md | 2 +- rust/src/msgpack_decode.rs | 17 ++++++++++------- tests/unit/test_standard_decode.py | 26 +++++++++++++++----------- 3 files changed, 26 insertions(+), 19 deletions(-) diff --git a/docs/performance.md b/docs/performance.md index dbe18141..8c50cadb 100644 --- a/docs/performance.md +++ b/docs/performance.md @@ -73,7 +73,7 @@ With `backend=None`, cachekit keeps the returned object itself in memory and han With an L2 backend configured, L1 holds the serialized bytes instead, and a hit deserializes them, so that hit costs more as the payload grows. For the 100-user dict it took 284 ± 2μs on the same runner, measured before the decode changes described below, about 910 times an L1-only hit of the same dict. Nearly all of the difference is deserializing 23.5KB of MessagePack back into about 400 nested dicts. It is still about half of a Redis L2 hit there (539 ± 3μs), which deserializes the same bytes after the round trip. -Every read with the default serializer decodes its MessagePack in the Rust extension, in one pass that also enforces the bounds on untrusted MessagePack: nesting at most 1,024 levels deep, and no header declaring more elements or bytes than the rest of the input can back. That covers L2 reads, L1 hits and the decrypted plaintext of encrypted reads. The decode builds the Python objects directly. It shares most repeated short strings, map keys up to 64 bytes and values up to 32, as one object within a value (a small table per decode; two strings that land in one slot replace each other). It runs the Python hook that restores datetimes only on maps holding a `__datetime__`, `__date__` or `__time__` key, because the hook hands back any other map unchanged. It returns the same values and refuses the same inputs with the same exception classes as the decode it replaced: cachekit-core's structure walk, then msgpack-python's C unpacker calling the hook once per map. On the 100-user dict (23,547 bytes of MessagePack, about 3,400 objects), the decode fell from 2,577,967 to 1,220,541 instructions on CPython 3.14 (−52.7%, callgrind main-thread instructions per call, the [Instruction Budgets](#instruction-budgets) method), level with `pickle.loads` of the same value (1,199,580). The 1,000-user dict (239 KB) fell 51.5% and a 1,090-byte one 45.8%. A list of 200 records each holding a datetime fell 33.9%, since every datetime still runs the hook. On the abi3t extension that CPython 3.15 and later load, the 100-user dict fell 51.0%. +Every read with the default serializer decodes its MessagePack in the Rust extension, in one pass that also enforces the bounds on untrusted MessagePack: nesting at most 1,024 levels deep, and no header declaring more elements or bytes than the rest of the input can back. That covers L2 reads, L1 hits and the decrypted plaintext of encrypted reads. The decode builds the Python objects directly. It shares most repeated short strings, map keys up to 64 bytes and values up to 32, as one object within a value (a small table per decode; two strings that land in one slot replace each other). It runs the Python hook that restores datetimes only on maps holding a `__datetime__`, `__date__` or `__time__` key, because the hook hands back any other map unchanged. It returns the same values and refuses the same inputs with the same exception classes as the decode it replaced: cachekit-core's structure walk, then msgpack-python's C unpacker calling the hook once per map. Running out of memory is the one difference. A large string, bin, ext or list still raises `MemoryError`. If CPython cannot allocate a dict, a list of up to 64 items, a float or an int outside -5 to 256, the process aborts. With Linux's default overcommit, the OOM killer usually ends the process before either outcome. On the 100-user dict (23,547 bytes of MessagePack, about 3,400 objects), the decode fell from 2,577,967 to 1,220,541 instructions on CPython 3.14 (−52.7%, callgrind main-thread instructions per call, the [Instruction Budgets](#instruction-budgets) method), level with `pickle.loads` of the same value (1,199,580). The 1,000-user dict (239 KB) fell 51.5% and a 1,090-byte one 45.8%. A list of 200 records each holding a datetime fell 33.9%, since every datetime still runs the hook. On the abi3t extension that CPython 3.15 and later load, the 100-user dict fell 51.0%. For a plaintext cache with the default serializer and integrity checking on, an entry's first L1 hit decodes it as every read does, and then, when its verified MessagePack is no longer than the envelope, which is true of small values and incompressible ones, L1 holds that MessagePack in the envelope's place. Later hits on it skip the envelope's LZ4 decompress and xxHash3-64 checksum. Every hit still runs the bounded decode, which costs no more than the unbounded one did. The 100-user dict compresses to a 3,171-byte envelope from 23.5KB of MessagePack, so L1 keeps its envelope; no budget gates this large-value hit, and `l1_hit_backed` in [Instruction Budgets](#instruction-budgets) gates the small-value hit on held MessagePack. Encrypted caches always keep the ciphertext envelope, and every hit on one decrypts and decodes. diff --git a/rust/src/msgpack_decode.rs b/rust/src/msgpack_decode.rs index bd3a96d9..4b77e6a2 100644 --- a/rust/src/msgpack_decode.rs +++ b/rust/src/msgpack_decode.rs @@ -15,11 +15,14 @@ //! //! Safe Rust only (`SECURITY.md`): every object is a PyO3 `Bound`, so each reference is released //! by its owner on any exit. PyO3's plain constructors abort the process when CPython is out of -//! memory (the release profile sets `panic = "abort"`), so every allocation the input sizes, a -//! long array's list and the bytes of a bin, an ext or `ExtraData`, goes through a constructor -//! that raises MemoryError instead (`nones`, `bytes_of`). - -// cargo-geiger checks the crate only after a merge; this fails any build, a PR's included. +//! memory (the release profile sets `panic = "abort"`). A single allocation the input sizes, a +//! long array's list or the bytes of a str, bin, ext or `ExtraData`, goes through a constructor +//! that raises MemoryError instead (`nones`, `bytes_of`, `PyString::from_bytes`). A dict, a list +//! of at most `SHORT_LIST` items, a float or an int outside -5..=256 does not: running out of +//! memory while building one aborts the process, where msgpack-python raises MemoryError. + +// cargo-geiger checks the crate only after a merge; this fails any build with the `python` +// feature (a default one), which every PR's clippy job and wheel builds compile. #![forbid(unsafe_code)] use std::convert::Infallible; @@ -204,9 +207,9 @@ impl<'a, 'py> Decoder<'a, 'py> { list.set_item(*index, value).map_err(|e| restore(py, e))?; *index += 1; } - Open::Map(dict, pending) => match pending.take() { + Open::Map(dict, awaiting_key) => match awaiting_key.take() { None => { - *pending = Some(value); + *awaiting_key = Some(value); break; } Some(key) => { diff --git a/tests/unit/test_standard_decode.py b/tests/unit/test_standard_decode.py index f6922f7e..8cd3ec8d 100644 --- a/tests/unit/test_standard_decode.py +++ b/tests/unit/test_standard_decode.py @@ -50,10 +50,18 @@ def reference(data: bytes) -> Any: return unpackb_bounded(data, raw=False, use_list=True, object_hook=_standard_object_hook) -def rust(data: bytes) -> Any: +def rust(data: bytes | bytearray | memoryview) -> Any: return unpackb_standard(data, MSGPACK_MAX_NESTING, _standard_object_hook) +def _decode_or_refuse(data: bytes | bytearray | memoryview) -> None: + """Decode ``data``, swallowing only the refusals CASES pins; any other error fails the caller.""" + try: + rust(data) + except (ValueError, SerializationError, TypeError): # noqa: S110 - the refusal is the expected outcome + pass + + def outcome(decode: Callable[[bytes], Any], data: bytes) -> tuple[Any, ...]: """What ``decode`` returned, as a repr, or the class it raised. A bound's message is kept: it names the bound.""" try: @@ -320,8 +328,10 @@ def test_an_amplifying_payload_costs_no_more_heap_than_before() -> None: sys.exit(0) sys.exit(3) """ -# Every allocation the input sizes (n is 128 MiB), as the `data` expression the script decodes. +# Each single allocation the input sizes (n is 128 MiB), as the `data` expression the script decodes. Many small +# objects (dicts, short lists, floats, ints past the small-int cache) still abort: see msgpack_decode.rs's module doc. _OUT_OF_MEMORY_SHAPES = { + "str": 'b"\\xdb" + n.to_bytes(4, "big") + b"a" * n', "bin": 'b"\\xc6" + n.to_bytes(4, "big") + bytes(n)', "ext": 'b"\\xc9" + n.to_bytes(4, "big") + b"\\x05" + bytes(n)', "trailing bytes": 'b"\\xc0" + bytes(n)', @@ -347,10 +357,7 @@ def _measure_leaks() -> list[int]: def every_path() -> None: for data in (*CASES.values(), _pack({"k": ["v" * 40, datetime(2024, 1, 15), b"b"] * 3}, True)): - try: - rust(data) - except (ValueError, SerializationError, TypeError): # noqa: S110 - every refusal CASES pins; any other error fails - pass + _decode_or_refuse(data) for _ in range(50): # warm caches and lazily imported msgpack classes every_path() @@ -384,7 +391,7 @@ def test_no_path_leaks_an_object() -> None: assert min(rounds) < LEAK_ROUND_BUDGET, sizes -# Empty and one-byte bytes are immortal singletons, whose count never moves. +# Empty and one-byte bytes are shared cached singletons (immortal from 3.12), so their counts are not the decode's alone. _MULTI_BYTE = {name: data for name, data in CASES.items() if len(data) > 1} @@ -398,10 +405,7 @@ def test_no_decode_keeps_a_reference_to_its_input(data: bytes) -> None: before = [sys.getrefcount(o) for o in watched] for source in (owned, view, buffer): for _ in range(100): - try: - rust(source) - except (ValueError, SerializationError, TypeError): # noqa: S110 - every refusal CASES pins; any other error fails - pass + _decode_or_refuse(source) del source # still bound to the last input assert [sys.getrefcount(o) for o in watched] == before buffer.append(0) # a buffer export never released refuses the resize with BufferError From 3d045a2da6b1e42ac93d27b84e22fb06fc660bea Mon Sep 17 00:00:00 2001 From: Ray Walker Date: Sun, 1 Nov 2026 12:25:58 +0000 Subject: [PATCH 9/9] test: run the out-of-memory subprocess on the parent's import path --- tests/unit/test_standard_decode.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/unit/test_standard_decode.py b/tests/unit/test_standard_decode.py index 8cd3ec8d..1b266ef5 100644 --- a/tests/unit/test_standard_decode.py +++ b/tests/unit/test_standard_decode.py @@ -22,6 +22,7 @@ from __future__ import annotations import gc +import os import subprocess import sys import sysconfig @@ -348,7 +349,8 @@ def test_an_amplifying_payload_costs_no_more_heap_than_before() -> None: def test_an_allocation_the_input_sizes_raises_memory_error(shape: str) -> None: """With panic = "abort", a PyO3 constructor that meets a NULL would kill the process, not raise.""" script = _OUT_OF_MEMORY.format(shape=shape) - run = subprocess.run([sys.executable, "-c", script], capture_output=True, check=False) # noqa: S603 - this module's own script + env = {**os.environ, "PYTHONPATH": os.pathsep.join(sys.path)} # import this checkout's cachekit, as pytest does + run = subprocess.run([sys.executable, "-c", script], capture_output=True, check=False, env=env) # noqa: S603 - this module's own script assert run.returncode == 0, run.stderr.decode()[-2000:]