Fix #2810 - handle the case where some registry calls return

ERROR_MORE_DATA, requiring another call to get the remaining data.

Patch by Daniel Stutzbach
This commit is contained in:
Brian Curtin 2010-05-25 15:06:15 +00:00
parent 264acad721
commit b64c89bd7a
3 changed files with 159 additions and 33 deletions

View file

@ -1153,7 +1153,14 @@ PyEnumKey(PyObject *self, PyObject *args)
int index;
long rc;
PyObject *retStr;
char tmpbuf[256]; /* max key name length is 255 */
/* The Windows docs claim that the max key name length is 255
* characters, plus a terminating nul character. However,
* empirical testing demonstrates that it is possible to
* create a 256 character key that is missing the terminating
* nul. RegEnumKeyEx requires a 257 character buffer to
* retrieve such a key name. */
char tmpbuf[257];
DWORD len = sizeof(tmpbuf); /* includes NULL terminator */
if (!PyArg_ParseTuple(args, "Oi:EnumKey", &obKey, &index))
@ -1180,8 +1187,8 @@ PyEnumValue(PyObject *self, PyObject *args)
long rc;
char *retValueBuf;
char *retDataBuf;
DWORD retValueSize;
DWORD retDataSize;
DWORD retValueSize, bufValueSize;
DWORD retDataSize, bufDataSize;
DWORD typ;
PyObject *obData;
PyObject *retVal;
@ -1199,6 +1206,8 @@ PyEnumValue(PyObject *self, PyObject *args)
"RegQueryInfoKey");
++retValueSize; /* include null terminators */
++retDataSize;
bufDataSize = retDataSize;
bufValueSize = retValueSize;
retValueBuf = (char *)PyMem_Malloc(retValueSize);
if (retValueBuf == NULL)
return PyErr_NoMemory();
@ -1208,16 +1217,33 @@ PyEnumValue(PyObject *self, PyObject *args)
return PyErr_NoMemory();
}
Py_BEGIN_ALLOW_THREADS
rc = RegEnumValue(hKey,
index,
retValueBuf,
&retValueSize,
NULL,
&typ,
(BYTE *)retDataBuf,
&retDataSize);
Py_END_ALLOW_THREADS
while (1) {
char *tmp;
Py_BEGIN_ALLOW_THREADS
rc = RegEnumValue(hKey,
index,
retValueBuf,
&retValueSize,
NULL,
&typ,
(BYTE *)retDataBuf,
&retDataSize);
Py_END_ALLOW_THREADS
if (rc != ERROR_MORE_DATA)
break;
bufDataSize *= 2;
tmp = (char *)PyMem_Realloc(retDataBuf, bufDataSize);
if (tmp == NULL) {
PyErr_NoMemory();
retVal = NULL;
goto fail;
}
retDataBuf = tmp;
retDataSize = bufDataSize;
retValueSize = bufValueSize;
}
if (rc != ERROR_SUCCESS) {
retVal = PyErr_SetFromWindowsErrWithFunction(rc,
@ -1374,28 +1400,56 @@ PyQueryValue(PyObject *self, PyObject *args)
long rc;
PyObject *retStr;
char *retBuf;
long bufSize = 0;
DWORD bufSize = 0;
DWORD retSize = 0;
char *tmp;
if (!PyArg_ParseTuple(args, "Oz:QueryValue", &obKey, &subKey))
return NULL;
if (!PyHKEY_AsHKEY(obKey, &hKey, FALSE))
return NULL;
if ((rc = RegQueryValue(hKey, subKey, NULL, &bufSize))
!= ERROR_SUCCESS)
rc = RegQueryValue(hKey, subKey, NULL, &retSize);
if (rc == ERROR_MORE_DATA)
retSize = 256;
else if (rc != ERROR_SUCCESS)
return PyErr_SetFromWindowsErrWithFunction(rc,
"RegQueryValue");
retStr = PyString_FromStringAndSize(NULL, bufSize);
if (retStr == NULL)
return NULL;
retBuf = PyString_AS_STRING(retStr);
if ((rc = RegQueryValue(hKey, subKey, retBuf, &bufSize))
!= ERROR_SUCCESS) {
Py_DECREF(retStr);
bufSize = retSize;
retBuf = (char *) PyMem_Malloc(bufSize);
if (retBuf == NULL)
return PyErr_NoMemory();
while (1) {
retSize = bufSize;
rc = RegQueryValue(hKey, subKey, retBuf, &retSize);
if (rc != ERROR_MORE_DATA)
break;
bufSize *= 2;
tmp = (char *) PyMem_Realloc(retBuf, bufSize);
if (tmp == NULL) {
PyMem_Free(retBuf);
return PyErr_NoMemory();
}
retBuf = tmp;
}
if (rc != ERROR_SUCCESS) {
PyMem_Free(retBuf);
return PyErr_SetFromWindowsErrWithFunction(rc,
"RegQueryValue");
}
_PyString_Resize(&retStr, strlen(retBuf));
if (retBuf[retSize-1] == '\x00')
retSize--;
retStr = PyString_FromStringAndSize(retBuf, retSize);
if (retStr == NULL) {
PyMem_Free(retBuf);
return NULL;
}
return retStr;
}
@ -1407,8 +1461,8 @@ PyQueryValueEx(PyObject *self, PyObject *args)
char *valueName;
long rc;
char *retBuf;
DWORD bufSize = 0;
char *retBuf, *tmp;
DWORD bufSize = 0, retSize;
DWORD typ;
PyObject *obData;
PyObject *result;
@ -1418,18 +1472,34 @@ PyQueryValueEx(PyObject *self, PyObject *args)
if (!PyHKEY_AsHKEY(obKey, &hKey, FALSE))
return NULL;
if ((rc = RegQueryValueEx(hKey, valueName,
NULL, NULL, NULL,
&bufSize))
!= ERROR_SUCCESS)
rc = RegQueryValueEx(hKey, valueName, NULL, NULL, NULL, &bufSize);
if (rc == ERROR_MORE_DATA)
bufSize = 256;
else if (rc != ERROR_SUCCESS)
return PyErr_SetFromWindowsErrWithFunction(rc,
"RegQueryValueEx");
retBuf = (char *)PyMem_Malloc(bufSize);
if (retBuf == NULL)
return PyErr_NoMemory();
if ((rc = RegQueryValueEx(hKey, valueName, NULL,
&typ, (BYTE *)retBuf, &bufSize))
!= ERROR_SUCCESS) {
while (1) {
retSize = bufSize;
rc = RegQueryValueEx(hKey, valueName, NULL, &typ,
(BYTE *)retBuf, &retSize);
if (rc != ERROR_MORE_DATA)
break;
bufSize *= 2;
tmp = (char *) PyMem_Realloc(retBuf, bufSize);
if (tmp == NULL) {
PyMem_Free(retBuf);
return PyErr_NoMemory();
}
retBuf = tmp;
}
if (rc != ERROR_SUCCESS) {
PyMem_Free(retBuf);
return PyErr_SetFromWindowsErrWithFunction(rc,
"RegQueryValueEx");