From e603e0eb299d46220a633fd8b405edf2ad2544a9 Mon Sep 17 00:00:00 2001 From: waterWang Date: Wed, 26 Aug 2026 20:57:23 +0800 Subject: [PATCH] fix(core): reject mismatched UUID in register() even under -O (#2697) DeviceMemoryResource.register() and PinnedMemoryResource.register() validated the supplied UUID against the resource's own with a bare `assert`, which CPython removes under `-O`. A mismatched UUID was then accepted, registered as a foreign registry key, and rewritten into the resource's own UUID, silently corrupting the IPC registry. Replace the assertion with an explicit ValueError so the check is enforced regardless of the interpreter's assertion mode, and convert the internal register() result check in MP_from_allocation_handle to an explicit RuntimeError as well. Signed-off-by: waterWang --- cuda_core/cuda/core/_memory/_ipc.pyx | 12 ++++++++-- cuda_core/tests/memory_ipc/test_errors.py | 28 +++++++++++++++++++++++ 2 files changed, 38 insertions(+), 2 deletions(-) diff --git a/cuda_core/cuda/core/_memory/_ipc.pyx b/cuda_core/cuda/core/_memory/_ipc.pyx index ae8db6589b4..f966793264a 100644 --- a/cuda_core/cuda/core/_memory/_ipc.pyx +++ b/cuda_core/cuda/core/_memory/_ipc.pyx @@ -269,7 +269,11 @@ cdef _MemPool MP_from_allocation_handle(cls, alloc_handle): # Register it. if uuid is not None: registered = self.register(uuid) - assert registered is self + if registered is not self: + raise RuntimeError( + "Internal error: register() returned a different memory " + "resource than the one being registered" + ) return self @@ -292,7 +296,11 @@ cdef _MemPool MP_register(_MemPool self, uuid): return existing if not self.is_ipc_enabled: raise RuntimeError("Memory resource is not IPC-enabled") - assert self.uuid is None or self.uuid == uuid + if self.uuid is not None and self.uuid != uuid: + raise ValueError( + f"Cannot register memory resource with UUID {uuid}: " + f"the resource already has UUID {self.uuid}" + ) registry[uuid] = self self._ipc_data._alloc_handle._uuid = uuid return self diff --git a/cuda_core/tests/memory_ipc/test_errors.py b/cuda_core/tests/memory_ipc/test_errors.py index 8038d62570c..93aceec34b1 100644 --- a/cuda_core/tests/memory_ipc/test_errors.py +++ b/cuda_core/tests/memory_ipc/test_errors.py @@ -72,6 +72,34 @@ def test_register_rejects_non_ipc_memory_resource(mempool_device): DeviceMemoryResource.from_registry(key) +@pytest.mark.human_authored +def test_register_rejects_mismatched_uuid(ipc_memory_resource): + """register() rejects a UUID that does not match the resource's own. + + The check is an explicit ValueError (not a bare ``assert``) so it remains + enforced even when CPython runs with ``-O``, where assertions are removed + (see #2697). Registering under a foreign key would otherwise rewrite the + resource's UUID and silently corrupt the IPC registry. + """ + mr = ipc_memory_resource + assert mr.is_ipc_enabled + + # The resource must already carry its own UUID (assigned on creation). + assert mr.uuid is not None + other = uuid.uuid4() + while other == mr.uuid: + other = uuid.uuid4() + + with pytest.raises(ValueError, match="already has UUID"): + mr.register(other) + + # A rejected registration must not rewrite the resource's own UUID. + assert mr.uuid != other + + # Registering under the resource's own UUID still succeeds. + assert mr.register(mr.uuid) is mr + + @pytest.mark.skipif(os.name == "nt", reason="IPC allocation handles are not supported on Windows") @pytest.mark.agent_authored(model="gpt-5.6") def test_ipc_allocation_handle_state_tracks_close():