BUG: PyDataMem_SetHandler check capsule name (numpy#26529)

Closes numpy#26137.

Tests capsule name and sets PyErr if not valid:

    if (!PyCapsule_IsValid(handler, MEM_HANDLER_CAPSULE_NAME)) {
        PyErr_SetString(PyExc_ValueError, "Capsule must be named 'mem_handler'")
        return NULL;
    }

Author: luxedo

numpy/_core/src/multiarray/alloc.c

@@ -354,7 +354,8 @@ NPY_NO_EXPORT void *
 PyDataMem_UserNEW(size_t size, PyObject *mem_handler)
 {
     void *result;
-    PyDataMem_Handler *handler = (PyDataMem_Handler *) PyCapsule_GetPointer(mem_handler, "mem_handler");
+    PyDataMem_Handler *handler = (PyDataMem_Handler *) PyCapsule_GetPointer(
+            mem_handler, MEM_HANDLER_CAPSULE_NAME);
     if (handler == NULL) {
         return NULL;
     }
@@ -368,7 +369,8 @@ NPY_NO_EXPORT void *
 PyDataMem_UserNEW_ZEROED(size_t nmemb, size_t size, PyObject *mem_handler)
 {
     void *result;
-    PyDataMem_Handler *handler = (PyDataMem_Handler *) PyCapsule_GetPointer(mem_handler, "mem_handler");
+    PyDataMem_Handler *handler = (PyDataMem_Handler *) PyCapsule_GetPointer(
+            mem_handler, MEM_HANDLER_CAPSULE_NAME);
     if (handler == NULL) {
         return NULL;
     }
@@ -381,7 +383,8 @@ PyDataMem_UserNEW_ZEROED(size_t nmemb, size_t size, PyObject *mem_handler)
 NPY_NO_EXPORT void
 PyDataMem_UserFREE(void *ptr, size_t size, PyObject *mem_handler)
 {
-    PyDataMem_Handler *handler = (PyDataMem_Handler *) PyCapsule_GetPointer(mem_handler, "mem_handler");
+    PyDataMem_Handler *handler = (PyDataMem_Handler *) PyCapsule_GetPointer(
+            mem_handler, MEM_HANDLER_CAPSULE_NAME);
     if (handler == NULL) {
         WARN_NO_RETURN(PyExc_RuntimeWarning,
                      "Could not get pointer to 'mem_handler' from PyCapsule");
@@ -395,7 +398,8 @@ NPY_NO_EXPORT void *
 PyDataMem_UserRENEW(void *ptr, size_t size, PyObject *mem_handler)
 {
     void *result;
-    PyDataMem_Handler *handler = (PyDataMem_Handler *) PyCapsule_GetPointer(mem_handler, "mem_handler");
+    PyDataMem_Handler *handler = (PyDataMem_Handler *) PyCapsule_GetPointer(
+            mem_handler, MEM_HANDLER_CAPSULE_NAME);
     if (handler == NULL) {
         return NULL;
     }
@@ -427,6 +431,10 @@ PyDataMem_SetHandler(PyObject *handler)
     if (handler == NULL) {
         handler = PyDataMem_DefaultHandler;
     }
+    if (!PyCapsule_IsValid(handler, MEM_HANDLER_CAPSULE_NAME)) {
+        PyErr_SetString(PyExc_ValueError, "Capsule must be named 'mem_handler'");
+        return NULL;
+    }
     token = PyContextVar_Set(current_handler, handler);
     if (token == NULL) {
         Py_DECREF(old_handler);
@@ -477,7 +485,8 @@ get_handler_name(PyObject *NPY_UNUSED(self), PyObject *args)
             return NULL;
         }
     }
-    handler = (PyDataMem_Handler *) PyCapsule_GetPointer(mem_handler, "mem_handler");
+    handler = (PyDataMem_Handler *) PyCapsule_GetPointer(
+            mem_handler, MEM_HANDLER_CAPSULE_NAME);
     if (handler == NULL) {
         Py_DECREF(mem_handler);
         return NULL;
@@ -514,7 +523,8 @@ get_handler_version(PyObject *NPY_UNUSED(self), PyObject *args)
             return NULL;
         }
     }
-    handler = (PyDataMem_Handler *) PyCapsule_GetPointer(mem_handler, "mem_handler");
+    handler = (PyDataMem_Handler *) PyCapsule_GetPointer(
+            mem_handler, MEM_HANDLER_CAPSULE_NAME);
     if (handler == NULL) {
         Py_DECREF(mem_handler);
         return NULL;

numpy/_core/src/multiarray/alloc.h

@@ -6,6 +6,7 @@
 #include "numpy/ndarraytypes.h"
 
 #define NPY_TRACE_DOMAIN 389047
+#define MEM_HANDLER_CAPSULE_NAME "mem_handler"
 
 NPY_NO_EXPORT PyObject *
 _get_madvise_hugepage(PyObject *NPY_UNUSED(self), PyObject *NPY_UNUSED(args));

numpy/_core/src/multiarray/multiarraymodule.c

@@ -5230,7 +5230,8 @@ PyMODINIT_FUNC PyInit__multiarray_umath(void) {
     /*
      * Initialize the default PyDataMem_Handler capsule singleton.
      */
-    PyDataMem_DefaultHandler = PyCapsule_New(&default_handler, "mem_handler", NULL);
+    PyDataMem_DefaultHandler = PyCapsule_New(
+            &default_handler, MEM_HANDLER_CAPSULE_NAME, NULL);
     if (PyDataMem_DefaultHandler == NULL) {
         goto err;
     }

numpy/_core/tests/__init__.py

@@ -46,6 +46,16 @@ def get_module(tmp_path):
              Py_DECREF(secret_data);
              return old;
          """),
+        ("set_wrong_capsule_name_data_policy", "METH_NOARGS", """
+             PyObject *wrong_name_capsule =
+                 PyCapsule_New(&secret_data_handler, "not_mem_handler", NULL);
+             if (wrong_name_capsule == NULL) {
+                 return NULL;
+             }
+             PyObject *old = PyDataMem_SetHandler(wrong_name_capsule);
+             Py_DECREF(wrong_name_capsule);
+             return old;
+         """),
         ("set_old_policy", "METH_O", """
              PyObject *old;
              if (args != NULL && PyCapsule_CheckExact(args)) {
@@ -252,6 +262,10 @@ def test_set_policy(get_module):
         get_module.set_old_policy(orig_policy)
         assert get_handler_name() == orig_policy_name
 
+    with pytest.raises(ValueError,
+                       match="Capsule must be named 'mem_handler'"):
+        get_module.set_wrong_capsule_name_data_policy()
+
 
 @pytest.mark.skipif(sys.version_info >= (3, 12), reason="no numpy.distutils")
 def test_default_policy_singleton(get_module):