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.
This commit is contained in:
Max Bachmann 2023-03-01 15:50:38 +01:00 committed by GitHub
parent d3d20743ee
commit c1748ed59d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 126 additions and 62 deletions

View File

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

View File

@ -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,
"RegQueryValue");
bufSize = retSize;
retBuf = (wchar_t *) PyMem_Malloc(bufSize);
if (retBuf == NULL)
return PyErr_NoMemory();
if (key == HKEY_PERFORMANCE_DATA) {
return PyErr_SetFromWindowsErrWithFunction(ERROR_INVALID_HANDLE,
"RegQueryValue");
}
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);
if (tmp == NULL) {
PyMem_Free(retBuf);
return PyErr_NoMemory();
}
retBuf = tmp;
void *tmp = PyMem_Realloc(pbuf != buf ? pbuf : NULL, size);
if (tmp == NULL) {
PyErr_NoMemory();
goto exit;
}
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]