From f79b0d33273cdad94358a12abb58bf6ba9a110e0 Mon Sep 17 00:00:00 2001 From: Mark S Date: Sun, 13 Sep 2026 16:36:48 +1000 Subject: [PATCH] fix(serializers): type ByteStorage.retrieve failures and collapse duplicated columnar decode (LAB-2736) ByteStorage.retrieve() flattened every corruption case (checksum mismatch, decompression bomb/failure, size mismatch) into the same PyValueError as "not a ByteStorage envelope at all", so AutoSerializer.deserialize() could not tell a genuinely corrupted cache entry from bytes that were legitimately written without an envelope (integrity checking off) - a checksum mismatch either fell through to a confusing "not decodable" error or, on the DataFrame/Series path, was duplicated into two near-identical retrieve+decode blocks with their own ad hoc error text. Add EnvelopeIntegrityError (rust/src/python_bindings.rs), a ValueError subclass raised for every ByteStorageError variant except DeserializationFailed (which stays a plain ValueError - the fall-through signal deserialize() depends on). AutoSerializer.deserialize() catches it specifically and re-raises as SerializationError without falling through. Collapse the DataFrame/Series metadata pre-branch and the verified-envelope branch onto the single _decode_columnar path, and drop the dead isinstance(data, dict) branches in _deserialize_dataframe/_deserialize_series now that no caller passes raw bytes. Expert-panel review caught a regression the collapse introduced: an entry written with integrity off and read by an integrity-on reader (same metadata routing to dataframe/series) fell through to the generic msgpack fallback and returned the raw wire dict instead of failing closed or reconstructing - fixed by routing that fallback through _decode_columnar too, with a regression test. --- rust/src/lib.rs | 7 ++ rust/src/python_bindings.rs | 39 +++++++- src/cachekit/serializers/auto_serializer.py | 98 ++++++++----------- ...auto_serializer_mutation_and_corruption.py | 43 ++++++++ tests/unit/test_auto_serializer_new_types.py | 17 ++-- 5 files changed, 141 insertions(+), 63 deletions(-) diff --git a/rust/src/lib.rs b/rust/src/lib.rs index fca8d5b7..dd57b441 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -39,6 +39,13 @@ fn _rust_serializer(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { // Add byte storage class m.add_class::()?; + // ByteStorage.retrieve() envelope-verification-failure taxonomy (LAB-2736). Same + // __module__ patch as KeyringConfigurationError below: create_exception! sets it to the + // bare "_rust_serializer", which breaks pickling back to a parent process otherwise. + let envelope_integrity_error = m.py().get_type::(); + envelope_integrity_error.setattr("__module__", "cachekit._rust_serializer")?; + m.add("EnvelopeIntegrityError", envelope_integrity_error)?; + // Standalone integrity primitive — registered unconditionally (usable with // the checksum feature alone; must not vanish when encryption is off) m.add_function(wrap_pyfunction!(python_bindings::checksum_py, m)?)?; diff --git a/rust/src/python_bindings.rs b/rust/src/python_bindings.rs index d15c96e7..e4d889cc 100644 --- a/rust/src/python_bindings.rs +++ b/rust/src/python_bindings.rs @@ -5,12 +5,49 @@ //! SDK-owned msgpack decode bound in `crate::msgpack_bounds`. use crate::msgpack_bounds::check_msgpack_structure; +use cachekit_core::byte_storage::ByteStorageError; use cachekit_core::ByteStorage; use pyo3::buffer::PyBuffer; use pyo3::exceptions::PyValueError; use pyo3::prelude::*; use pyo3::types::PyBytes; +pyo3::create_exception!( + _rust_serializer, + EnvelopeIntegrityError, + PyValueError, + "A ByteStorage envelope parsed but failed verification: checksum mismatch, decompression\n\ + bomb/failure, or a decoded size mismatch against the envelope header.\n\ + \n\ + Distinguishes a verified-but-corrupt envelope (this exception) from bytes that were never\n\ + a ByteStorage envelope at all (`DeserializationFailed`, e.g. written with integrity\n\ + checking off) — the latter stays a plain `ValueError` so callers keep falling through to\n\ + the plain-msgpack/NumPy decode paths for it, while this one must fail closed.\n\ + \n\ + Subclasses ValueError so existing `pytest.raises(ValueError)` assertions on `retrieve()`\n\ + failures stay valid. `AutoSerializer.deserialize` catches this specifically and re-raises\n\ + it as `SerializationError` without falling through." +); + +/// Map a cachekit-core `retrieve()` failure onto the Python exception taxonomy. +/// +/// `DeserializationFailed` means `envelope_bytes` never parsed as a `StorageEnvelope` — not +/// corruption, just "not an envelope" — so it stays a plain `ValueError`, the fall-through +/// signal `AutoSerializer.deserialize` depends on. Every other variant is mapped to +/// `EnvelopeIntegrityError` and must fail closed: the post-parse checks (checksum, decompressed +/// size, decompression itself, the compression-ratio bomb guard) are genuine corruption or +/// tampering, and `InputTooLarge` — raised on the raw `envelope_bytes` length before parsing is +/// even attempted — is deliberately bucketed the same way rather than treated as "not an +/// envelope": falling through would hand an oversized blob to the plain-msgpack decode path +/// instead of rejecting it outright, trading one size guard for a weaker one. +fn retrieve_error_to_py(err: ByteStorageError) -> PyErr { + let message = format!("Retrieval failed: {}", err); + match err { + ByteStorageError::DeserializationFailed(_) => PyValueError::new_err(message), + _ => EnvelopeIntegrityError::new_err(message), + } +} + /// Python wrapper for ByteStorage #[pyclass(name = "ByteStorage")] pub struct PyByteStorage { @@ -155,7 +192,7 @@ impl PyByteStorage { let data = view.as_slice(); // Detach from the GIL for decompression + checksum (see store()). py.detach(|| self.inner.retrieve(data)) - .map_err(|e| PyValueError::new_err(format!("Retrieval failed: {}", e))) + .map_err(retrieve_error_to_py) } /// Get compression ratio for given data diff --git a/src/cachekit/serializers/auto_serializer.py b/src/cachekit/serializers/auto_serializer.py index f75f7bb0..124d2df2 100644 --- a/src/cachekit/serializers/auto_serializer.py +++ b/src/cachekit/serializers/auto_serializer.py @@ -58,7 +58,7 @@ HAS_ARROW_SERIALIZER = False ArrowSerializer = None # type: ignore[assignment,misc] -from cachekit._rust_serializer import ByteStorage +from cachekit._rust_serializer import ByteStorage, EnvelopeIntegrityError from .base import PAYLOAD_DECODE_ERRORS, SerializationError, SerializationFormat, SerializationMetadata, unpackb_bounded @@ -557,6 +557,14 @@ def deserialize(self, data: bytes | memoryview, metadata: Optional[Serialization Returns: Any: Deserialized Python object + + Raises: + SerializationError: A ByteStorage envelope was present but failed verification + (checksum mismatch, decompression bomb/failure, size mismatch) — genuine + corruption or tampering, or a payload that failed to decode inside a + verified envelope. Bytes that were never a ByteStorage envelope (e.g. written + with integrity checking off) are not an error here: they fall through to the + plain-msgpack/NumPy decode paths and only raise if none of those decode either. """ # coerce unwrap's zero-copy memoryview; no-op when already bytes (enables .startswith below + Rust retrieve) data = bytes(data) @@ -584,30 +592,13 @@ def deserialize(self, data: bytes | memoryview, metadata: Optional[Serialization "Cannot deserialize Arrow format: ArrowSerializer not available. " "Install with: pip install 'cachekit[data]'" ) - elif detected_format == "dataframe": - if self.enable_integrity_checking and len(data) > 4: - # Unwrap the ByteStorage envelope. A checksum mismatch (raised by retrieve) fails - # closed with a clear corruption error instead of being swallowed and re-parsed as - # raw msgpack, which lost the diagnostic and produced a confusing error (#156). - # The unpack/build sits OUTSIDE this guard so a genuine post-retrieve error surfaces - # as itself rather than being mistaken for corruption. - try: - original_data, _ = self._byte_storage.retrieve(data) - except (ValueError, SerializationError) as e: - raise SerializationError(f"DataFrame integrity check failed (corrupted cache entry): {e}") from e - return self._decode_columnar(original_data, detected_format) - # Integrity off: data is direct msgpack (no envelope) - return self._decode_columnar(data, detected_format) - elif detected_format == "series": - if self.enable_integrity_checking and len(data) > 4: - # Same fail-closed contract as the DataFrame branch above (#156). - try: - original_data, _ = self._byte_storage.retrieve(data) - except (ValueError, SerializationError) as e: - raise SerializationError(f"Series integrity check failed (corrupted cache entry): {e}") from e - return self._decode_columnar(original_data, detected_format) - # Integrity off: data is direct msgpack (no envelope) - return self._decode_columnar(data, detected_format) + elif detected_format in ("dataframe", "series"): + if not (self.enable_integrity_checking and len(data) > 4): + # Integrity off: data is direct msgpack (no envelope) + return self._decode_columnar(data, detected_format) + # Integrity on: fall through to the shared Rust-envelope retrieve below, which + # re-derives this same detected_format from metadata (#156, LAB-2736 collapse — + # one retrieve+decode path instead of a second copy here). # For Rust-envelope formats, use the Rust layer envelope_error: Exception | None = None @@ -615,14 +606,17 @@ def deserialize(self, data: bytes | memoryview, metadata: Optional[Serialization try: # Use Rust layer for decompression and validation original_data, format_id = self._byte_storage.retrieve(data) - except SerializationError: - # Re-raise SerializationError (corruption detection) without swallowing - raise + except EnvelopeIntegrityError as e: + # The envelope parsed but failed verification (checksum, decompression bomb, + # size mismatch) — genuine corruption or tampering. Must fail closed, never + # fall through to a re-parse as plain msgpack/NumPy (that would either raise a + # confusing "not decodable" error or, worse, decode envelope bytes as if they + # were the payload). + raise SerializationError(f"Cache entry failed envelope verification (corrupted cache entry): {e}") from e except Exception as e: - # Not a ByteStorage envelope (e.g. written with integrity checking off): + # Not a ByteStorage envelope at all (e.g. written with integrity checking off): # fall through to the Python-only paths below, keeping the reason for the - # final error (a checksum mismatch also lands here — retrieve raises a plain - # ValueError for both; distinguishing them is a Rust-extension follow-up). + # final error. envelope_error = e logger.debug(f"Rust envelope parsing failed, falling back to Python-only deserialization: {e}") else: @@ -637,10 +631,7 @@ def deserialize(self, data: bytes | memoryview, metadata: Optional[Serialization if detected_format == "numpy": return self._deserialize_numpy(original_data) if detected_format in ("dataframe", "series"): - unpacked_data = unpackb_bounded(original_data, **self._msgpack_unpack_opts) - if detected_format == "dataframe": - return self._deserialize_dataframe(unpacked_data) - return self._deserialize_series(unpacked_data) + return self._decode_columnar(original_data, detected_format) return unpackb_bounded(original_data, **self._msgpack_unpack_opts) except PAYLOAD_DECODE_ERRORS as e: raise SerializationError( @@ -666,6 +657,17 @@ def deserialize(self, data: bytes | memoryview, metadata: Optional[Serialization "Cannot deserialize Arrow format: ArrowSerializer not available. Install with: pip install 'cachekit[data]'" ) + # Metadata says dataframe/series but the envelope attempt above either wasn't + # taken (integrity off) or failed as "not an envelope" (envelope_error set, + # e.g. cross-config read of an entry written with integrity off): the bytes + # are direct columnar msgpack, not a bare object, and MUST still reconstruct + # through _decode_columnar — falling through to the generic branch below would + # return the raw wire dict instead of a DataFrame/Series (LAB-2736 regression + # caught by expert-panel review: silently wrong-typed data, not merely a + # confusing error). + if metadata and hasattr(metadata, "original_type") and metadata.original_type in ("dataframe", "series"): + return self._decode_columnar(data, metadata.original_type) + # Python-only path (no Rust compression) - direct msgpack deserialization try: return unpackb_bounded(data, **self._msgpack_unpack_opts) @@ -821,8 +823,8 @@ def _serialize_dataframe(self, df: pd.DataFrame) -> bytes: else: return msgpack_data # type: ignore[return-value] - def _deserialize_dataframe(self, data) -> pd.DataFrame: - """Deserialize DataFrame from column-wise data. + def _deserialize_dataframe(self, data: dict) -> pd.DataFrame: + """Deserialize DataFrame from an already-unpacked column-wise document. Requires: pandas installed (HAS_PANDAS=True) @@ -833,14 +835,7 @@ def _deserialize_dataframe(self, data) -> pd.DataFrame: if not HAS_PANDAS: raise RuntimeError("Pandas not installed. Install with: pip install cachekit[data]") - # If data is already unpacked (from Rust layer), use it directly - if isinstance(data, dict): - serialized = data - else: - # Otherwise unpack msgpack - serialized = unpackb_bounded(data, **self._msgpack_unpack_opts) - - serialized = _expect(serialized, dict, "document") + serialized = _expect(data, dict, "document") columns_data = {} for col, col_info in _expect(serialized["data"], dict, "data").items(): what = f"column {col!r:.40}" # col is attacker-chosen: cap the echo @@ -883,8 +878,8 @@ def _serialize_series(self, series: pd.Series) -> bytes: else: return msgpack_data # type: ignore[return-value] - def _deserialize_series(self, data) -> pd.Series: - """Deserialize Pandas Series. + def _deserialize_series(self, data: dict) -> pd.Series: + """Deserialize Pandas Series from an already-unpacked document. Requires: pandas installed (HAS_PANDAS=True) @@ -895,14 +890,7 @@ def _deserialize_series(self, data) -> pd.Series: if not HAS_PANDAS: raise RuntimeError("Pandas not installed. Install with: pip install cachekit[data]") - # If data is already unpacked (from Rust layer), use it directly - if isinstance(data, dict): - serialized = data - else: - # Otherwise unpack msgpack - serialized = unpackb_bounded(data, **self._msgpack_unpack_opts) - - serialized = _expect(serialized, dict, "document") + serialized = _expect(data, dict, "document") series = pd.Series(_column_values(serialized, "series"), name=serialized["name"]) # Restore index if it was serialized diff --git a/tests/unit/test_auto_serializer_mutation_and_corruption.py b/tests/unit/test_auto_serializer_mutation_and_corruption.py index bfb41f30..7be00e75 100644 --- a/tests/unit/test_auto_serializer_mutation_and_corruption.py +++ b/tests/unit/test_auto_serializer_mutation_and_corruption.py @@ -73,6 +73,19 @@ def test_roundtrip_without_metadata_via_envelope_format_id(self, value: pd.DataF data, _ = s.serialize(value) _assert_equal(s.deserialize(data), value) + @pytest.mark.parametrize("value", [FRAME, SERIES], ids=["dataframe", "series"]) + def test_roundtrip_cross_config_written_off_read_on(self, value: pd.DataFrame | pd.Series) -> None: + """LAB-2736 regression: an entry written with integrity off (no ByteStorage envelope) + must still reconstruct through the columnar decoder when read by a reader with + integrity on — not fall through to returning the raw wire dict unchecked.""" + writer = _no_arrow(enable_integrity_checking=False) + reader = _no_arrow(enable_integrity_checking=True) + data, meta = writer.serialize(value) + + out = reader.deserialize(data, meta) + assert type(out) is type(value) + _assert_equal(out, value) + @pytest.mark.unit class TestDeserializedArraysAreWritable: @@ -136,6 +149,36 @@ def test_dataframe_corruption_raises_serialization_error(self) -> None: s.deserialize(bytes(corrupted), meta) +@pytest.mark.unit +class TestEnvelopeVerificationVsNotAnEnvelope: + """LAB-2736: ``retrieve()`` raises a distinct type for a verified-but-corrupt envelope + (checksum/decompression/size failure) vs. bytes that were never a ByteStorage envelope + at all (e.g. written with integrity checking off). ``deserialize`` must fail closed on + the former and keep falling through to the plain-msgpack path only on the latter. + """ + + def test_corrupted_payload_names_the_integrity_failure(self) -> None: + s = AutoSerializer() + data, meta = s.serialize({"nums": list(range(2000))}) + + corrupted = bytearray(data) + corrupted[len(corrupted) // 2] ^= 0xFF + with pytest.raises(SerializationError) as exc_info: + s.deserialize(bytes(corrupted), meta) + + message = str(exc_info.value) + assert "envelope verification" in message + assert "not a decodable MessagePack" not in message + + def test_plain_msgpack_written_with_integrity_off_still_falls_through(self) -> None: + off = AutoSerializer(enable_integrity_checking=False) + payload = {"a": 1, "b": [1, 2, 3]} + data, _ = off.serialize(payload) + + on = AutoSerializer(enable_integrity_checking=True) + assert on.deserialize(data) == payload + + # A well-formed __ndarray__ marker: the object hook turns it into an ndarray wherever it sits, so a # forged document can put an array where the writer only ever puts a list or a dict. M8[2s] is a # dtype numpy accepts and pandas then asserts on (AssertionError, outside PAYLOAD_DECODE_ERRORS). diff --git a/tests/unit/test_auto_serializer_new_types.py b/tests/unit/test_auto_serializer_new_types.py index 7009f920..c8c4bc92 100644 --- a/tests/unit/test_auto_serializer_new_types.py +++ b/tests/unit/test_auto_serializer_new_types.py @@ -719,6 +719,7 @@ class TestColumnarFallbackExtensionDtypes: def test_nullable_dtypes_and_objects_roundtrip(self): pd = pytest.importorskip("pandas") ser = AutoSerializer(enable_integrity_checking=False) + ser._arrow_serializer = None # force the msgpack-columnar DataFrame fallback path df = pd.DataFrame( { "ints": pd.array([1, 2, None, 4], dtype="Int64"), # capital -> missed numeric branch before @@ -728,8 +729,8 @@ def test_nullable_dtypes_and_objects_roundtrip(self): } ) - data = ser._serialize_dataframe(df) # previously raised: msgpack can't pack pd.NA - out = ser._deserialize_dataframe(data) + data, metadata = ser.serialize(df) # previously raised: msgpack can't pack pd.NA + out = ser.deserialize(data, metadata) assert list(out.columns) == ["ints", "floats", "objs", "plain"] assert out.shape == (4, 4) @@ -744,12 +745,13 @@ def test_pyarrow_backed_dtype_does_not_crash(self): pd = pytest.importorskip("pandas") pytest.importorskip("pyarrow") ser = AutoSerializer(enable_integrity_checking=False) + ser._arrow_serializer = None # force the msgpack-columnar DataFrame fallback path # "int64[pyarrow]".startswith("int") was True -> hit the numeric branch -> # .values.tobytes() AttributeError on the Arrow extension array, before the fix. df = pd.DataFrame({"x": pd.array([1, 2, 3], dtype="int64[pyarrow]")}) - data = ser._serialize_dataframe(df) - out = ser._deserialize_dataframe(data) + data, metadata = ser.serialize(df) + out = ser.deserialize(data, metadata) assert out["x"].tolist() == [1, 2, 3] @@ -758,8 +760,8 @@ def test_nullable_series_roundtrip(self): ser = AutoSerializer(enable_integrity_checking=False) s = pd.Series(pd.array([1, 2, None, 4], dtype="Int64"), name="n") - data = ser._serialize_series(s) # previously raised on the pd.NA sentinel - out = ser._deserialize_series(data) + data, metadata = ser.serialize(s) # previously raised on the pd.NA sentinel + out = ser.deserialize(data, metadata) assert out.name == "n" assert out.iloc[0] == 1 and out.iloc[3] == 4 @@ -771,6 +773,7 @@ def test_pyarrow_backed_series_does_not_crash(self): ser = AutoSerializer(enable_integrity_checking=False) s = pd.Series(pd.array([1, 2, 3], dtype="int64[pyarrow]"), name="x") - out = ser._deserialize_series(ser._serialize_series(s)) + data, metadata = ser.serialize(s) + out = ser.deserialize(data, metadata) assert out.tolist() == [1, 2, 3]