Merged revisions 81517 via svnmerge from

svn+ssh://pythondev@svn.python.org/python/trunk

........
  r81517 | brian.curtin | 2010-05-25 10:06:15 -0500 (Tue, 25 May 2010) | 5 lines

  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
........

diff --git a/Lib/test/test_winreg.py b/Lib/test/test_winreg.py
index f2ec2a8c0f89c76517b3b3fefaf5f96776fe732d..902309a558e388caeba98e49291bb73a1e338739 100644 (file)
--- a/Lib/test/test_winreg.py
+++ b/Lib/test/test_winreg.py
@@ -6,6 +6,8 @@ import os, sys
 import unittest
 
 from test import test_support
+threading = test_support.import_module("threading")
+from platform import machine
 
 test_key_name = "SOFTWARE\\Python Registry Test Key - Delete Me"
 
@@ -166,6 +168,58 @@ class WinregTests(unittest.TestCase):
         self.assertEqual(type(r), unicode)
         self.assertEqual(r, os.environ["windir"] + "\\test")
 
+    def test_changing_value(self):
+        # Issue2810: A race condition in 2.6 and 3.1 may cause
+        # EnumValue or QueryValue to throw "WindowsError: More data is
+        # available"
+        done = False
+
+        class VeryActiveThread(threading.Thread):
+            def run(self):
+                with CreateKey(HKEY_CURRENT_USER, test_key_name) as key:
+                    use_short = True
+                    long_string = 'x'*2000
+                    while not done:
+                        s = 'x' if use_short else long_string
+                        use_short = not use_short
+                        SetValue(key, 'changing_value', REG_SZ, s)
+
+        thread = VeryActiveThread()
+        thread.start()
+        try:
+            with CreateKey(HKEY_CURRENT_USER,
+                           test_key_name+'\\changing_value') as key:
+                for _ in range(1000):
+                    num_subkeys, num_values, t = QueryInfoKey(key)
+                    for i in range(num_values):
+                        name = EnumValue(key, i)
+                        QueryValue(key, name[0])
+        finally:
+            done = True
+            thread.join()
+            DeleteKey(HKEY_CURRENT_USER, test_key_name+'\\changing_value')
+            DeleteKey(HKEY_CURRENT_USER, test_key_name)
+
+    def test_long_key(self):
+        # Issue2810, in 2.6 and 3.1 when the key name was exactly 256
+        # characters, EnumKey threw "WindowsError: More data is
+        # available"
+        name = 'x'*256
+        try:
+            with CreateKey(HKEY_CURRENT_USER, test_key_name) as key:
+                SetValue(key, name, REG_SZ, 'x')
+                num_subkeys, num_values, t = QueryInfoKey(key)
+                EnumKey(key, 0)
+        finally:
+            DeleteKey(HKEY_CURRENT_USER, '\\'.join((test_key_name, name)))
+            DeleteKey(HKEY_CURRENT_USER, test_key_name)
+
+    def test_dynamic_key(self):
+        # Issue2810, when the value is dynamically generated, these
+        # throw "WindowsError: More data is available" in 2.6 and 3.1
+        EnumValue(HKEY_PERFORMANCE_DATA, 0)
+        QueryValueEx(HKEY_PERFORMANCE_DATA, None)
+
 def test_main():
     test_support.run_unittest(WinregTests)
 

diff --git a/Misc/NEWS b/Misc/NEWS
index 34a7c859bdff17008ac54def521751cb8bf8154e..770a7f1fc5baf466b2512c489554ecc167bf4321 100644 (file)
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -225,6 +225,12 @@ Library
 - Issue #7356: ctypes.util: Make parsing of ldconfig output independent of
   the locale.
 
+Extension Modules
+-----------------
+
+- Issue #2810: Fix cases where the Windows registry API returns
+  ERROR_MORE_DATA, requiring a re-try in order to get the complete result.
+
 Build
 -----
 

diff --git a/PC/_winreg.c b/PC/_winreg.c
index 1fb656bda88d604601e92cd88c048e4a36ef90ca..eac596e2f33cfa1cd9b1710400fe63dd06a9f9c8 100644 (file)
--- a/PC/_winreg.c
+++ b/PC/_winreg.c
@@ -1071,7 +1071,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))
@@ -1098,8 +1105,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;
@@ -1117,6 +1124,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();
@@ -1126,16 +1135,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,
@@ -1292,28 +1318,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;
 }
 
@@ -1325,8 +1379,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;
@@ -1336,18 +1390,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");