gh-102344: Reimplement winreg QueryValue / SetValue using QueryValueEx / SetValueEx (GH-102345)

The newer APIs are more widely available than the old ones, and are called in a way to preserve functionality.

Author: maxbachmann

Misc/NEWS.d/next/Windows/2023-03-01-01-36-39.gh-issue-102344.Dgfux4.rst

@@ -0,0 +1,2 @@
+Implement ``winreg.QueryValue`` using ``QueryValueEx`` and
+``winreg.SetValue`` using ``SetValueEx``. Patch by Max Bachmann.

PC/winreg.c

@@ -561,7 +561,7 @@ Py2Reg(PyObject *value, DWORD typ, BYTE **retDataBuf, DWORD *retDataSize)
 {
     Py_ssize_t i,j;
     switch (typ) {
-        case REG_DWORD: 
+        case REG_DWORD:
             {
                 if (value != Py_None && !PyLong_Check(value)) {
                     return FALSE;
@@ -585,7 +585,7 @@ Py2Reg(PyObject *value, DWORD typ, BYTE **retDataBuf, DWORD *retDataSize)
                 *retDataSize = sizeof(DWORD);
                 break;
             }
-        case REG_QWORD: 
+        case REG_QWORD:
             {
                 if (value != Py_None && !PyLong_Check(value)) {
                     return FALSE;
@@ -1488,53 +1488,77 @@ static PyObject *
 winreg_QueryValue_impl(PyObject *module, HKEY key, const Py_UNICODE *sub_key)
 /*[clinic end generated code: output=c655810ae50c63a9 input=41cafbbf423b21d6]*/
 {
-    long rc;
-    PyObject *retStr;
-    wchar_t *retBuf;
-    DWORD bufSize = 0;
-    DWORD retSize = 0;
-    wchar_t *tmp;
+    LONG rc;
+    HKEY childKey = key;
+    WCHAR buf[256], *pbuf = buf;
+    DWORD size = sizeof(buf);
+    DWORD type;
+    Py_ssize_t length;
+    PyObject *result = NULL;
 
     if (PySys_Audit("winreg.QueryValue", "nuu",
-                    (Py_ssize_t)key, sub_key, NULL) < 0) {
+                    (Py_ssize_t)key, sub_key, NULL) < 0)
+    {
         return NULL;
     }
-    rc = RegQueryValueW(key, sub_key, NULL, &retSize);
-    if (rc == ERROR_MORE_DATA)
-        retSize = 256;
-    else if (rc != ERROR_SUCCESS)
-        return PyErr_SetFromWindowsErrWithFunction(rc,
+
+    if (key == HKEY_PERFORMANCE_DATA) {
+        return PyErr_SetFromWindowsErrWithFunction(ERROR_INVALID_HANDLE,
                                                    "RegQueryValue");
+    }
 
-    bufSize = retSize;
-    retBuf = (wchar_t *) PyMem_Malloc(bufSize);
-    if (retBuf == NULL)
-        return PyErr_NoMemory();
+    if (sub_key && sub_key[0]) {
+        Py_BEGIN_ALLOW_THREADS
+        rc = RegOpenKeyExW(key, sub_key, 0, KEY_QUERY_VALUE, &childKey);
+        Py_END_ALLOW_THREADS
+        if (rc != ERROR_SUCCESS) {
+            return PyErr_SetFromWindowsErrWithFunction(rc, "RegOpenKeyEx");
+        }
+    }
 
     while (1) {
-        retSize = bufSize;
-        rc = RegQueryValueW(key, sub_key, retBuf, &retSize);
-        if (rc != ERROR_MORE_DATA)
+        Py_BEGIN_ALLOW_THREADS
+        rc = RegQueryValueExW(childKey, NULL, NULL, &type, (LPBYTE)pbuf,
+                              &size);
+        Py_END_ALLOW_THREADS
+        if (rc != ERROR_MORE_DATA) {
             break;
-
-        bufSize *= 2;
-        tmp = (wchar_t *) PyMem_Realloc(retBuf, bufSize);
+        }
+        void *tmp = PyMem_Realloc(pbuf != buf ? pbuf : NULL, size);
         if (tmp == NULL) {
-            PyMem_Free(retBuf);
-            return PyErr_NoMemory();
+            PyErr_NoMemory();
+            goto exit;
         }
-        retBuf = tmp;
+        pbuf = tmp;
     }
 
-    if (rc != ERROR_SUCCESS) {
-        PyMem_Free(retBuf);
-        return PyErr_SetFromWindowsErrWithFunction(rc,
-                                                   "RegQueryValue");
+    if (rc == ERROR_SUCCESS) {
+        if (type != REG_SZ) {
+            PyErr_SetFromWindowsErrWithFunction(ERROR_INVALID_DATA,
+                                                "RegQueryValue");
+            goto exit;
+        }
+        length = wcsnlen(pbuf, size / sizeof(WCHAR));
+    }
+    else if (rc == ERROR_FILE_NOT_FOUND) {
+        // Return an empty string if there's no default value.
+        length = 0;
+    }
+    else {
+        PyErr_SetFromWindowsErrWithFunction(rc, "RegQueryValueEx");
+        goto exit;
     }
 
-    retStr = PyUnicode_FromWideChar(retBuf, wcslen(retBuf));
-    PyMem_Free(retBuf);
-    return retStr;
+    result = PyUnicode_FromWideChar(pbuf, length);
+
+exit:
+    if (pbuf != buf) {
+        PyMem_Free(pbuf);
+    }
+    if (childKey != key) {
+        RegCloseKey(childKey);
+    }
+    return result;
 }
 
 
@@ -1687,38 +1711,69 @@ winreg_SetValue_impl(PyObject *module, HKEY key, const Py_UNICODE *sub_key,
                      DWORD type, PyObject *value_obj)
 /*[clinic end generated code: output=d4773dc9c372311a input=bf088494ae2d24fd]*/
 {
-    Py_ssize_t value_length;
-    long rc;
+    LONG rc;
+    HKEY childKey = key;
+    LPWSTR value;
+    Py_ssize_t size;
+    Py_ssize_t length;
+    PyObject *result = NULL;
 
     if (type != REG_SZ) {
         PyErr_SetString(PyExc_TypeError, "type must be winreg.REG_SZ");
         return NULL;
     }
 
-    wchar_t *value = PyUnicode_AsWideCharString(value_obj, &value_length);
+    value = PyUnicode_AsWideCharString(value_obj, &length);
     if (value == NULL) {
         return NULL;
     }
-    if ((Py_ssize_t)(DWORD)value_length != value_length) {
+
+    size = (length + 1) * sizeof(WCHAR);
+    if ((Py_ssize_t)(DWORD)size != size) {
         PyErr_SetString(PyExc_OverflowError, "value is too long");
-        PyMem_Free(value);
-        return NULL;
+        goto exit;
     }
 
     if (PySys_Audit("winreg.SetValue", "nunu#",
                     (Py_ssize_t)key, sub_key, (Py_ssize_t)type,
-                    value, value_length) < 0) {
-        PyMem_Free(value);
-        return NULL;
+                    value, length) < 0)
+    {
+        goto exit;
+    }
+
+    if (key == HKEY_PERFORMANCE_DATA) {
+        PyErr_SetFromWindowsErrWithFunction(ERROR_INVALID_HANDLE,
+                                            "RegSetValue");
+        goto exit;
+    }
+
+    if (sub_key && sub_key[0]) {
+        Py_BEGIN_ALLOW_THREADS
+        rc = RegCreateKeyExW(key, sub_key, 0, NULL, 0, KEY_SET_VALUE, NULL,
+                             &childKey, NULL);
+        Py_END_ALLOW_THREADS
+        if (rc != ERROR_SUCCESS) {
+            PyErr_SetFromWindowsErrWithFunction(rc, "RegCreateKeyEx");
+            goto exit;
+        }
     }
 
     Py_BEGIN_ALLOW_THREADS
-    rc = RegSetValueW(key, sub_key, REG_SZ, value, (DWORD)(value_length + 1));
+    rc = RegSetValueExW(childKey, NULL, 0, REG_SZ, (LPBYTE)value, (DWORD)size);
     Py_END_ALLOW_THREADS
+    if (rc == ERROR_SUCCESS) {
+        result = Py_NewRef(Py_None);
+    }
+    else {
+        PyErr_SetFromWindowsErrWithFunction(rc, "RegSetValueEx");
+    }
+
+exit:
     PyMem_Free(value);
-    if (rc != ERROR_SUCCESS)
-        return PyErr_SetFromWindowsErrWithFunction(rc, "RegSetValue");
-    Py_RETURN_NONE;
+    if (childKey != key) {
+        RegCloseKey(childKey);
+    }
+    return result;
 }
 
 /*[clinic input]
@@ -1771,32 +1826,39 @@ winreg_SetValueEx_impl(PyObject *module, HKEY key,
                        DWORD type, PyObject *value)
 /*[clinic end generated code: output=811b769a66ae11b7 input=900a9e3990bfb196]*/
 {
-    BYTE *data;
-    DWORD len;
-
     LONG rc;
+    BYTE *data = NULL;
+    DWORD size;
+    PyObject *result = NULL;
 
-    if (!Py2Reg(value, type, &data, &len))
+    if (!Py2Reg(value, type, &data, &size))
     {
-        if (!PyErr_Occurred())
+        if (!PyErr_Occurred()) {
             PyErr_SetString(PyExc_ValueError,
                      "Could not convert the data to the specified type.");
+        }
         return NULL;
     }
     if (PySys_Audit("winreg.SetValue", "nunO",
                     (Py_ssize_t)key, value_name, (Py_ssize_t)type,
-                    value) < 0) {
-        PyMem_Free(data);
-        return NULL;
+                    value) < 0)
+    {
+        goto exit;
     }
+
     Py_BEGIN_ALLOW_THREADS
-    rc = RegSetValueExW(key, value_name, 0, type, data, len);
+    rc = RegSetValueExW(key, value_name, 0, type, data, size);
     Py_END_ALLOW_THREADS
+    if (rc == ERROR_SUCCESS) {
+        result = Py_NewRef(Py_None);
+    }
+    else {
+        PyErr_SetFromWindowsErrWithFunction(rc, "RegSetValueEx");
+    }
+
+exit:
     PyMem_Free(data);
-    if (rc != ERROR_SUCCESS)
-        return PyErr_SetFromWindowsErrWithFunction(rc,
-                                                   "RegSetValueEx");
-    Py_RETURN_NONE;
+    return result;
 }
 
 /*[clinic input]