Skip to content

Commit c1748ed

Browse files
maxbachmannmaxbachmann
authored
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.
1 parent d3d2074 commit c1748ed

File tree

2 files changed

+124
-60
lines changed

2 files changed

+124
-60
lines changed

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
Implement ``winreg.QueryValue`` using ``QueryValueEx`` and
2+
``winreg.SetValue`` using ``SetValueEx``. Patch by Max Bachmann.

‎PC/winreg.c

Lines changed: 122 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -561,7 +561,7 @@ Py2Reg(PyObject *value, DWORD typ, BYTE **retDataBuf, DWORD *retDataSize)
561561
{
562562
Py_ssize_t i,j;
563563
switch (typ) {
564-
case REG_DWORD:
564+
case REG_DWORD:
565565
{
566566
if (value != Py_None && !PyLong_Check(value)) {
567567
return FALSE;
@@ -585,7 +585,7 @@ Py2Reg(PyObject *value, DWORD typ, BYTE **retDataBuf, DWORD *retDataSize)
585585
*retDataSize = sizeof(DWORD);
586586
break;
587587
}
588-
case REG_QWORD:
588+
case REG_QWORD:
589589
{
590590
if (value != Py_None && !PyLong_Check(value)) {
591591
return FALSE;
@@ -1488,53 +1488,77 @@ static PyObject *
14881488
winreg_QueryValue_impl(PyObject *module, HKEY key, const Py_UNICODE *sub_key)
14891489
/*[clinic end generated code: output=c655810ae50c63a9 input=41cafbbf423b21d6]*/
14901490
{
1491-
long rc;
1492-
PyObject *retStr;
1493-
wchar_t *retBuf;
1494-
DWORD bufSize = 0;
1495-
DWORD retSize = 0;
1496-
wchar_t *tmp;
1491+
LONG rc;
1492+
HKEY childKey = key;
1493+
WCHAR buf[256], *pbuf = buf;
1494+
DWORD size = sizeof(buf);
1495+
DWORD type;
1496+
Py_ssize_t length;
1497+
PyObject *result = NULL;
14971498

14981499
if (PySys_Audit("winreg.QueryValue", "nuu",
1499-
(Py_ssize_t)key, sub_key, NULL) < 0) {
1500+
(Py_ssize_t)key, sub_key, NULL) < 0)
1501+
{
15001502
return NULL;
15011503
}
1502-
rc = RegQueryValueW(key, sub_key, NULL, &retSize);
1503-
if (rc == ERROR_MORE_DATA)
1504-
retSize = 256;
1505-
else if (rc != ERROR_SUCCESS)
1506-
return PyErr_SetFromWindowsErrWithFunction(rc,
1504+
1505+
if (key == HKEY_PERFORMANCE_DATA) {
1506+
return PyErr_SetFromWindowsErrWithFunction(ERROR_INVALID_HANDLE,
15071507
"RegQueryValue");
1508+
}
15081509

1509-
bufSize = retSize;
1510-
retBuf = (wchar_t *) PyMem_Malloc(bufSize);
1511-
if (retBuf == NULL)
1512-
return PyErr_NoMemory();
1510+
if (sub_key && sub_key[0]) {
1511+
Py_BEGIN_ALLOW_THREADS
1512+
rc = RegOpenKeyExW(key, sub_key, 0, KEY_QUERY_VALUE, &childKey);
1513+
Py_END_ALLOW_THREADS
1514+
if (rc != ERROR_SUCCESS) {
1515+
return PyErr_SetFromWindowsErrWithFunction(rc, "RegOpenKeyEx");
1516+
}
1517+
}
15131518

15141519
while (1) {
1515-
retSize = bufSize;
1516-
rc = RegQueryValueW(key, sub_key, retBuf, &retSize);
1517-
if (rc != ERROR_MORE_DATA)
1520+
Py_BEGIN_ALLOW_THREADS
1521+
rc = RegQueryValueExW(childKey, NULL, NULL, &type, (LPBYTE)pbuf,
1522+
&size);
1523+
Py_END_ALLOW_THREADS
1524+
if (rc != ERROR_MORE_DATA) {
15181525
break;
1519-
1520-
bufSize *= 2;
1521-
tmp = (wchar_t *) PyMem_Realloc(retBuf, bufSize);
1526+
}
1527+
void *tmp = PyMem_Realloc(pbuf != buf ? pbuf : NULL, size);
15221528
if (tmp == NULL) {
1523-
PyMem_Free(retBuf);
1524-
return PyErr_NoMemory();
1529+
PyErr_NoMemory();
1530+
goto exit;
15251531
}
1526-
retBuf = tmp;
1532+
pbuf = tmp;
15271533
}
15281534

1529-
if (rc != ERROR_SUCCESS) {
1530-
PyMem_Free(retBuf);
1531-
return PyErr_SetFromWindowsErrWithFunction(rc,
1532-
"RegQueryValue");
1535+
if (rc == ERROR_SUCCESS) {
1536+
if (type != REG_SZ) {
1537+
PyErr_SetFromWindowsErrWithFunction(ERROR_INVALID_DATA,
1538+
"RegQueryValue");
1539+
goto exit;
1540+
}
1541+
length = wcsnlen(pbuf, size / sizeof(WCHAR));
1542+
}
1543+
else if (rc == ERROR_FILE_NOT_FOUND) {
1544+
// Return an empty string if there's no default value.
1545+
length = 0;
1546+
}
1547+
else {
1548+
PyErr_SetFromWindowsErrWithFunction(rc, "RegQueryValueEx");
1549+
goto exit;
15331550
}
15341551

1535-
retStr = PyUnicode_FromWideChar(retBuf, wcslen(retBuf));
1536-
PyMem_Free(retBuf);
1537-
return retStr;
1552+
result = PyUnicode_FromWideChar(pbuf, length);
1553+
1554+
exit:
1555+
if (pbuf != buf) {
1556+
PyMem_Free(pbuf);
1557+
}
1558+
if (childKey != key) {
1559+
RegCloseKey(childKey);
1560+
}
1561+
return result;
15381562
}
15391563

15401564

@@ -1687,38 +1711,69 @@ winreg_SetValue_impl(PyObject *module, HKEY key, const Py_UNICODE *sub_key,
16871711
DWORD type, PyObject *value_obj)
16881712
/*[clinic end generated code: output=d4773dc9c372311a input=bf088494ae2d24fd]*/
16891713
{
1690-
Py_ssize_t value_length;
1691-
long rc;
1714+
LONG rc;
1715+
HKEY childKey = key;
1716+
LPWSTR value;
1717+
Py_ssize_t size;
1718+
Py_ssize_t length;
1719+
PyObject *result = NULL;
16921720

16931721
if (type != REG_SZ) {
16941722
PyErr_SetString(PyExc_TypeError, "type must be winreg.REG_SZ");
16951723
return NULL;
16961724
}
16971725

1698-
wchar_t *value = PyUnicode_AsWideCharString(value_obj, &value_length);
1726+
value = PyUnicode_AsWideCharString(value_obj, &length);
16991727
if (value == NULL) {
17001728
return NULL;
17011729
}
1702-
if ((Py_ssize_t)(DWORD)value_length != value_length) {
1730+
1731+
size = (length + 1) * sizeof(WCHAR);
1732+
if ((Py_ssize_t)(DWORD)size != size) {
17031733
PyErr_SetString(PyExc_OverflowError, "value is too long");
1704-
PyMem_Free(value);
1705-
return NULL;
1734+
goto exit;
17061735
}
17071736

17081737
if (PySys_Audit("winreg.SetValue", "nunu#",
17091738
(Py_ssize_t)key, sub_key, (Py_ssize_t)type,
1710-
value, value_length) < 0) {
1711-
PyMem_Free(value);
1712-
return NULL;
1739+
value, length) < 0)
1740+
{
1741+
goto exit;
1742+
}
1743+
1744+
if (key == HKEY_PERFORMANCE_DATA) {
1745+
PyErr_SetFromWindowsErrWithFunction(ERROR_INVALID_HANDLE,
1746+
"RegSetValue");
1747+
goto exit;
1748+
}
1749+
1750+
if (sub_key && sub_key[0]) {
1751+
Py_BEGIN_ALLOW_THREADS
1752+
rc = RegCreateKeyExW(key, sub_key, 0, NULL, 0, KEY_SET_VALUE, NULL,
1753+
&childKey, NULL);
1754+
Py_END_ALLOW_THREADS
1755+
if (rc != ERROR_SUCCESS) {
1756+
PyErr_SetFromWindowsErrWithFunction(rc, "RegCreateKeyEx");
1757+
goto exit;
1758+
}
17131759
}
17141760

17151761
Py_BEGIN_ALLOW_THREADS
1716-
rc = RegSetValueW(key, sub_key, REG_SZ, value, (DWORD)(value_length + 1));
1762+
rc = RegSetValueExW(childKey, NULL, 0, REG_SZ, (LPBYTE)value, (DWORD)size);
17171763
Py_END_ALLOW_THREADS
1764+
if (rc == ERROR_SUCCESS) {
1765+
result = Py_NewRef(Py_None);
1766+
}
1767+
else {
1768+
PyErr_SetFromWindowsErrWithFunction(rc, "RegSetValueEx");
1769+
}
1770+
1771+
exit:
17181772
PyMem_Free(value);
1719-
if (rc != ERROR_SUCCESS)
1720-
return PyErr_SetFromWindowsErrWithFunction(rc, "RegSetValue");
1721-
Py_RETURN_NONE;
1773+
if (childKey != key) {
1774+
RegCloseKey(childKey);
1775+
}
1776+
return result;
17221777
}
17231778

17241779
/*[clinic input]
@@ -1771,32 +1826,39 @@ winreg_SetValueEx_impl(PyObject *module, HKEY key,
17711826
DWORD type, PyObject *value)
17721827
/*[clinic end generated code: output=811b769a66ae11b7 input=900a9e3990bfb196]*/
17731828
{
1774-
BYTE *data;
1775-
DWORD len;
1776-
17771829
LONG rc;
1830+
BYTE *data = NULL;
1831+
DWORD size;
1832+
PyObject *result = NULL;
17781833

1779-
if (!Py2Reg(value, type, &data, &len))
1834+
if (!Py2Reg(value, type, &data, &size))
17801835
{
1781-
if (!PyErr_Occurred())
1836+
if (!PyErr_Occurred()) {
17821837
PyErr_SetString(PyExc_ValueError,
17831838
"Could not convert the data to the specified type.");
1839+
}
17841840
return NULL;
17851841
}
17861842
if (PySys_Audit("winreg.SetValue", "nunO",
17871843
(Py_ssize_t)key, value_name, (Py_ssize_t)type,
1788-
value) < 0) {
1789-
PyMem_Free(data);
1790-
return NULL;
1844+
value) < 0)
1845+
{
1846+
goto exit;
17911847
}
1848+
17921849
Py_BEGIN_ALLOW_THREADS
1793-
rc = RegSetValueExW(key, value_name, 0, type, data, len);
1850+
rc = RegSetValueExW(key, value_name, 0, type, data, size);
17941851
Py_END_ALLOW_THREADS
1852+
if (rc == ERROR_SUCCESS) {
1853+
result = Py_NewRef(Py_None);
1854+
}
1855+
else {
1856+
PyErr_SetFromWindowsErrWithFunction(rc, "RegSetValueEx");
1857+
}
1858+
1859+
exit:
17951860
PyMem_Free(data);
1796-
if (rc != ERROR_SUCCESS)
1797-
return PyErr_SetFromWindowsErrWithFunction(rc,
1798-
"RegSetValueEx");
1799-
Py_RETURN_NONE;
1861+
return result;
18001862
}
18011863

18021864
/*[clinic input]

0 commit comments

Comments
 (0)