gh-103092: Isolate winreg (#103250)

1 files changed, 154 insertions, 117 deletions

diff --git a/PC/winreg.c b/PC/winreg.c
index 15d32e7..4884125 100644
--- a/PC/winreg.c
+++ b/PC/winreg.c
@@ -15,15 +15,22 @@
 #define PY_SSIZE_T_CLEAN
 #include "Python.h"
 #include "pycore_object.h"        // _PyObject_Init()
+#include "pycore_moduleobject.h"
 #include "structmember.h"         // PyMemberDef
 #include <windows.h>
 
 #if defined(MS_WINDOWS_DESKTOP) || defined(MS_WINDOWS_SYSTEM) || defined(MS_WINDOWS_GAMES)
 
-static BOOL PyHKEY_AsHKEY(PyObject *ob, HKEY *pRes, BOOL bNoneOK);
-static BOOL clinic_HKEY_converter(PyObject *ob, void *p);
-static PyObject *PyHKEY_FromHKEY(HKEY h);
-static BOOL PyHKEY_Close(PyObject *obHandle);
+typedef struct {
+    PyTypeObject *PyHKEY_Type;
+} winreg_state;
+
+/* Forward declares */
+
+static BOOL PyHKEY_AsHKEY(winreg_state *st, PyObject *ob, HKEY *pRes, BOOL bNoneOK);
+static BOOL clinic_HKEY_converter(winreg_state *st, PyObject *ob, void *p);
+static PyObject *PyHKEY_FromHKEY(winreg_state *st, HKEY h);
+static BOOL PyHKEY_Close(winreg_state *st, PyObject *obHandle);
 
 static char errNotAHandle[] = "Object is not a handle";
 
@@ -35,8 +42,6 @@ static char errNotAHandle[] = "Object is not a handle";
 #define PyErr_SetFromWindowsErrWithFunction(rc, fnname) \
     PyErr_SetFromWindowsErr(rc)
 
-/* Forward declares */
-
 /* Doc strings */
 PyDoc_STRVAR(module_doc,
 "This module provides access to the Windows registry API.\n"
@@ -114,7 +119,7 @@ typedef struct {
     HKEY hkey;
 } PyHKEYObject;
 
-#define PyHKEY_Check(op) Py_IS_TYPE(op, &PyHKEY_Type)
+#define PyHKEY_Check(st, op) Py_IS_TYPE(op, st->PyHKEY_Type)
 
 static char *failMsg = "bad operand type";
 
@@ -147,7 +152,18 @@ PyHKEY_deallocFunc(PyObject *ob)
     PyHKEYObject *obkey = (PyHKEYObject *)ob;
     if (obkey->hkey)
         RegCloseKey((HKEY)obkey->hkey);
-    PyObject_Free(ob);
+
+    PyTypeObject *tp = Py_TYPE(ob);
+    PyObject_GC_UnTrack(ob);
+    PyObject_GC_Del(ob);
+    Py_DECREF(tp);
+}
+
+static int
+PyHKEY_traverseFunc(PyHKEYObject *self, visitproc visit, void *arg)
+{
+    Py_VISIT(Py_TYPE(self));
+    return 0;
 }
 
 static int
@@ -189,29 +205,6 @@ PyHKEY_hashFunc(PyObject *ob)
 }
 
 
-static PyNumberMethods PyHKEY_NumberMethods =
-{
-    PyHKEY_binaryFailureFunc,           /* nb_add */
-    PyHKEY_binaryFailureFunc,           /* nb_subtract */
-    PyHKEY_binaryFailureFunc,           /* nb_multiply */
-    PyHKEY_binaryFailureFunc,           /* nb_remainder */
-    PyHKEY_binaryFailureFunc,           /* nb_divmod */
-    PyHKEY_ternaryFailureFunc,          /* nb_power */
-    PyHKEY_unaryFailureFunc,            /* nb_negative */
-    PyHKEY_unaryFailureFunc,            /* nb_positive */
-    PyHKEY_unaryFailureFunc,            /* nb_absolute */
-    PyHKEY_boolFunc,                    /* nb_bool */
-    PyHKEY_unaryFailureFunc,            /* nb_invert */
-    PyHKEY_binaryFailureFunc,           /* nb_lshift */
-    PyHKEY_binaryFailureFunc,           /* nb_rshift */
-    PyHKEY_binaryFailureFunc,           /* nb_and */
-    PyHKEY_binaryFailureFunc,           /* nb_xor */
-    PyHKEY_binaryFailureFunc,           /* nb_or */
-    PyHKEY_intFunc,                     /* nb_int */
-    0,                                  /* nb_reserved */
-    PyHKEY_unaryFailureFunc,            /* nb_float */
-};
-
 /*[clinic input]
 module winreg
 class winreg.HKEYType "PyHKEYObject *" "&PyHKEY_Type"
@@ -229,6 +222,14 @@ class HKEY_converter(CConverter):
     type = 'HKEY'
     converter = 'clinic_HKEY_converter'
 
+    def parse_arg(self, argname, displayname):
+        return """
+        if (!{converter}(_PyModule_GetState(module), {argname}, &{paramname})) {{{{
+            goto exit;
+        }}}}
+        """.format(argname=argname, paramname=self.parser_name,
+                   converter=self.converter)
+
 class HKEY_return_converter(CReturnConverter):
     type = 'HKEY'
 
@@ -236,7 +237,7 @@ class HKEY_return_converter(CReturnConverter):
         self.declare(data)
         self.err_occurred_if_null_pointer("_return_value", data)
         data.return_conversion.append(
-            'return_value = PyHKEY_FromHKEY(_return_value);\n')
+            'return_value = PyHKEY_FromHKEY(_PyModule_GetState(module), _return_value);\n')
 
 # HACK: this only works for PyHKEYObjects, nothing else.
 #       Should this be generalized and enshrined in clinic.py,
@@ -249,7 +250,7 @@ class self_return_converter(CReturnConverter):
         data.return_conversion.append(
             'return_value = (PyObject *)_return_value;\n')
 [python start generated code]*/
-/*[python end generated code: output=da39a3ee5e6b4b0d input=2ebb7a4922d408d6]*/
+/*[python end generated code: output=da39a3ee5e6b4b0d input=17e645060c7b8ae1]*/
 
 #include "clinic/winreg.c.h"
 
@@ -270,8 +271,11 @@ static PyObject *
 winreg_HKEYType_Close_impl(PyHKEYObject *self)
 /*[clinic end generated code: output=fced3a624fb0c344 input=6786ac75f6b89de6]*/
 {
-    if (!PyHKEY_Close((PyObject *)self))
+    winreg_state *st = _PyType_GetModuleState(Py_TYPE(self));
+    assert(st != NULL);
+    if (!PyHKEY_Close(st, (PyObject *)self)) {
         return NULL;
+    }
     Py_RETURN_NONE;
 }
 
@@ -327,8 +331,11 @@ winreg_HKEYType___exit___impl(PyHKEYObject *self, PyObject *exc_type,
                               PyObject *exc_value, PyObject *traceback)
 /*[clinic end generated code: output=923ebe7389e6a263 input=fb32489ee92403c7]*/
 {
-    if (!PyHKEY_Close((PyObject *)self))
+    winreg_state *st = _PyType_GetModuleState(Py_TYPE(self));
+    assert(st != NULL);
+    if (!PyHKEY_Close(st, (PyObject *)self)) {
         return NULL;
+    }
     Py_RETURN_NONE;
 }
 
@@ -350,62 +357,71 @@ static PyMemberDef PyHKEY_memberlist[] = {
     {NULL}    /* Sentinel */
 };
 
-/* The type itself */
-PyTypeObject PyHKEY_Type =
-{
-    PyVarObject_HEAD_INIT(0, 0) /* fill in type at module init */
-    "PyHKEY",
-    sizeof(PyHKEYObject),
-    0,
-    PyHKEY_deallocFunc,                 /* tp_dealloc */
-    0,                                  /* tp_vectorcall_offset */
-    0,                                  /* tp_getattr */
-    0,                                  /* tp_setattr */
-    0,                                  /* tp_as_async */
-    0,                                  /* tp_repr */
-    &PyHKEY_NumberMethods,              /* tp_as_number */
-    0,                                  /* tp_as_sequence */
-    0,                                  /* tp_as_mapping */
-    PyHKEY_hashFunc,                    /* tp_hash */
-    0,                                  /* tp_call */
-    PyHKEY_strFunc,                     /* tp_str */
-    0,                                  /* tp_getattro */
-    0,                                  /* tp_setattro */
-    0,                                  /* tp_as_buffer */
-    0,                                  /* tp_flags */
-    PyHKEY_doc,                         /* tp_doc */
-    0,                                  /*tp_traverse*/
-    0,                                  /*tp_clear*/
-    0,                                  /*tp_richcompare*/
-    0,                                  /*tp_weaklistoffset*/
-    0,                                  /*tp_iter*/
-    0,                                  /*tp_iternext*/
-    PyHKEY_methods,                     /*tp_methods*/
-    PyHKEY_memberlist,                  /*tp_members*/
+static PyType_Slot pyhkey_type_slots[] = {
+    {Py_tp_dealloc, PyHKEY_deallocFunc},
+    {Py_tp_members, PyHKEY_memberlist},
+    {Py_tp_methods, PyHKEY_methods},
+    {Py_tp_doc, (char *)PyHKEY_doc},
+    {Py_tp_traverse, PyHKEY_traverseFunc},
+    {Py_tp_hash, PyHKEY_hashFunc},
+    {Py_tp_str, PyHKEY_strFunc},
+
+    // Number protocol
+    {Py_nb_add, PyHKEY_binaryFailureFunc},
+    {Py_nb_subtract, PyHKEY_binaryFailureFunc},
+    {Py_nb_multiply, PyHKEY_binaryFailureFunc},
+    {Py_nb_remainder, PyHKEY_binaryFailureFunc},
+    {Py_nb_divmod, PyHKEY_binaryFailureFunc},
+    {Py_nb_power, PyHKEY_ternaryFailureFunc},
+    {Py_nb_negative, PyHKEY_unaryFailureFunc},
+    {Py_nb_positive, PyHKEY_unaryFailureFunc},
+    {Py_nb_absolute, PyHKEY_unaryFailureFunc},
+    {Py_nb_bool, PyHKEY_boolFunc},
+    {Py_nb_invert, PyHKEY_unaryFailureFunc},
+    {Py_nb_lshift, PyHKEY_binaryFailureFunc},
+    {Py_nb_rshift, PyHKEY_binaryFailureFunc},
+    {Py_nb_and, PyHKEY_binaryFailureFunc},
+    {Py_nb_xor, PyHKEY_binaryFailureFunc},
+    {Py_nb_or, PyHKEY_binaryFailureFunc},
+    {Py_nb_int, PyHKEY_intFunc},
+    {Py_nb_float, PyHKEY_unaryFailureFunc},
+    {0, NULL},
+};
+
+static PyType_Spec pyhkey_type_spec = {
+    .name = "winreg.PyHKEY",
+    .basicsize = sizeof(PyHKEYObject),
+    .flags = (Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_IMMUTABLETYPE |
+              Py_TPFLAGS_DISALLOW_INSTANTIATION),
+    .slots = pyhkey_type_slots,
 };
 
 /************************************************************************
    The public PyHKEY API (well, not public yet :-)
 ************************************************************************/
 PyObject *
-PyHKEY_New(HKEY hInit)
+PyHKEY_New(PyObject *m, HKEY hInit)
 {
-    PyHKEYObject *key = PyObject_New(PyHKEYObject, &PyHKEY_Type);
-    if (key)
-        key->hkey = hInit;
+    winreg_state *st = _PyModule_GetState(m);
+    PyHKEYObject *key = PyObject_GC_New(PyHKEYObject, st->PyHKEY_Type);
+    if (key == NULL) {
+        return NULL;
+    }
+    key->hkey = hInit;
+    PyObject_GC_Track(key);
     return (PyObject *)key;
 }
 
 BOOL
-PyHKEY_Close(PyObject *ob_handle)
+PyHKEY_Close(winreg_state *st, PyObject *ob_handle)
 {
     LONG rc;
     HKEY key;
 
-    if (!PyHKEY_AsHKEY(ob_handle, &key, TRUE)) {
+    if (!PyHKEY_AsHKEY(st, ob_handle, &key, TRUE)) {
         return FALSE;
     }
-    if (PyHKEY_Check(ob_handle)) {
+    if (PyHKEY_Check(st, ob_handle)) {
         ((PyHKEYObject*)ob_handle)->hkey = 0;
     }
     rc = key ? RegCloseKey(key) : ERROR_SUCCESS;
@@ -415,7 +431,7 @@ PyHKEY_Close(PyObject *ob_handle)
 }
 
 BOOL
-PyHKEY_AsHKEY(PyObject *ob, HKEY *pHANDLE, BOOL bNoneOK)
+PyHKEY_AsHKEY(winreg_state *st, PyObject *ob, HKEY *pHANDLE, BOOL bNoneOK)
 {
     if (ob == Py_None) {
         if (!bNoneOK) {
@@ -426,7 +442,7 @@ PyHKEY_AsHKEY(PyObject *ob, HKEY *pHANDLE, BOOL bNoneOK)
         }
         *pHANDLE = (HKEY)0;
     }
-    else if (PyHKEY_Check(ob)) {
+    else if (PyHKEY_Check(st ,ob)) {
         PyHKEYObject *pH = (PyHKEYObject *)ob;
         *pHANDLE = pH->hkey;
     }
@@ -447,23 +463,24 @@ PyHKEY_AsHKEY(PyObject *ob, HKEY *pHANDLE, BOOL bNoneOK)
 }
 
 BOOL
-clinic_HKEY_converter(PyObject *ob, void *p)
+clinic_HKEY_converter(winreg_state *st, PyObject *ob, void *p)
 {
-    if (!PyHKEY_AsHKEY(ob, (HKEY *)p, FALSE))
+    if (!PyHKEY_AsHKEY(st, ob, (HKEY *)p, FALSE)) {
         return FALSE;
+    }
     return TRUE;
 }
 
 PyObject *
-PyHKEY_FromHKEY(HKEY h)
+PyHKEY_FromHKEY(winreg_state *st, HKEY h)
 {
-    /* Inline PyObject_New */
-    PyHKEYObject *op = (PyHKEYObject *) PyObject_Malloc(sizeof(PyHKEYObject));
+    PyHKEYObject *op = (PyHKEYObject *)PyObject_GC_New(PyHKEYObject,
+                                                       st->PyHKEY_Type);
     if (op == NULL) {
-        return PyErr_NoMemory();
+        return NULL;
     }
-    _PyObject_Init((PyObject*)op, &PyHKEY_Type);
     op->hkey = h;
+    PyObject_GC_Track(op);
     return (PyObject *)op;
 }
 
@@ -472,11 +489,11 @@ PyHKEY_FromHKEY(HKEY h)
   The module methods
 ************************************************************************/
 BOOL
-PyWinObject_CloseHKEY(PyObject *obHandle)
+PyWinObject_CloseHKEY(winreg_state *st, PyObject *obHandle)
 {
     BOOL ok;
-    if (PyHKEY_Check(obHandle)) {
-        ok = PyHKEY_Close(obHandle);
+    if (PyHKEY_Check(st, obHandle)) {
+        ok = PyHKEY_Close(st, obHandle);
     }
 #if SIZEOF_LONG >= SIZEOF_HKEY
     else if (PyLong_Check(obHandle)) {
@@ -826,8 +843,9 @@ static PyObject *
 winreg_CloseKey(PyObject *module, PyObject *hkey)
 /*[clinic end generated code: output=a4fa537019a80d15 input=5b1aac65ba5127ad]*/
 {
-    if (!PyHKEY_Close(hkey))
+    if (!PyHKEY_Close(_PyModule_GetState(module), hkey)) {
         return NULL;
+    }
     Py_RETURN_NONE;
 }
 
@@ -2061,7 +2079,7 @@ static struct PyMethodDef winreg_methods[] = {
 
 #define ADD_INT(VAL) do {                               \
     if (PyModule_AddIntConstant(m, #VAL, VAL) < 0) {    \
-        goto error;                                     \
+        return -1;                                      \
     }                                                   \
 } while (0)
 
@@ -2079,38 +2097,25 @@ inskey(PyObject *mod, char *name, HKEY key)
 
 #define ADD_KEY(VAL) do {           \
     if (inskey(m, #VAL, VAL) < 0) { \
-        goto error;                 \
+        return -1;                  \
     }                               \
 } while (0)
 
-
-static struct PyModuleDef winregmodule = {
-    PyModuleDef_HEAD_INIT,
-    "winreg",
-    module_doc,
-    -1,
-    winreg_methods,
-    NULL,
-    NULL,
-    NULL,
-    NULL
-};
-
-PyMODINIT_FUNC PyInit_winreg(void)
+static int
+exec_module(PyObject *m)
 {
-    PyObject *m = PyModule_Create(&winregmodule);
-    if (m == NULL) {
-        return NULL;
-    }
-    PyHKEY_Type.tp_doc = PyHKEY_doc;
-    if (PyType_Ready(&PyHKEY_Type) < 0) {
-        goto error;
+    winreg_state *st = (winreg_state *)_PyModule_GetState(m);
+
+    st->PyHKEY_Type = (PyTypeObject *)
+                       PyType_FromModuleAndSpec(m, &pyhkey_type_spec, NULL);
+    if (st->PyHKEY_Type == NULL) {
+        return -1;
     }
-    if (PyModule_AddObjectRef(m, "HKEYType", (PyObject *)&PyHKEY_Type) < 0) {
-        goto error;
+    if (PyModule_AddObjectRef(m, "HKEYType", (PyObject *)st->PyHKEY_Type) < 0) {
+        return -1;
     }
     if (PyModule_AddObjectRef(m, "error", PyExc_OSError) < 0) {
-        goto error;
+        return -1;
     }
 
     /* Add the relevant constants */
@@ -2174,12 +2179,44 @@ PyMODINIT_FUNC PyInit_winreg(void)
     ADD_INT(REG_RESOURCE_REQUIREMENTS_LIST);
 
 #undef ADD_INT
+    return 0;
+}
 
-    return m;
+static PyModuleDef_Slot winreg_slots[] = {
+    {Py_mod_exec, exec_module},
+    {0, NULL}
+};
 
-error:
-    Py_DECREF(m);
-    return NULL;
+static int
+winreg_traverse(PyObject *module, visitproc visit, void *arg)
+{
+    winreg_state *state = _PyModule_GetState(module);
+    Py_VISIT(state->PyHKEY_Type);
+    return 0;
+}
+
+static int
+winreg_clear(PyObject *module)
+{
+    winreg_state *state = _PyModule_GetState(module);
+    Py_CLEAR(state->PyHKEY_Type);
+    return 0;
+}
+
+static struct PyModuleDef winregmodule = {
+    .m_base = PyModuleDef_HEAD_INIT,
+    .m_name = "winreg",
+    .m_doc = module_doc,
+    .m_size = sizeof(winreg_state),
+    .m_methods = winreg_methods,
+    .m_slots = winreg_slots,
+    .m_traverse = winreg_traverse,
+    .m_clear = winreg_clear,
+};
+
+PyMODINIT_FUNC PyInit_winreg(void)
+{
+    return PyModuleDef_Init(&winregmodule);
 }
 
 #endif /* MS_WINDOWS_DESKTOP || MS_WINDOWS_SYSTEM || MS_WINDOWS_GAMES */