fix(serializers): type ByteStorage.retrieve failures and collapse duplicated columnar decode (LAB-2736)

rust/src/lib.rs

@@ -39,6 +39,13 @@ fn _rust_serializer(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
     // Add byte storage class
     m.add_class::<python_bindings::PyByteStorage>()?;
 
+    // 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::<python_bindings::EnvelopeIntegrityError>();
+    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)?)?;

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

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 cachekit.hash_utils import redact_error_for_log
 
 from .base import PAYLOAD_DECODE_ERRORS, SerializationError, SerializationFormat, SerializationMetadata, unpackb_bounded
@@ -575,6 +575,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)
@@ -602,45 +610,31 @@ 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
         if self.enable_integrity_checking:
             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: {redact_error_for_log(e)}"
@@ -657,10 +651,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(
@@ -686,6 +677,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)

tests/unit/test_auto_serializer_mutation_and_corruption.py

@@ -77,6 +77,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:
@@ -140,6 +153,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).

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,10 +729,8 @@ def test_nullable_dtypes_and_objects_roundtrip(self):
             }
         )
 
-        data = ser._serialize_dataframe(df)  # previously raised: msgpack can't pack pd.NA
-        # _decode_columnar decodes the msgpack body then hands the document to
-        # _deserialize_dataframe (which now takes a decoded document, not bytes).
-        out = ser._decode_columnar(data, "dataframe")
+        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)
@@ -746,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._decode_columnar(data, "dataframe")
+        data, metadata = ser.serialize(df)
+        out = ser.deserialize(data, metadata)
 
         assert out["x"].tolist() == [1, 2, 3]
 
@@ -760,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._decode_columnar(data, "series")
+        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
@@ -773,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._decode_columnar(ser._serialize_series(s), "series")
+        data, metadata = ser.serialize(s)
+        out = ser.deserialize(data, metadata)
 
         assert out.tolist() == [1, 2, 3]