bpo-36085: Enable better DLL resolution on Windows (GH-12302)

4 files changed, 159 insertions, 1 deletions

diff --git a/Lib/ctypes/__init__.py b/Lib/ctypes/__init__.py
index 5f78bed..4107db3 100644
--- a/Lib/ctypes/__init__.py
+++ b/Lib/ctypes/__init__.py
@@ -326,7 +326,8 @@ class CDLL(object):
 
     def __init__(self, name, mode=DEFAULT_MODE, handle=None,
                  use_errno=False,
-                 use_last_error=False):
+                 use_last_error=False,
+                 winmode=None):
         self._name = name
         flags = self._func_flags_
         if use_errno:
@@ -341,6 +342,15 @@ class CDLL(object):
             """
             if name and name.endswith(")") and ".a(" in name:
                 mode |= ( _os.RTLD_MEMBER | _os.RTLD_NOW )
+        if _os.name == "nt":
+            if winmode is not None:
+                mode = winmode
+            else:
+                import nt
+                mode = nt._LOAD_LIBRARY_SEARCH_DEFAULT_DIRS
+                if '/' in name or '\\' in name:
+                    self._name = nt._getfullpathname(self._name)
+                    mode |= nt._LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR
 
         class _FuncPtr(_CFuncPtr):
             _flags_ = flags
diff --git a/Lib/ctypes/test/test_loading.py b/Lib/ctypes/test/test_loading.py
index f3b65b9..be367c6 100644
--- a/Lib/ctypes/test/test_loading.py
+++ b/Lib/ctypes/test/test_loading.py
@@ -1,6 +1,9 @@
 from ctypes import *
 import os
+import shutil
+import subprocess
 import sys
+import sysconfig
 import unittest
 import test.support
 from ctypes.util import find_library
@@ -112,5 +115,65 @@ class LoaderTest(unittest.TestCase):
         # This is the real test: call the function via 'call_function'
         self.assertEqual(0, call_function(proc, (None,)))
 
+    @unittest.skipUnless(os.name == "nt",
+                         'test specific to Windows')
+    def test_load_dll_with_flags(self):
+        _sqlite3 = test.support.import_module("_sqlite3")
+        src = _sqlite3.__file__
+        if src.lower().endswith("_d.pyd"):
+            ext = "_d.dll"
+        else:
+            ext = ".dll"
+
+        with test.support.temp_dir() as tmp:
+            # We copy two files and load _sqlite3.dll (formerly .pyd),
+            # which has a dependency on sqlite3.dll. Then we test
+            # loading it in subprocesses to avoid it starting in memory
+            # for each test.
+            target = os.path.join(tmp, "_sqlite3.dll")
+            shutil.copy(src, target)
+            shutil.copy(os.path.join(os.path.dirname(src), "sqlite3" + ext),
+                        os.path.join(tmp, "sqlite3" + ext))
+
+            def should_pass(command):
+                with self.subTest(command):
+                    subprocess.check_output(
+                        [sys.executable, "-c",
+                         "from ctypes import *; import nt;" + command],
+                        cwd=tmp
+                    )
+
+            def should_fail(command):
+                with self.subTest(command):
+                    with self.assertRaises(subprocess.CalledProcessError):
+                        subprocess.check_output(
+                            [sys.executable, "-c",
+                             "from ctypes import *; import nt;" + command],
+                            cwd=tmp, stderr=subprocess.STDOUT,
+                        )
+
+            # Default load should not find this in CWD
+            should_fail("WinDLL('_sqlite3.dll')")
+
+            # Relative path (but not just filename) should succeed
+            should_pass("WinDLL('./_sqlite3.dll')")
+
+            # Insecure load flags should succeed
+            should_pass("WinDLL('_sqlite3.dll', winmode=0)")
+
+            # Full path load without DLL_LOAD_DIR shouldn't find dependency
+            should_fail("WinDLL(nt._getfullpathname('_sqlite3.dll'), " +
+                        "winmode=nt._LOAD_LIBRARY_SEARCH_SYSTEM32)")
+
+            # Full path load with DLL_LOAD_DIR should succeed
+            should_pass("WinDLL(nt._getfullpathname('_sqlite3.dll'), " +
+                        "winmode=nt._LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR)")
+
+            # User-specified directory should succeed
+            should_pass("import os; p = os.add_dll_directory(os.getcwd());" +
+                        "WinDLL('_sqlite3.dll'); p.close()")
+
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/Lib/os.py b/Lib/os.py
index 7741c75..79ff7a2 100644
--- a/Lib/os.py
+++ b/Lib/os.py
@@ -1070,3 +1070,40 @@ class PathLike(abc.ABC):
     @classmethod
     def __subclasshook__(cls, subclass):
         return hasattr(subclass, '__fspath__')
+
+
+if name == 'nt':
+    class _AddedDllDirectory:
+        def __init__(self, path, cookie, remove_dll_directory):
+            self.path = path
+            self._cookie = cookie
+            self._remove_dll_directory = remove_dll_directory
+        def close(self):
+            self._remove_dll_directory(self._cookie)
+            self.path = None
+        def __enter__(self):
+            return self
+        def __exit__(self, *args):
+            self.close()
+        def __repr__(self):
+            if self.path:
+                return "<AddedDllDirectory({!r})>".format(self.path)
+            return "<AddedDllDirectory()>"
+
+    def add_dll_directory(path):
+        """Add a path to the DLL search path.
+
+        This search path is used when resolving dependencies for imported
+        extension modules (the module itself is resolved through sys.path),
+        and also by ctypes.
+
+        Remove the directory by calling close() on the returned object or
+        using it in a with statement.
+        """
+        import nt
+        cookie = nt._add_dll_directory(path)
+        return _AddedDllDirectory(
+            path,
+            cookie,
+            nt._remove_dll_directory
+        )
diff --git a/Lib/test/test_import/__init__.py b/Lib/test/test_import/__init__.py
index 7306e0f..a0bfe1a 100644
--- a/Lib/test/test_import/__init__.py
+++ b/Lib/test/test_import/__init__.py
@@ -8,6 +8,8 @@ import os
 import platform
 import py_compile
 import random
+import shutil
+import subprocess
 import stat
 import sys
 import threading
@@ -17,6 +19,7 @@ import unittest.mock as mock
 import textwrap
 import errno
 import contextlib
+import glob
 
 import test.support
 from test.support import (
@@ -460,6 +463,51 @@ class ImportTests(unittest.TestCase):
         finally:
             del sys.path[0]
 
+    @unittest.skipUnless(sys.platform == "win32", "Windows-specific")
+    def test_dll_dependency_import(self):
+        from _winapi import GetModuleFileName
+        dllname = GetModuleFileName(sys.dllhandle)
+        pydname = importlib.util.find_spec("_sqlite3").origin
+        depname = os.path.join(
+            os.path.dirname(pydname),
+            "sqlite3{}.dll".format("_d" if "_d" in pydname else ""))
+
+        with test.support.temp_dir() as tmp:
+            tmp2 = os.path.join(tmp, "DLLs")
+            os.mkdir(tmp2)
+
+            pyexe = os.path.join(tmp, os.path.basename(sys.executable))
+            shutil.copy(sys.executable, pyexe)
+            shutil.copy(dllname, tmp)
+            for f in glob.glob(os.path.join(sys.prefix, "vcruntime*.dll")):
+                shutil.copy(f, tmp)
+
+            shutil.copy(pydname, tmp2)
+
+            env = None
+            env = {k.upper(): os.environ[k] for k in os.environ}
+            env["PYTHONPATH"] = tmp2 + ";" + os.path.dirname(os.__file__)
+
+            # Test 1: import with added DLL directory
+            subprocess.check_call([
+                pyexe, "-Sc", ";".join([
+                    "import os",
+                    "p = os.add_dll_directory({!r})".format(
+                        os.path.dirname(depname)),
+                    "import _sqlite3",
+                    "p.close"
+                ])],
+                stderr=subprocess.STDOUT,
+                env=env,
+                cwd=os.path.dirname(pyexe))
+
+            # Test 2: import with DLL adjacent to PYD
+            shutil.copy(depname, tmp2)
+            subprocess.check_call([pyexe, "-Sc", "import _sqlite3"],
+                                    stderr=subprocess.STDOUT,
+                                    env=env,
+                                    cwd=os.path.dirname(pyexe))
+
 
 @skip_if_dont_write_bytecode
 class FilePermissionTests(unittest.TestCase):