kernel32: Move Get/SetDllDirectory() implementation to ntdll.

dlls/kernel32/module.c

@@ -56,7 +56,6 @@ struct dll_dir_entry
 };
 
 static struct list dll_dir_list = LIST_INIT( dll_dir_list );  /* extra dirs from AddDllDirectory */
-static WCHAR *dll_directory;  /* extra path for SetDllDirectoryW */
 static DWORD default_search_flags;  /* default flags set by SetDefaultDllDirectories */
 
 /* to keep track of LOAD_LIBRARY_AS_DATAFILE_EXCLUSIVE file handles */
@@ -82,21 +81,40 @@ static CRITICAL_SECTION dlldir_section = { &critsect_debug, -1, 0, 0, 0, 0 };
  */
 DWORD WINAPI GetDllDirectoryA( DWORD buf_len, LPSTR buffer )
 {
+    UNICODE_STRING str;
+    NTSTATUS status;
+    WCHAR data[MAX_PATH];
     DWORD len;
 
-    RtlEnterCriticalSection( &dlldir_section );
-    len = dll_directory ? FILE_name_WtoA( dll_directory, strlenW(dll_directory), NULL, 0 ) : 0;
+    str.Buffer = data;
+    str.MaximumLength = sizeof(data);
+
+    for (;;)
+    {
+        status = LdrGetDllDirectory( &str );
+        if (status != STATUS_BUFFER_TOO_SMALL) break;
+        if (str.Buffer != data) HeapFree( GetProcessHeap(), 0, str.Buffer );
+        str.MaximumLength = str.Length;
+        if (!(str.Buffer = HeapAlloc( GetProcessHeap(), 0, str.MaximumLength )))
+        {
+            status = STATUS_NO_MEMORY;
+            break;
+        }
+    }
+
+    if (!set_ntstatus( status )) return 0;
+
+    len = FILE_name_WtoA( str.Buffer, str.Length / sizeof(WCHAR), NULL, 0 );
     if (buffer && buf_len > len)
     {
-        if (dll_directory) FILE_name_WtoA( dll_directory, -1, buffer, buf_len );
-        else *buffer = 0;
+        FILE_name_WtoA( str.Buffer, -1, buffer, buf_len );
     }
     else
     {
         len++;  /* for terminating null */
         if (buffer) *buffer = 0;
     }
-    RtlLeaveCriticalSection( &dlldir_section );
+    if (str.Buffer != data) HeapFree( GetProcessHeap(), 0, str.Buffer );
     return len;
 }
 
@@ -106,22 +124,15 @@ DWORD WINAPI GetDllDirectoryA( DWORD buf_len, LPSTR buffer )
  */
 DWORD WINAPI GetDllDirectoryW( DWORD buf_len, LPWSTR buffer )
 {
-    DWORD len;
+    UNICODE_STRING str;
+    NTSTATUS status;
 
-    RtlEnterCriticalSection( &dlldir_section );
-    len = dll_directory ? strlenW( dll_directory ) : 0;
-    if (buffer && buf_len > len)
-    {
-        if (dll_directory) memcpy( buffer, dll_directory, (len + 1) * sizeof(WCHAR) );
-        else *buffer = 0;
-    }
-    else
-    {
-        len++;  /* for terminating null */
-        if (buffer) *buffer = 0;
-    }
-    RtlLeaveCriticalSection( &dlldir_section );
-    return len;
+    str.Buffer = buffer;
+    str.MaximumLength = min( buf_len, UNICODE_STRING_MAX_CHARS ) * sizeof(WCHAR);
+    status = LdrGetDllDirectory( &str );
+    if (status == STATUS_BUFFER_TOO_SMALL) status = STATUS_SUCCESS;
+    if (!set_ntstatus( status )) return 0;
+    return str.Length / sizeof(WCHAR);
 }
 
 
@@ -145,24 +156,10 @@ BOOL WINAPI SetDllDirectoryA( LPCSTR dir )
  */
 BOOL WINAPI SetDllDirectoryW( LPCWSTR dir )
 {
-    WCHAR *newdir = NULL;
+    UNICODE_STRING str;
 
-    if (dir)
-    {
-        DWORD len = (strlenW(dir) + 1) * sizeof(WCHAR);
-        if (!(newdir = HeapAlloc( GetProcessHeap(), 0, len )))
-        {
-            SetLastError( ERROR_NOT_ENOUGH_MEMORY );
-            return FALSE;
-        }
-        memcpy( newdir, dir, len );
-    }
-
-    RtlEnterCriticalSection( &dlldir_section );
-    HeapFree( GetProcessHeap(), 0, dll_directory );
-    dll_directory = newdir;
-    RtlLeaveCriticalSection( &dlldir_section );
-    return TRUE;
+    RtlInitUnicodeString( &str, dir );
+    return set_ntstatus( LdrSetDllDirectory( &str ));
 }
 
 
@@ -508,7 +505,7 @@ WCHAR *MODULE_get_dll_load_path( LPCWSTR module, int safe_mode )
     const WCHAR *mod_end = NULL;
     UNICODE_STRING name, value;
     WCHAR *p, *ret;
-    int len = 0, path_len = 0;
+    int len = 0, path_len = 0, dlldir_len;
 
     /* adjust length for module name */
 
@@ -534,22 +531,27 @@ WCHAR *MODULE_get_dll_load_path( LPCWSTR module, int safe_mode )
     if (RtlQueryEnvironmentVariable_U( NULL, &name, &value ) == STATUS_BUFFER_TOO_SMALL)
         path_len = value.Length;
 
-    RtlEnterCriticalSection( &dlldir_section );
+    dlldir_len = GetDllDirectoryW( 0, NULL );
+
     if (safe_mode == -1) safe_mode = get_dll_safe_mode();
-    if (dll_directory) len += strlenW(dll_directory) + 1;
+    if (dlldir_len > 1) len += dlldir_len;
     else len += 2;  /* current directory */
     if ((p = ret = HeapAlloc( GetProcessHeap(), 0, path_len + len * sizeof(WCHAR) )))
     {
         if (module) p = append_path_len( p, module, mod_end - module );
 
-        if (dll_directory) p = append_path( p, dll_directory );
+        if (dlldir_len > 1)
+        {
+            GetDllDirectoryW( len - (p - ret), p );
+            p += strlenW(p);
+            *p++ = ';';
+        }
         else if (!safe_mode) p = append_path( p, dotW );
 
         p = append_path( p, system_path );
 
-        if (!dll_directory && safe_mode) p = append_path( p, dotW );
+        if (dlldir_len <= 1 && safe_mode) p = append_path( p, dotW );
     }
-    RtlLeaveCriticalSection( &dlldir_section );
     if (!ret) return NULL;
 
     value.Buffer = p;
@@ -616,7 +618,7 @@ static WCHAR *get_dll_load_path_search_flags( LPCWSTR module, DWORD flags )
     {
         LIST_FOR_EACH_ENTRY( dir, &dll_dir_list, struct dll_dir_entry, entry )
             len += strlenW( dir->dir ) + 1;
-        if (dll_directory) len += strlenW(dll_directory) + 1;
+        len += GetDllDirectoryW( 0, NULL );
     }
 
     if (flags & LOAD_LIBRARY_SEARCH_SYSTEM32) len += GetSystemDirectoryW( NULL, 0 );
@@ -629,7 +631,12 @@ static WCHAR *get_dll_load_path_search_flags( LPCWSTR module, DWORD flags )
         {
             LIST_FOR_EACH_ENTRY( dir, &dll_dir_list, struct dll_dir_entry, entry )
                 p = append_path( p, dir->dir );
-            if (dll_directory) p = append_path( p, dll_directory );
+            GetDllDirectoryW( ret + len - p, p );
+            if (*p)
+            {
+                p += strlenW(p);
+                *p++ = ';';
+            }
         }
         if (flags & LOAD_LIBRARY_SEARCH_SYSTEM32) GetSystemDirectoryW( p, ret + len - p );
         else

dlls/kernel32/tests/module.c

@@ -18,10 +18,15 @@
  * Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301, USA
  */
 
-#include "wine/test.h"
-#include <windows.h>
 #include <stdio.h>
+#include "ntstatus.h"
+#define WIN32_NO_STATUS
+#include "windef.h"
+#include "winbase.h"
+#include "winnls.h"
+#include "winternl.h"
 #include <psapi.h>
+#include "wine/test.h"
 
 static DWORD (WINAPI *pGetDllDirectoryA)(DWORD,LPSTR);
 static DWORD (WINAPI *pGetDllDirectoryW)(DWORD,LPWSTR);
@@ -32,6 +37,9 @@ static BOOL (WINAPI *pSetDefaultDllDirectories)(DWORD);
 static BOOL (WINAPI *pK32GetModuleInformation)(HANDLE process, HMODULE module,
                                                MODULEINFO *modinfo, DWORD cb);
 
+static NTSTATUS (WINAPI *pLdrGetDllDirectory)(UNICODE_STRING*);
+static NTSTATUS (WINAPI *pLdrSetDllDirectory)(UNICODE_STRING*);
+
 static BOOL is_unicode_enabled = TRUE;
 
 static BOOL cmpStrAW(const char* a, const WCHAR* b, DWORD lenA, DWORD lenB)
@@ -601,11 +609,14 @@ static void test_LoadLibraryEx_search_flags(void)
     ok( !mod, "LoadLibrary succeeded\n" );
     ok( GetLastError() == ERROR_MOD_NOT_FOUND, "wrong error %u\n", GetLastError() );
 
+    if (0)  /* crashes on win10 */
+    {
     SetLastError( 0xdeadbeef );
     mod = LoadLibraryExA( "winetestdll.dll", 0, LOAD_LIBRARY_SEARCH_USER_DIRS );
     ok( !mod, "LoadLibrary succeeded\n" );
     ok( GetLastError() == ERROR_MOD_NOT_FOUND || broken(GetLastError() == ERROR_NOT_ENOUGH_MEMORY),
         "wrong error %u\n", GetLastError() );
+    }
 
     SetLastError( 0xdeadbeef );
     mod = LoadLibraryExA( "winetestdll.dll", 0, LOAD_LIBRARY_SEARCH_SYSTEM32 );
@@ -753,8 +764,8 @@ static void testGetDllDirectory(void)
         bufferW[0] = 'A';
         ret = pGetDllDirectoryW(0, bufferW);
         ok(ret == length + 1, "i=%d, Expected %u, got %u\n", i, length + 1, ret);
-        ok(bufferW[0] == 0 || /* XP, 2003 */
-           broken(bufferW[0] == 'A'), "i=%d, Buffer overflow\n", i);
+        ok(bufferW[0] == 'A' || broken(bufferW[0] == 0), /* XP, 2003 */
+           "i=%d, Buffer overflow\n", i);
 
         /* buffer just one too short */
         bufferA[0] = 'A';
@@ -766,8 +777,8 @@ static void testGetDllDirectory(void)
         bufferW[0] = 'A';
         ret = pGetDllDirectoryW(length, bufferW);
         ok(ret == length + 1, "i=%d, Expected %u, got %u\n", i, length + 1, ret);
-        ok(bufferW[0] == 0 || /* XP, 2003 */
-           broken(bufferW[0] == 'A'), "i=%d, Buffer overflow\n", i);
+        if (length != 0)
+            ok(bufferW[0] == 0, "i=%d, Buffer overflow\n", i);
 
         if (0)
         {
@@ -779,6 +790,30 @@ static void testGetDllDirectory(void)
             ret = pGetDllDirectoryW(length, NULL);
             ok(ret == length + 1, "i=%d, Expected %u, got %u\n", i, length + 1, ret);
         }
+
+        if (pLdrGetDllDirectory)
+        {
+            UNICODE_STRING str;
+            NTSTATUS status;
+            str.Buffer = bufferW;
+            str.MaximumLength = sizeof(bufferW);
+            status = pLdrGetDllDirectory( &str );
+            ok( !status, "LdrGetDllDirectory failed %x\n", status );
+            ok( cmpStrAW( dll_directories[i], bufferW, strlen(dll_directories[i]),
+                          str.Length / sizeof(WCHAR) ), "%u: got %s instead of %s\n",
+                i, wine_dbgstr_w(bufferW), dll_directories[i] );
+            if (dll_directories[i][0])
+            {
+                memset( bufferW, 0xcc, sizeof(bufferW) );
+                str.MaximumLength = (strlen( dll_directories[i] ) - 1) * sizeof(WCHAR);
+                status = pLdrGetDllDirectory( &str );
+                ok( status == STATUS_BUFFER_TOO_SMALL, "%u: LdrGetDllDirectory failed %x\n", i, status );
+                ok( bufferW[0] == 0 && bufferW[1] == 0xcccc,
+                    "%u: buffer %x %x\n", i, bufferW[0], bufferW[1] );
+                length = (strlen( dll_directories[i] ) + 1) * sizeof(WCHAR);
+                ok( str.Length == length, "%u: wrong len %u / %u\n", i, str.Length, length );
+            }
+        }
     }
 
     /* unset whatever we did so following tests won't be affected */
@@ -787,9 +822,9 @@ static void testGetDllDirectory(void)
 
 static void init_pointers(void)
 {
-    HMODULE hKernel32 = GetModuleHandleA("kernel32.dll");
+    HMODULE mod = GetModuleHandleA("kernel32.dll");
 
-#define MAKEFUNC(f) (p##f = (void*)GetProcAddress(hKernel32, #f))
+#define MAKEFUNC(f) (p##f = (void*)GetProcAddress(mod, #f))
     MAKEFUNC(GetDllDirectoryA);
     MAKEFUNC(GetDllDirectoryW);
     MAKEFUNC(SetDllDirectoryA);
@@ -797,6 +832,9 @@ static void init_pointers(void)
     MAKEFUNC(RemoveDllDirectory);
     MAKEFUNC(SetDefaultDllDirectories);
     MAKEFUNC(K32GetModuleInformation);
+    mod = GetModuleHandleA( "ntdll.dll" );
+    MAKEFUNC(LdrGetDllDirectory);
+    MAKEFUNC(LdrSetDllDirectory);
 #undef MAKEFUNC
 
     /* before Windows 7 this was not exported in kernel32 */

dlls/ntdll/loader.c

@@ -78,6 +78,7 @@ static BOOL imports_fixup_done = FALSE;  /* set once the imports have been fixed
 static BOOL process_detaching = FALSE;  /* set on process detach to avoid deadlocks with thread detach */
 static int free_lib_count;   /* recursion depth of LdrUnloadDll calls */
 static ULONG path_safe_mode;  /* path mode set by RtlSetSearchPathMode */
+static UNICODE_STRING dll_directory;  /* extra path for LdrSetDllDirectory */
 
 struct ldr_notification
 {
@@ -137,6 +138,15 @@ static RTL_CRITICAL_SECTION_DEBUG critsect_debug =
 };
 static RTL_CRITICAL_SECTION loader_section = { &critsect_debug, -1, 0, 0, 0, 0 };
 
+static CRITICAL_SECTION dlldir_section;
+static CRITICAL_SECTION_DEBUG dlldir_critsect_debug =
+{
+    0, 0, &dlldir_section,
+    { &dlldir_critsect_debug.ProcessLocksList, &dlldir_critsect_debug.ProcessLocksList },
+      0, 0, { (DWORD_PTR)(__FILE__ ": dlldir_section") }
+};
+static CRITICAL_SECTION dlldir_section = { &dlldir_critsect_debug, -1, 0, 0, 0, 0 };
+
 static WINE_MODREF *cached_modref;
 static WINE_MODREF *current_modref;
 static WINE_MODREF *last_failed_modref;
@@ -3827,6 +3837,45 @@ PVOID WINAPI RtlPcToFileHeader( PVOID pc, PVOID *address )
 }
 
 
+/****************************************************************************
+ *		LdrGetDllDirectory  (NTDLL.@)
+ */
+NTSTATUS WINAPI LdrGetDllDirectory( UNICODE_STRING *dir )
+{
+    NTSTATUS status = STATUS_SUCCESS;
+
+    RtlEnterCriticalSection( &dlldir_section );
+    dir->Length = dll_directory.Length + sizeof(WCHAR);
+    if (dir->MaximumLength >= dir->Length) RtlCopyUnicodeString( dir, &dll_directory );
+    else
+    {
+        status = STATUS_BUFFER_TOO_SMALL;
+        if (dir->MaximumLength) dir->Buffer[0] = 0;
+    }
+    RtlLeaveCriticalSection( &dlldir_section );
+    return status;
+}
+
+
+/****************************************************************************
+ *		LdrSetDllDirectory  (NTDLL.@)
+ */
+NTSTATUS WINAPI LdrSetDllDirectory( const UNICODE_STRING *dir )
+{
+    NTSTATUS status = STATUS_SUCCESS;
+    UNICODE_STRING new;
+
+    if (!dir->Buffer) RtlInitUnicodeString( &new, NULL );
+    else if ((status = RtlDuplicateUnicodeString( 1, dir, &new ))) return status;
+
+    RtlEnterCriticalSection( &dlldir_section );
+    RtlFreeUnicodeString( &dll_directory );
+    dll_directory = new;
+    RtlLeaveCriticalSection( &dlldir_section );
+    return status;
+}
+
+
 /*************************************************************************
  *		RtlSetSearchPathMode (NTDLL.@)
  */

dlls/ntdll/ntdll.spec

@@ -86,6 +86,7 @@
 @ stub LdrFlushAlternateResourceModules
 @ stdcall LdrGetDllHandle(wstr long ptr ptr)
 # @ stub LdrGetDllHandleEx
+@ stdcall LdrGetDllDirectory(ptr)
 @ stdcall LdrGetProcedureAddress(ptr ptr long ptr)
 # @ stub LdrHotPatchRoutine
 @ stub LdrInitShimEngineDynamic
@@ -99,6 +100,7 @@
 @ stdcall LdrRegisterDllNotification(long ptr ptr ptr)
 @ stdcall LdrResolveDelayLoadedAPI(ptr ptr ptr ptr ptr long)
 @ stub LdrSetAppCompatDllRedirectionCallback
+@ stdcall LdrSetDllDirectory(ptr)
 @ stub LdrSetDllManifestProber
 @ stdcall LdrShutdownProcess()
 @ stdcall LdrShutdownThread()

include/winternl.h

@@ -2310,10 +2310,11 @@ NTSYSAPI void      WINAPI DbgUiRemoteBreakin(void*);
 NTSYSAPI void      WINAPI DbgUserBreakPoint(void);
 NTSYSAPI NTSTATUS  WINAPI LdrAccessResource(HMODULE,const IMAGE_RESOURCE_DATA_ENTRY*,void**,PULONG);
 NTSYSAPI NTSTATUS  WINAPI LdrAddRefDll(ULONG,HMODULE);
-NTSYSAPI NTSTATUS  WINAPI LdrFindResourceDirectory_U(HMODULE,const LDR_RESOURCE_INFO*,ULONG,const IMAGE_RESOURCE_DIRECTORY**);
-NTSYSAPI NTSTATUS  WINAPI LdrFindResource_U(HMODULE,const LDR_RESOURCE_INFO*,ULONG,const IMAGE_RESOURCE_DATA_ENTRY**);
 NTSYSAPI NTSTATUS  WINAPI LdrDisableThreadCalloutsForDll(HMODULE);
 NTSYSAPI NTSTATUS  WINAPI LdrFindEntryForAddress(const void*, PLDR_MODULE*);
+NTSYSAPI NTSTATUS  WINAPI LdrFindResourceDirectory_U(HMODULE,const LDR_RESOURCE_INFO*,ULONG,const IMAGE_RESOURCE_DIRECTORY**);
+NTSYSAPI NTSTATUS  WINAPI LdrFindResource_U(HMODULE,const LDR_RESOURCE_INFO*,ULONG,const IMAGE_RESOURCE_DATA_ENTRY**);
+NTSYSAPI NTSTATUS  WINAPI LdrGetDllDirectory(UNICODE_STRING*);
 NTSYSAPI NTSTATUS  WINAPI LdrGetDllHandle(LPCWSTR, ULONG, const UNICODE_STRING*, HMODULE*);
 NTSYSAPI NTSTATUS  WINAPI LdrGetProcedureAddress(HMODULE, const ANSI_STRING*, ULONG, void**);
 NTSYSAPI NTSTATUS  WINAPI LdrLoadDll(LPCWSTR, DWORD, const UNICODE_STRING*, HMODULE*);
@@ -2321,6 +2322,7 @@ NTSYSAPI NTSTATUS  WINAPI LdrLockLoaderLock(ULONG,ULONG*,ULONG_PTR*);
 IMAGE_BASE_RELOCATION * WINAPI LdrProcessRelocationBlock(void*,UINT,USHORT*,INT_PTR);
 NTSYSAPI NTSTATUS  WINAPI LdrQueryImageFileExecutionOptions(const UNICODE_STRING*,LPCWSTR,ULONG,void*,ULONG,ULONG*);
 NTSYSAPI NTSTATUS  WINAPI LdrQueryProcessModuleInformation(SYSTEM_MODULE_INFORMATION*, ULONG, ULONG*);
+NTSYSAPI NTSTATUS  WINAPI LdrSetDllDirectory(const UNICODE_STRING*);
 NTSYSAPI void      WINAPI LdrShutdownProcess(void);
 NTSYSAPI void      WINAPI LdrShutdownThread(void);
 NTSYSAPI NTSTATUS  WINAPI LdrUnloadDll(HMODULE);