diff --git a/cuda_core/cuda/core/_event.pyx b/cuda_core/cuda/core/_event.pyx index d38f88149f..e5cb81ac41 100644 --- a/cuda_core/cuda/core/_event.pyx +++ b/cuda_core/cuda/core/_event.pyx @@ -5,6 +5,7 @@ from __future__ import annotations cimport cpython +from libc.stddef cimport size_t from libc.string cimport memcpy from cuda.bindings cimport cydriver from cuda.core._context cimport Context @@ -232,6 +233,13 @@ cdef class Event: A new event backed by the imported IPC handle. """ + cdef size_t reserved_size = len(ipc_descriptor._reserved) + cdef size_t expected_size = sizeof(cydriver.CUipcEventHandle) + if reserved_size < expected_size: + raise ValueError( + f"IPC event descriptor reserved field is {reserved_size} bytes; " + f"expected at least {expected_size}" + ) cdef cydriver.CUipcEventHandle data memcpy(data.reserved, (ipc_descriptor._reserved), sizeof(data.reserved)) cdef Event self = Event.__new__(cls) diff --git a/cuda_core/cuda/core/_memory/_ipc.pyx b/cuda_core/cuda/core/_memory/_ipc.pyx index 61fd8b086e..879a198d08 100644 --- a/cuda_core/cuda/core/_memory/_ipc.pyx +++ b/cuda_core/cuda/core/_memory/_ipc.pyx @@ -4,6 +4,7 @@ cimport cpython +from libc.stddef cimport size_t from cuda.bindings cimport cydriver from cuda.core._memory._buffer cimport Buffer, Buffer_from_deviceptr_handle from cuda.core._memory._memory_pool cimport _MemPool @@ -171,6 +172,13 @@ cdef Buffer Buffer_from_ipc_descriptor( """Import a buffer that was exported from another process.""" if not mr.is_ipc_enabled: raise RuntimeError("Memory resource is not IPC-enabled") + cdef size_t payload_size = len(ipc_descriptor._payload) + cdef size_t expected_size = sizeof(cydriver.CUmemPoolPtrExportData) + if payload_size < expected_size: + raise ValueError( + f"IPC buffer descriptor payload is {payload_size} bytes; " + f"expected at least {expected_size}" + ) cdef Stream s = Stream_accept(stream) cdef DevicePtrHandle h_ptr = deviceptr_import_ipc( mr._h_pool, diff --git a/cuda_core/tests/memory_ipc/test_errors.py b/cuda_core/tests/memory_ipc/test_errors.py index 84fb9093e1..326a077f6e 100644 --- a/cuda_core/tests/memory_ipc/test_errors.py +++ b/cuda_core/tests/memory_ipc/test_errors.py @@ -9,6 +9,7 @@ from helpers.child_processes import child_timeout_sec, kill_subprocesses from cuda.core import Buffer, Device, DeviceMemoryResource, DeviceMemoryResourceOptions +from cuda.core._memory import IPCBufferDescriptor from cuda.core._utils.cuda_utils import CUDAError CHILD_TIMEOUT_SEC = child_timeout_sec() @@ -31,6 +32,13 @@ def test_outer_timeout_marker_is_applied(request): assert marker.args == (expected,), f"unexpected timeout value: {marker.args!r}" +def test_import_truncated_buffer_descriptor(ipc_device, ipc_memory_resource): + """Truncated IPC buffer descriptor payload is rejected before driver import.""" + desc = IPCBufferDescriptor._init(b"\x00" * 8, NBYTES) + with pytest.raises(ValueError, match=r"payload is 8 bytes; expected at least 64"): + Buffer.from_ipc_descriptor(ipc_memory_resource, desc, stream=ipc_device.default_stream) + + def test_ipc_allocation_handle_rejects_negative_fd(): """Negative fds are rejected even when CPython runs with -O (Glasswing V3.2).""" from cuda.core._memory._ipc import IPCAllocationHandle diff --git a/cuda_core/tests/test_event.py b/cuda_core/tests/test_event.py index c85efa1925..79f0090ace 100644 --- a/cuda_core/tests/test_event.py +++ b/cuda_core/tests/test_event.py @@ -235,6 +235,15 @@ def test_event_is_done_false(init_cuda): event.sync() +def test_import_truncated_event_descriptor(): + """Truncated IPC event descriptor reserved field is rejected before memcpy.""" + import cuda.core._event as _event_module + + desc = _event_module.IPCEventDescriptor._init(b"\x00" * 8, True) + with pytest.raises(ValueError, match=r"reserved field is 8 bytes; expected at least 64"): + Event.from_ipc_descriptor(desc) + + def test_ipc_event_descriptor_direct_init(): """IPCEventDescriptor cannot be instantiated directly.""" import cuda.core._event as _event_module