From 50c71aa1cba89c8ac4a0b87cd409ef990abd4626 Mon Sep 17 00:00:00 2001 From: Gagan Dhakrey Date: Mon, 3 Aug 2026 05:12:50 +0530 Subject: [PATCH 1/4] Fix NPZ buffer lifetime to prevent use-after-free Signed-off-by: Gagan Dhakrey --- src/spdl/io/_array.py | 35 ++++++++++++++++++++++++++++++++--- 1 file changed, 32 insertions(+), 3 deletions(-) diff --git a/src/spdl/io/_array.py b/src/spdl/io/_array.py index dadd9b8ac..5482bf21b 100644 --- a/src/spdl/io/_array.py +++ b/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 ) From 24878107af54fc071f5afb162dff70395ff665ab Mon Sep 17 00:00:00 2001 From: Gagan Dhakrey Date: Wed, 5 Aug 2026 18:30:59 +0530 Subject: [PATCH 2/4] adding testcase Signed-off-by: Gagan Dhakrey --- tests/io/array_test.py | 95 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 95 insertions(+) diff --git a/tests/io/array_test.py b/tests/io/array_test.py index 701ce5037..aadf0467f 100644 --- a/tests/io/array_test.py +++ b/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,95 @@ 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]: + """Allocate over recently freed memory. + + If the archive was released, a stale pointer into it reads this pattern + instead of the original data. + """ + 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. + """ + + 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() + 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() + clobber = _reuse_freed_memory(size) + + np.testing.assert_array_equal(arr, ref) + + del clobber From b09a8d674f40ce43330ca6f6864c4165099aec3a Mon Sep 17 00:00:00 2001 From: Gagan Dhakrey Date: Thu, 13 Aug 2026 04:28:26 +0530 Subject: [PATCH 3/4] adding comments for better understanding Signed-off-by: Gagan Dhakrey --- tests/io/array_test.py | 29 +++++++++++++++++++++++++---- 1 file changed, 25 insertions(+), 4 deletions(-) diff --git a/tests/io/array_test.py b/tests/io/array_test.py index aadf0467f..124d28608 100644 --- a/tests/io/array_test.py +++ b/tests/io/array_test.py @@ -262,10 +262,22 @@ def test_load_npy_cpp(self) -> None: def _reuse_freed_memory(size: int, count: int = 2000) -> list[bytearray]: - """Allocate over recently freed memory. - - If the archive was released, a stale pointer into it reads this pattern - instead of the original data. + """Fill recently freed memory with a recognizable pattern. + + This relies on an implementation detail of CPython: memory released by a + deallocated object is returned to the allocator, which hands it back out to + later allocations of a similar size. So allocating many `size`-byte buffers + right after the archive was freed is likely to land one of them on the + block the archive used to occupy. `count` repetitions raise that chance. + + The buffers are filled with `0xAB` so that a stale pointer into the freed + block reads this pattern instead of the original data, and the subsequent + `assert_array_equal` fails. Without it, a use-after-free would most likely + still read the original bytes and the test would pass. + + This is best-effort: it makes a regression *likely* to be caught, never + guaranteed. The deterministic guarantee comes from the refcount assertions + in `TestNpzBufferLifetime`. """ return [bytearray(b"\xab" * size) for _ in range(count)] @@ -276,6 +288,11 @@ class TestNpzBufferLifetime(unittest.TestCase): 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: @@ -333,6 +350,8 @@ def test_load_npz_source_may_be_temporary( # 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) @@ -346,6 +365,8 @@ def test_array_outlives_npz_file(self) -> None: 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) From cf43f8315668aaa52cc022b25f9de94dc728772d Mon Sep 17 00:00:00 2001 From: Gagan Dhakrey Date: Thu, 13 Aug 2026 04:32:21 +0530 Subject: [PATCH 4/4] shorting comment Signed-off-by: Gagan Dhakrey --- tests/io/array_test.py | 22 ++++++++-------------- 1 file changed, 8 insertions(+), 14 deletions(-) diff --git a/tests/io/array_test.py b/tests/io/array_test.py index 124d28608..de015ada9 100644 --- a/tests/io/array_test.py +++ b/tests/io/array_test.py @@ -264,20 +264,14 @@ def test_load_npy_cpp(self) -> None: def _reuse_freed_memory(size: int, count: int = 2000) -> list[bytearray]: """Fill recently freed memory with a recognizable pattern. - This relies on an implementation detail of CPython: memory released by a - deallocated object is returned to the allocator, which hands it back out to - later allocations of a similar size. So allocating many `size`-byte buffers - right after the archive was freed is likely to land one of them on the - block the archive used to occupy. `count` repetitions raise that chance. - - The buffers are filled with `0xAB` so that a stale pointer into the freed - block reads this pattern instead of the original data, and the subsequent - `assert_array_equal` fails. Without it, a use-after-free would most likely - still read the original bytes and the test would pass. - - This is best-effort: it makes a regression *likely* to be caught, never - guaranteed. The deterministic guarantee comes from the refcount assertions - in `TestNpzBufferLifetime`. + 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)]