Fix NPZ buffer lifetime to prevent use-after-free

src/spdl/io/_array.py

@@ -10,7 +10,7 @@
     "NpzFile",
 ]
 from collections.abc import Iterator, Mapping
-from typing import TypeAlias
+from typing import Any, TypeAlias
 
 import numpy as np
 from numpy.typing import NDArray
@@ -28,6 +28,26 @@ def _get_pointer(data: Buffer) -> int:
     return np.frombuffer(data, dtype=np.byte).ctypes.data
 
 
+class _OwnedArrayInterface:
+    """Expose the array interface of ``obj`` while keeping ``owner`` alive.
+
+    The ``NPYArray`` object returned by the C++ extension points at a memory
+    region it does not own. NumPy sets the ``base`` of an array to the object it
+    was created from, so creating the array from this wrapper (instead of from
+    the ``NPYArray`` directly) keeps the object owning the memory alive for as
+    long as the array is.
+    """
+
+    def __init__(self, obj: object, owner: object) -> None:
+        self._obj = obj
+        self._owner = owner
+
+    @property
+    def __array_interface__(self) -> dict[str, Any]:
+        # pyre-ignore[16]
+        return self._obj.__array_interface__
+
+
 def load_npy(data: Buffer, *, copy: bool = False) -> NDArray:
     """Load NumPy NDArray from bytes.
 
@@ -97,7 +117,11 @@ def __init__(
         data: "bytes | memoryview[bytes]",
         meta: dict[str, tuple[int, int, int, int]],
     ) -> None:
-        self._data: int = _get_pointer(data)
+        # `_data` is a raw pointer into `data`, so the archive must be kept
+        # alive. A memoryview also blocks resizing a mutable source, which
+        # would reallocate the buffer and leave `_data` dangling.
+        self._buf: "memoryview[bytes]" = memoryview(data)
+        self._data: int = _get_pointer(self._buf)
         self._len: int = len(data)
         self._meta = meta
         self.files: list[str] = [f.removesuffix(".npy") for f in meta]
@@ -125,11 +149,16 @@ def __getitem__(self, key: str) -> NDArray:
         offset, compressed_size, uncompressed_size, compression_method = self._meta[key]
         match compression_method:
             case 0:
+                # The data is stored uncompressed, so the resulting array refers
+                # to the archive itself. It must keep the archive alive, as it
+                # can outlive this `NpzFile` object.
                 buffer = _libspdl._archive.load_npy(
                     self._data, size=compressed_size, offset=offset
                 )
-                return np.array(buffer, copy=False)
+                return np.array(_OwnedArrayInterface(buffer, self._buf), copy=False)
             case 8:
+                # The data is inflated into a buffer owned by the `NPYArray`
+                # object, which NumPy keeps alive as the base of the array.
                 buffer = _libspdl._archive.load_npy_compressed(
                     self._data, offset, compressed_size, uncompressed_size
                 )

tests/io/array_test.py

@@ -6,8 +6,11 @@
 
 # pyre-strict
 
+import gc
 import io
+import sys
 import unittest
+from collections.abc import Callable
 from io import BytesIO
 
 import numpy as np
@@ -256,3 +259,110 @@ def test_load_npy_cpp(self) -> None:
             buffer = spdl.io.load_npy(data)
             hyp = np.array(buffer, copy=False)
             np.testing.assert_array_equal(hyp, ref)
+
+
+def _reuse_freed_memory(size: int, count: int = 2000) -> list[bytearray]:
+    """Fill recently freed memory with a recognizable pattern.
+
+    CPython hands memory from deallocated objects back out to later
+    allocations of a similar size, so `count` buffers of `size` bytes are
+    likely to land on the block the archive just freed. A stale pointer into
+    it then reads `0xAB` and the caller's `assert_array_equal` fails; without
+    this, it would likely read the original bytes and pass.
+
+    Best-effort by nature -- the refcount assertions below are the
+    deterministic check.
+    """
+    return [bytearray(b"\xab" * size) for _ in range(count)]
+
+
+class TestNpzBufferLifetime(unittest.TestCase):
+    """`NpzFile` does not copy the archive.
+
+    It holds a raw pointer into the source buffer, and the arrays it returns for
+    stored (uncompressed) entries are views into the same memory. Both read
+    freed memory unless the source buffer is kept alive.
+
+    The tests come in two flavors. The ones asserting on `sys.getrefcount`
+    check the contract directly, and are the authoritative check. The ones
+    calling `_reuse_freed_memory` additionally try to turn a violation into an
+    observable data corruption; see that function for the caveats.
+    """
+
+    def test_load_npz_retains_source(self) -> None:
+        """`load_npz` keeps a reference to the source buffer."""
+        ref = np.arange(10)
+        data = _dump_npz(x=ref)
+
+        num_refs = sys.getrefcount(data)
+        npz = spdl.io.load_npz(data)
+
+        self.assertGreater(
+            sys.getrefcount(data),
+            num_refs,
+            "`NpzFile` must keep a reference to the source buffer, "
+            "as it holds a pointer into it.",
+        )
+        np.testing.assert_array_equal(npz["x"], ref)
+
+    def test_getitem_retains_source(self) -> None:
+        """Arrays of stored entries keep the source buffer alive.
+
+        Such an array is a view into the archive, so it can outlive the
+        `NpzFile` it was retrieved from.
+        """
+        ref = np.arange(10)
+        data = _dump_npz(x=ref)
+
+        num_refs = sys.getrefcount(data)
+        # The `NpzFile` is released as soon as the entry is retrieved.
+        arr = spdl.io.load_npz(data)["x"]
+        gc.collect()
+
+        self.assertGreater(
+            sys.getrefcount(data),
+            num_refs,
+            "The array must keep a reference to the source buffer, "
+            "as it is a view into it.",
+        )
+        np.testing.assert_array_equal(arr, ref)
+
+    @parameterized.expand(
+        [
+            ("stored", _dump_npz),
+            ("deflated", _dump_npz_compressed),
+        ]
+    )
+    def test_load_npz_source_may_be_temporary(
+        self, _: str, dump: Callable[..., bytes]
+    ) -> None:
+        """Entries are readable when the caller does not hold the source."""
+        ref = np.arange(1000, dtype=np.int64)
+        size = len(dump(x=ref))
+
+        # The source is a temporary, so it is released when `load_npz` returns
+        # unless `NpzFile` retains it.
+        npz = spdl.io.load_npz(dump(x=ref))
+        gc.collect()
+        # If `NpzFile` failed to retain the temporary, `npz["x"]` now points
+        # into freed memory, and reads `0xAB` instead of `ref`.
+        clobber = _reuse_freed_memory(size)
+
+        np.testing.assert_array_equal(npz["x"], ref)
+
+        del clobber
+
+    def test_array_outlives_npz_file(self) -> None:
+        """A stored entry stays valid after the `NpzFile` is released."""
+        ref = np.arange(1000, dtype=np.int64)
+        size = len(_dump_npz(x=ref))
+
+        arr = spdl.io.load_npz(_dump_npz(x=ref))["x"]
+        gc.collect()
+        # If the source buffer was not retained, `arr` now views freed memory,
+        # and reads `0xAB` instead of `ref`.
+        clobber = _reuse_freed_memory(size)
+
+        np.testing.assert_array_equal(arr, ref)
+
+        del clobber