merge, document py2.7-winreg which backports winreg changes from py3.7

4 files changed, 169 insertions, 16 deletions

diff --git a/pypy/doc/whatsnew-head.rst b/pypy/doc/whatsnew-head.rst
index 2ea03cbe7d..f8adb5e3c2 100644
--- a/pypy/doc/whatsnew-head.rst
+++ b/pypy/doc/whatsnew-head.rst
@@ -44,3 +44,8 @@ Refactor documentation of win64 from future plans to what was executed
 Backport msvc detection from python3, which probably breaks using Visual Studio
 2008 (MSVC9, or the version that used to be used to build CPython2.7 on
 Windows)
+
+.. branch: py2.7-winreg
+
+Backport fixes to winreg adding reflection and fix for passing None (bpo
+21151).
diff --git a/pypy/module/_winreg/interp_winreg.py b/pypy/module/_winreg/interp_winreg.py
index c7af44d6da..dc29cd844c 100644
--- a/pypy/module/_winreg/interp_winreg.py
+++ b/pypy/module/_winreg/interp_winreg.py
@@ -6,6 +6,61 @@ from pypy.interpreter.error import OperationError, oefmt, wrap_windowserror
 from rpython.rtyper.lltypesystem import rffi, lltype
 from rpython.rlib import rwinreg, rwin32
 from rpython.rlib.rarithmetic import r_uint, intmask
+from rpython.translator.tool.cbuild import ExternalCompilationInfo
+
+
+# wrappers needed to call the reflection functions loaded at runtime
+# using WINAPI convention
+eci = ExternalCompilationInfo(
+    includes=['windows.h'],
+    post_include_bits=[
+        "RPY_EXTERN LONG\n"
+        "pypy_RegChangeReflectionKey(FARPROC address, HKEY key);\n"
+        "RPY_EXTERN LONG\n"
+        "pypy_RegQueryReflectionKey(FARPROC address, HKEY key, LPBOOL isDisabled);\n"
+        "RPY_EXTERN LONG\n"
+        "pypy_RegDeleteKeyExA(FARPROC address, HKEY key, LPCSTR subkey,\n"
+        "                    REGSAM sam, DWORD reserved);\n"
+    ],
+    separate_module_sources=['''
+        LONG
+        pypy_RegChangeReflectionKey(FARPROC address, HKEY key) {
+            LONG (WINAPI *func)(HKEY);
+            *(FARPROC*)&func = address;
+            return func(key);
+        }
+
+        LONG
+        pypy_RegQueryReflectionKey(FARPROC address, HKEY key, LPBOOL isDisabled) {
+            LONG (WINAPI *func)(HKEY, LPBOOL);
+            *(FARPROC*)&func = address;
+            return func(key, isDisabled);
+        }
+
+        LONG
+        pypy_RegDeleteKeyExA(FARPROC address, HKEY key, LPCSTR subkey,
+                            REGSAM sam, DWORD reserved) {
+            LONG (WINAPI *func)(HKEY, LPCSTR, REGSAM, DWORD);
+            *(FARPROC*)&func = address;
+            return func(key, subkey, sam, reserved);
+        }
+    '''],
+)
+pypy_RegChangeReflectionKey = rffi.llexternal(
+    'pypy_RegChangeReflectionKey',
+    [rffi.VOIDP, rwinreg.HKEY],
+    rffi.LONG, compilation_info=eci)
+
+pypy_RegQueryReflectionKey = rffi.llexternal(
+    'pypy_RegQueryReflectionKey',
+    [rffi.VOIDP, rwinreg.HKEY, rwin32.LPBOOL],
+    rffi.LONG, compilation_info=eci)
+
+pypy_RegDeleteKeyExA = rffi.llexternal(
+    'pypy_RegDeleteKeyExA',
+    [rffi.VOIDP, rwinreg.HKEY, rffi.CCHARP, rwinreg.REGSAM, rwin32.DWORD],
+    rffi.LONG, compilation_info=eci)
+
 
 def raiseWindowsError(space, errcode, context):
     message = rwin32.FormatError(errcode)
@@ -258,6 +313,8 @@ But the underlying API call doesn't return the type, Lame Lame Lame, DONT USE TH
                 if ret != 0:
                     raiseWindowsError(space, ret, 'RegQueryValue')
                 length = intmask(bufsize_p[0])
+                if length == 0:
+                    return space.w_None
                 return space.newtext(rffi.charp2strn(buf, length - 1))
 
 def convert_to_regdata(space, w_value, typ):
@@ -328,8 +385,7 @@ def convert_to_regdata(space, w_value, typ):
     else: # REG_BINARY and ALL unknown data types.
         if space.is_w(w_value, space.w_None):
             buflen = 0
-            buf = lltype.malloc(rffi.CCHARP.TO, 1, flavor='raw')
-            buf[0] = '\0'
+            buf = lltype.nullptr(rffi.CCHARP.TO)
         else:
             try:
                 value = w_value.readbuf_w(space)
@@ -385,7 +441,10 @@ def convert_from_regdata(space, buf, buflen, typ):
         return space.newlist(l)
 
     else: # REG_BINARY and all other types
-        return space.newbytes(rffi.charpsize2str(buf, buflen))
+        if buflen == 0:
+            return space.w_None
+        else:
+            return space.newbytes(rffi.charpsize2str(buf, buflen))
 
 @unwrap_spec(value_name="text", typ=int)
 def SetValueEx(space, w_hkey, value_name, w_reserved, typ, w_value):
@@ -424,7 +483,8 @@ the configuration registry.  This helps the registry perform efficiently."""
     try:
         ret = rwinreg.RegSetValueExA(hkey, value_name, 0, typ, buf, buflen)
     finally:
-        lltype.free(buf, flavor='raw')
+        if buf != lltype.nullptr(rffi.CCHARP.TO):
+            lltype.free(buf, flavor='raw')
     if ret != 0:
         raiseWindowsError(space, ret, 'RegSetValueEx')
 
@@ -707,30 +767,84 @@ def ExpandEnvironmentStrings(space, w_source):
     except WindowsError as e:
         raise wrap_windowserror(space, e)
 
+
+class ReflectionFunction(object):
+    def __init__(self, name, stdcall_wrapper):
+        self.name = name
+        self.handle = lltype.nullptr(rffi.VOIDP.TO)
+        self.wrapper = stdcall_wrapper
+
+    def check(self):
+        if self.handle != lltype.nullptr(rffi.VOIDP.TO):
+            return True
+        from rpython.rlib.rdynload import GetModuleHandle, dlsym
+        lib = GetModuleHandle("advapi32.dll")
+        try:
+            handle = dlsym(lib, self.name)
+        except KeyError:
+            return False
+        self.handle = handle
+        return True
+
+    def call(self, *args):
+        assert self.handle != lltype.nullptr(rffi.VOIDP.TO)
+        return self.wrapper(self.handle, *args)
+
+
+_RegDisableReflectionKey = ReflectionFunction(
+    "RegDisableReflectionKey", pypy_RegChangeReflectionKey)
+_RegEnableReflectionKey = ReflectionFunction(
+    "RegEnableReflectionKey", pypy_RegChangeReflectionKey)
+_RegQueryReflectionKey = ReflectionFunction(
+    "RegQueryReflectionKey", pypy_RegQueryReflectionKey)
+_RegDeleteKeyExA = ReflectionFunction("RegDeleteKeyExA", pypy_RegDeleteKeyExA)
+
+
 def DisableReflectionKey(space, w_key):
     """Disables registry reflection for 32-bit processes running on a 64-bit
     Operating System.  Will generally raise NotImplemented if executed on
     a 32-bit Operating System.
     If the key is not on the reflection list, the function succeeds but has no effect.
     Disabling reflection for a key does not affect reflection of any subkeys."""
-    raise oefmt(space.w_NotImplementedError,
-                "not implemented on this platform")
+    if not _RegDisableReflectionKey.check():
+        raise oefmt(space.w_NotImplementedError,
+                    "not implemented on this platform")
+    else:
+        hkey = hkey_w(w_key, space)
+        ret = _RegDisableReflectionKey.call(hkey)
+        if ret != 0:
+            raiseWindowsError(space, ret, 'RegDisableReflectionKey')
 
 def EnableReflectionKey(space, w_key):
     """Restores registry reflection for the specified disabled key.
     Will generally raise NotImplemented if executed on a 32-bit Operating System.
     Restoring reflection for a key does not affect reflection of any subkeys."""
-    raise oefmt(space.w_NotImplementedError,
-                "not implemented on this platform")
+    if not _RegEnableReflectionKey.check():
+        raise oefmt(space.w_NotImplementedError,
+                    "not implemented on this platform")
+    else:
+        hkey = hkey_w(w_key, space)
+        ret = _RegEnableReflectionKey.call(hkey)
+        if ret != 0:
+            raiseWindowsError(space, ret, 'RegEnableReflectionKey')
 
 def QueryReflectionKey(space, w_key):
     """bool = QueryReflectionKey(hkey) - Determines the reflection state for the specified key.
     Will generally raise NotImplemented if executed on a 32-bit Operating System."""
-    raise oefmt(space.w_NotImplementedError,
-                "not implemented on this platform")
+    if not _RegQueryReflectionKey.check():
+        raise oefmt(space.w_NotImplementedError,
+                    "not implemented on this platform")
+    else:
+        hkey = hkey_w(w_key, space)
+        with lltype.scoped_alloc(rwin32.LPBOOL.TO, 1) as isDisabled:
+            ret = _RegQueryReflectionKey.call(hkey, isDisabled)
+            if ret != 0:
+                raiseWindowsError(space, ret, 'RegQueryReflectionKey')
+            return space.newbool(intmask(isDisabled[0]) != 0)
 
-@unwrap_spec(subkey="text")
-def DeleteKeyEx(space, w_key, subkey):
+
+@unwrap_spec(sub_key="text", access=r_uint, reserved=int)
+def DeleteKeyEx(space, w_key, sub_key, access=rwinreg.KEY_WOW64_64KEY, reserved=0):
     """DeleteKeyEx(key, sub_key, sam, res) - Deletes the specified key.
 
     key is an already open key, or any one of the predefined HKEY_* constants.
@@ -744,5 +858,11 @@ def DeleteKeyEx(space, w_key, subkey):
     If the method succeeds, the entire key, including all of its values,
     is removed.  If the method fails, a WindowsError exception is raised.
     On unsupported Windows versions, NotImplementedError is raised."""
-    raise oefmt(space.w_NotImplementedError,
-                "not implemented on this platform")
+    if not _RegDeleteKeyExA.check():
+        raise oefmt(space.w_NotImplementedError,
+                    "not implemented on this platform")
+    else:
+        hkey = hkey_w(w_key, space)
+        ret = _RegDeleteKeyExA.call(hkey, sub_key, access, reserved)
+        if ret != 0:
+            raiseWindowsError(space, ret, 'RegDeleteKeyEx')
diff --git a/pypy/module/_winreg/moduledef.py b/pypy/module/_winreg/moduledef.py
index d865ca691e..b20bd15d62 100644
--- a/pypy/module/_winreg/moduledef.py
+++ b/pypy/module/_winreg/moduledef.py
@@ -72,3 +72,8 @@ to see what constants are used, and where."""
 
     for name, value in constants.iteritems():
         interpleveldefs[name] = "space.wrap(%s)" % (value,)
+
+    import pypy.module.sys.version
+    if pypy.module.sys.version.CPYTHON_VERSION < (3, 6):
+        del interpleveldefs["REG_QWORD"]
+        del interpleveldefs["REG_QWORD_LITTLE_ENDIAN"]
diff --git a/pypy/module/_winreg/test/test_winreg.py b/pypy/module/_winreg/test/test_winreg.py
index ae2d8ee227..1d25cbb900 100644
--- a/pypy/module/_winreg/test/test_winreg.py
+++ b/pypy/module/_winreg/test/test_winreg.py
@@ -31,6 +31,7 @@ class AppTestFfi:
 
     def setup_class(cls):
         import _winreg
+        from platform import machine
         space = cls.space
         cls.root_key = _winreg.HKEY_CURRENT_USER
         cls.test_key_name = "SOFTWARE\\Pypy Registry Test Key - Delete Me"
@@ -38,6 +39,7 @@ class AppTestFfi:
         cls.w_test_key_name = space.wrap(cls.test_key_name)
         cls.w_canSaveKey = space.wrap(canSaveKey)
         cls.w_tmpfilename = space.wrap(str(udir.join('winreg-temp')))
+        cls.w_win64_machine = space.wrap(machine() == "AMD64")
 
         test_data = [
             ("Int Value", 0xFEDCBA98, _winreg.REG_DWORD),
@@ -45,6 +47,7 @@ class AppTestFfi:
             ("Unicode Value", u"A unicode Value", _winreg.REG_SZ),
             ("Str Expand", "The path is %path%", _winreg.REG_EXPAND_SZ),
             ("Multi Str", ["Several", "string", u"values"], _winreg.REG_MULTI_SZ),
+            ("Raw None", None, _winreg.REG_BINARY),
             ("Raw data", "binary"+chr(0)+"data", _winreg.REG_BINARY),
             ]
         cls.w_test_data = space.wrap(test_data)
@@ -175,14 +178,19 @@ class AppTestFfi:
 
     def test_delete(self):
         # must be run after test_SetValueEx
-        from _winreg import OpenKey, KEY_ALL_ACCESS, DeleteValue, DeleteKey
+        from _winreg import OpenKey, KEY_ALL_ACCESS, DeleteValue, DeleteKey, DeleteKeyEx
         key = OpenKey(self.root_key, self.test_key_name, 0, KEY_ALL_ACCESS)
         sub_key = OpenKey(key, "sub_key", 0, KEY_ALL_ACCESS)
 
         for name, value, type in self.test_data:
             DeleteValue(sub_key, name)
 
-        DeleteKey(key, "sub_key")
+        if self.win64_machine:
+            DeleteKeyEx(key, "sub_key", KEY_ALL_ACCESS, 0)
+        else:
+            DeleteKey(key, "sub_key")
+
+        raises(OSError, OpenKey, key, "sub_key")
 
     def test_connect(self):
         from _winreg import ConnectRegistry, HKEY_LOCAL_MACHINE
@@ -255,3 +263,18 @@ class AppTestFfi:
             raises(NotImplementedError, DeleteKeyEx, self.root_key,
                    self.test_key_name)
 
+    def test_reflection(self):
+        import sys
+        from _winreg import DisableReflectionKey, EnableReflectionKey, \
+                           QueryReflectionKey, OpenKey, HKEY_LOCAL_MACHINE
+        # Adapted from lib-python test
+        if not self.win64_machine:
+            skip("Requires 64-bit host")
+        # Test that we can call the query, enable, and disable functions
+        # on a key which isn't on the reflection list with no consequences.
+        with OpenKey(HKEY_LOCAL_MACHINE, "Software") as key:
+            # HKLM\Software is redirected but not reflected in all OSes
+            assert QueryReflectionKey(key)
+            assert EnableReflectionKey(key) is None
+            assert DisableReflectionKey(key) is None
+            assert QueryReflectionKey(key)