Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions cuda_core/cuda/core/_event.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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, <const void*><const char*>(ipc_descriptor._reserved), sizeof(data.reserved))
cdef Event self = Event.__new__(cls)
Expand Down
8 changes: 8 additions & 0 deletions cuda_core/cuda/core/_memory/_ipc.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down
8 changes: 8 additions & 0 deletions cuda_core/tests/memory_ipc/test_errors.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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
Expand Down
9 changes: 9 additions & 0 deletions cuda_core/tests/test_event.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading