diff --git a/capnp/lib/capnp.pyx b/capnp/lib/capnp.pyx index 69dc653f3..fd84debc8 100644 --- a/capnp/lib/capnp.pyx +++ b/capnp/lib/capnp.pyx @@ -16,7 +16,7 @@ from capnp.includes.schema_cpp cimport (MessageReader,) from builtins import memoryview as BuiltinsMemoryview from cpython cimport array, Py_buffer, PyObject_CheckBuffer from cpython.buffer cimport PyBUF_SIMPLE, PyBUF_WRITABLE, PyBUF_WRITE, PyBUF_READ, PyBUF_CONTIG_RO, PyBuffer_FillInfo -from cpython.memoryview cimport PyMemoryView_FromMemory, PyMemoryView_FromBuffer +from cpython.memoryview cimport PyMemoryView_FromMemory, PyMemoryView_FromObject from cpython.bytes cimport PyBytes_FromStringAndSize from cpython.exc cimport PyErr_Clear from cpython.pyport cimport PY_SSIZE_T_MAX @@ -54,6 +54,8 @@ _CAPNP_VERSION_MINOR = capnp.CAPNP_VERSION_MINOR _CAPNP_VERSION_MICRO = capnp.CAPNP_VERSION_MICRO _CAPNP_VERSION = capnp.CAPNP_VERSION +cdef char _EMPTY_DATA_VIEW_SENTINEL = 0 + cdef dict _type_registry = {} @@ -1176,20 +1178,24 @@ cdef class _MessageSize: @cython.internal -cdef class _SegmentView: - cdef object _builder +cdef class _BorrowedBufferView: + """Buffer-protocol exporter that pins an owner while a view borrows its memory.""" + cdef object _owner cdef const char* _ptr cdef Py_ssize_t _size + cdef bint _readonly - cdef _init(self, object builder, const char* ptr, Py_ssize_t size): - self._builder = builder - self._ptr = ptr + cdef _init(self, object owner, const void* ptr, Py_ssize_t size, bint readonly): + self._owner = owner + self._ptr = ptr self._size = size + self._readonly = readonly return self def __getbuffer__(self, Py_buffer *buffer, int flags): - if PyBuffer_FillInfo(buffer, self, self._ptr, self._size, 1, flags) < 0: - raise BufferError("Failed to create segment buffer view") + if PyBuffer_FillInfo(buffer, self, self._ptr, self._size, + self._readonly, flags) < 0: + raise BufferError("Failed to create borrowed buffer view") def __releasebuffer__(self, Py_buffer *buffer): pass @@ -1198,7 +1204,56 @@ cdef class _SegmentView: return self._size def __repr__(self): - return '' % self._size + if self._readonly: + return '' % self._size + return '' % self._size + + +cdef inline object _memoryview_borrowing(object owner, void* ptr, Py_ssize_t size, + bint readonly): + cdef _BorrowedBufferView exporter + exporter = _BorrowedBufferView()._init(owner, ptr, size, readonly) + return PyMemoryView_FromObject(exporter) + + +cdef void _data_field_ptr_reader(_DynamicStructReader self, field, + void** data_ptr, size_t* data_size) except *: + cdef C_DynamicValue.Reader val + cdef capnp.Data.Reader temp_data + + try: + val = self.thisptr.get(field) + except KjException as e: + raise e._to_python() from None + + if val.getType() != capnp.TYPE_DATA: + raise TypeError("Field '{}' is not a DATA field".format(field)) + + temp_data = val.asData() + data_ptr[0] = temp_data.begin() + data_size[0] = temp_data.size() + if data_size[0] == 0 and data_ptr[0] == NULL: + data_ptr[0] = &_EMPTY_DATA_VIEW_SENTINEL + + +cdef void _data_field_ptr_builder(_DynamicStructBuilder self, field, + void** data_ptr, size_t* data_size) except *: + cdef C_DynamicValue.Builder val + cdef capnp.Data.Builder temp_data + + try: + val = self.thisptr.get(field) + except KjException as e: + raise e._to_python() from None + + if val.getType() != capnp.TYPE_DATA: + raise TypeError("Field '{}' is not a DATA field".format(field)) + + temp_data = val.asData() + data_ptr[0] = temp_data.begin() + data_size[0] = temp_data.size() + if data_size[0] == 0 and data_ptr[0] == NULL: + data_ptr[0] = &_EMPTY_DATA_VIEW_SENTINEL @cython.internal @@ -1219,10 +1274,11 @@ cdef class _SegmentViews: if word_count > (PY_SSIZE_T_MAX // 8): raise OverflowError("segment is too large to expose as a Python buffer") byte_count = (8 * word_count) - self._views.append(_SegmentView()._init( + self._views.append(_BorrowedBufferView()._init( builder, - segments[i].begin(), - byte_count)) + segments[i].begin(), + byte_count, + True)) return self def __getitem__(self, index): @@ -1292,27 +1348,21 @@ cdef class _DynamicStructReader: return self.thisptr.hasByField(field.thisptr) cpdef get_data_as_view(self, field): - """ - Efficiently get a read-only memoryview for a DATA field without copying. - """ - cdef C_DynamicValue.Reader val - cdef capnp.Data.Reader temp_data - - try: - val = self.thisptr.get(field) - except KjException as e: - raise e._to_python() from None + """Efficiently get a read-only memoryview for a DATA field without copying. - if val.getType() != capnp.TYPE_DATA: - raise TypeError("Field '{}' is not a DATA field".format(field)) + .. warning:: + The returned memoryview *borrows* memory owned by this message. It stays valid while + the memoryview (and any object derived from it) is alive; an internal exporter pins + this reader for that duration. Do not let the message be mutated underneath an + outstanding view. - temp_data = val.asData() + An unset/empty DATA field yields a valid, zero-length view (it does not raise). + """ + cdef void* data_ptr + cdef size_t data_size - # Return read-only memoryview - cdef Py_buffer buf - if PyBuffer_FillInfo(&buf, self, temp_data.begin(), temp_data.size(), 1, PyBUF_CONTIG_RO) < 0: - raise KjException("Failed to create buffer info") - return PyMemoryView_FromBuffer(&buf) + _data_field_ptr_reader(self, field, &data_ptr, &data_size) + return _memoryview_borrowing(self, data_ptr, data_size, True) cpdef _which_str(self): try: @@ -1531,8 +1581,21 @@ cdef class _DynamicStructBuilder: cpdef to_segment_views(_DynamicStructBuilder self): """Returns the struct's containing message as zero-copy, read-only segment views. - The returned views borrow memory from the message builder. Do not mutate, reset, or reuse - the builder while the views are still in use. + The returned object is a sequence of read-only buffer-protocol views, one per output + segment. Each view borrows memory owned by the message builder; the segment pointers and + sizes are captured eagerly at call time (a snapshot). + + .. warning:: + The views (and any buffer exported from them, e.g. by ``memoryview()`` or by a + consumer that holds them) keep the builder pinned and remain valid only while no + mutation happens. Do NOT mutate, re-set, reset, or reuse the builder while any view or + exported buffer is still alive -- this includes calls that may allocate (e.g. getting + an unset pointer/struct/list field). Mutating after the snapshot can grow/relocate + segments, leaving the views pointing at stale or truncated data. Sharing the views + across threads or ``await`` points while the builder may change is a data race. + + Lifetime is enforced only by buffer-protocol reference counting (memory is not freed + while a view is held); correctness of the *contents* is the caller's responsibility. :rtype: sequence """ @@ -1731,30 +1794,29 @@ cdef class _DynamicStructBuilder: return _DynamicOrphan()._init(self.thisptr.disown(field), self._parent) cpdef get_data_as_view(self, field): - """ - Efficiently get a writable memoryview for a DATA field without copying. + """Efficiently get a writable memoryview for a DATA field without copying. - This allows in-place modification of the underlying buffer: - msg.get_data_as_view('myField')[0] = 0xFF - """ - cdef C_DynamicValue.Builder val - cdef capnp.Data.Builder temp_data + This allows in-place modification of the underlying buffer:: - try: - val = self.thisptr.get(field) - except KjException as e: - raise e._to_python() from None - - if val.getType() != capnp.TYPE_DATA: - raise TypeError("Field '{}' is not a DATA field".format(field)) + msg.get_data_as_view('myField')[0] = 0xFF - temp_data = val.asData() + .. warning:: + The returned memoryview *borrows* mutable memory owned by this message builder. It is + valid while the memoryview (and any object derived from it) is alive; an internal + exporter pins this builder for that duration. Do NOT mutate, re-set, reset, or reuse the + builder while a view is outstanding -- including calls that may allocate (e.g. getting + an unset pointer/struct/list field), since those can relocate or stale the borrowed + memory. Sharing a view across threads or ``await`` points while the builder may change + is a data race. + + An unset/empty DATA field yields a valid but zero-length view (writes are no-ops); to write + into the field, initialize it to the desired size first. + """ + cdef void* data_ptr + cdef size_t data_size - # Return writable memoryview - cdef Py_buffer buf - if PyBuffer_FillInfo(&buf, self, temp_data.begin(), temp_data.size(), 0, PyBUF_WRITABLE) < 0: - raise KjException("Failed to create buffer info") - return PyMemoryView_FromBuffer(&buf) + _data_field_ptr_builder(self, field, &data_ptr, &data_size) + return _memoryview_borrowing(self, data_ptr, data_size, False) cpdef as_reader(self): """A method for casting this Builder to a Reader diff --git a/test/test_get_data_view.py b/test/test_get_data_view.py index 22b88cbef..776366faf 100644 --- a/test/test_get_data_view.py +++ b/test/test_get_data_view.py @@ -1,4 +1,8 @@ import os +import tempfile +import weakref +from pathlib import Path + import pytest import capnp import sys @@ -152,6 +156,33 @@ def test_corner_cases_values(all_types): assert msg.get_data_as_view("dataField").tobytes() == binary_data +def test_uninitialized_data_get_view(all_types): + """ + Default DATA fields should expose an empty memoryview instead of failing on a NULL buffer pointer. + """ + builder = all_types.TestAllTypes.new_message() + builder_view = builder.get_data_as_view("dataField") + + assert isinstance(builder_view, memoryview) + assert builder_view.readonly is False + assert len(builder_view) == 0 + assert builder_view.tobytes() == b"" + + reader = all_types.TestAllTypes.new_message().as_reader() + reader_view = reader.get_data_as_view("dataField") + + assert isinstance(reader_view, memoryview) + assert reader_view.readonly is True + assert len(reader_view) == 0 + assert reader_view.tobytes() == b"" + + with pytest.raises(IndexError): + builder_view[0] = 0xFF + + with pytest.raises(ValueError): + builder_view[0:1] = b"\xff" + + def test_error_wrong_type(all_types): """ Test error handling: Calling get_data_as_view on non-Data fields. @@ -206,3 +237,56 @@ def test_view_keeps_message_alive(all_types): gc.collect() assert view.tobytes() == expected_data + + +def test_data_view_exports_through_buffer_exporter(all_types): + """Returned memoryviews should pin an internal exporter, not bare pointers.""" + msg = all_types.TestAllTypes.new_message() + msg.dataField = b"exporter_check" + view = msg.get_data_as_view("dataField") + + assert isinstance(view, memoryview) + assert view.obj is not None + assert len(view.obj) == len(view) + + +def test_data_view_survives_del_builder(all_types): + msg = all_types.TestAllTypes.new_message() + msg.dataField = b"persistence_check" + view = msg.get_data_as_view("dataField") + + del msg + gc.collect() + + assert view.tobytes() == b"persistence_check" + + +def test_data_view_releases_packed_payload(): + schema_text = """ + @0x9d7d4f087df9b6e1; + struct BlobMsg { + data @0 :Data; + } + """ + + class Payload(bytearray): + pass + + td = tempfile.TemporaryDirectory() + path = Path(td.name) / "blob.capnp" + path.write_text(schema_text) + schema = capnp.load(str(path)) + try: + payload = Payload(schema.BlobMsg.new_message(data=b"x" * 4096).to_bytes_packed()) + payload_ref = weakref.ref(payload) + + reader = schema.BlobMsg.from_bytes_packed(payload) + view = reader.get_data_as_view("data") + view.release() + + del view, reader, payload + gc.collect() + + assert payload_ref() is None + finally: + td.cleanup()