[PATCH v2 06/22] ntdll: Use shared loader locking in LdrQueryProcessModuleInformation().

[Paul Gofman]

Signed-off-by: Paul Gofman <pgofman at codeweavers.com>
---
    - remove redundant 'volatile' qualifier from locked_exclusive.

 dlls/kernel32/tests/loader.c | 46 ++++++++++++++++++++++++++++++++++++
 dlls/ntdll/loader.c          | 37 ++++++++++++++++++++++++++---
 2 files changed, 80 insertions(+), 3 deletions(-)

diff --git a/dlls/kernel32/tests/loader.c b/dlls/kernel32/tests/loader.c
index 045338c1289..e97c114c5e9 100644
--- a/dlls/kernel32/tests/loader.c
+++ b/dlls/kernel32/tests/loader.c
@@ -4049,6 +4049,16 @@ static volatile LONG test_loader_lock_timeout_count;
 
 #define BLOCKING_TESTS_ENABLED 0
 
+static void CALLBACK test_loader_lock_module_enum_locking_callback(LDR_DATA_TABLE_ENTRY *mod,
+        void *context, BOOLEAN *stop)
+{
+    SetEvent(lock_ready_event);
+    if (WaitForSingleObject(next_test_event, 3000) == WAIT_TIMEOUT)
+        ++test_loader_lock_timeout_count;
+
+    *stop = TRUE;
+}
+
 static DWORD WINAPI test_loader_lock_thread(void *param)
 {
     ULONG_PTR magic;
@@ -4085,6 +4095,20 @@ static DWORD WINAPI test_loader_lock_thread(void *param)
     SetEvent(lock_ready_event);
     WaitForSingleObject(next_test_event, INFINITE);
 
+    /* 2. Test with the thread blocked in LdrEnumerateLoadedModules callback. */
+    do
+    {
+        pLdrEnumerateLoadedModules(NULL, test_loader_lock_module_enum_locking_callback, NULL);
+        if (test_loader_lock_timeout_count)
+        {
+            WaitForSingleObject(next_test_event, INFINITE);
+            test_loader_lock_timeout_count = 0;
+        }
+    } while (test_loader_lock_repeat_lock);
+
+    SetEvent(lock_ready_event);
+    WaitForSingleObject(next_test_event, INFINITE);
+
     SetEvent(lock_ready_event);
     return 0;
 }
@@ -4131,6 +4155,7 @@ static void test_loader_lock(void)
     static const WCHAR not_loaded_dll_name[] = L"authz.dll";
     static const WCHAR preloaded_dll_name[] = L"winmm.dll";
     HMODULE hmodule_preloaded, hmodule;
+    ULONG_PTR magic;
     NTSTATUS status;
     HANDLE thread;
     void *cookie;
@@ -4200,6 +4225,27 @@ static void test_loader_lock(void)
     LdrUnregisterDllNotification( cookie );
     check_timeout(FALSE);
 
+    /* 2. Test with the thread blocked in LdrEnumerateLoadedModules callback. */
+    trace("Test 2.\n");
+    test_loader_lock_next_test();
+
+    status = LdrRegisterDllNotification(0, test_loader_lock_ldr_notify, NULL, &cookie);
+    ok(!status, "Got unexpected status %#x.\n", status);
+    check_timeout(FALSE);
+
+    if (BLOCKING_TESTS_ENABLED)
+    {
+        status = pLdrLockLoaderLock(0, NULL, &magic);
+        ok(!status, "Got unexpected status %#x.\n", status);
+        status = pLdrUnlockLoaderLock(0, magic);
+        ok(!status, "Got unexpected status %#x.\n", status);
+
+        check_timeout(TRUE);
+    }
+
+    LdrUnregisterDllNotification( cookie );
+    check_timeout(FALSE);
+
     test_loader_lock_next_test();
 
 done:
diff --git a/dlls/ntdll/loader.c b/dlls/ntdll/loader.c
index d0d4e5448ed..092805cc945 100644
--- a/dlls/ntdll/loader.c
+++ b/dlls/ntdll/loader.c
@@ -161,6 +161,7 @@ static RTL_CRITICAL_SECTION_DEBUG critsect_debug =
 static RTL_CRITICAL_SECTION loader_section = { &critsect_debug, -1, 0, 0, 0, 0 };
 
 static RTL_SRWLOCK loader_srw_lock = RTL_SRWLOCK_INIT;
+static BOOL locked_exclusive;
 
 static CRITICAL_SECTION dlldir_section;
 static CRITICAL_SECTION_DEBUG dlldir_critsect_debug =
@@ -250,8 +251,32 @@ static void lock_loader_exclusive(void)
     ULONG recursion_count = inc_recursion_count();
 
     TRACE( "recursion_count %u.\n", recursion_count );
+    if (!recursion_count)
+    {
+        if (!RtlDllShutdownInProgress())
+            RtlAcquireSRWLockExclusive( &loader_srw_lock );
+        locked_exclusive = TRUE;
+    }
+    else
+    {
+        assert( locked_exclusive );
+    }
+}
+
+/*************************************************************************
+ *		lock_loader_shared
+ *
+ * Take shared ownership of internal loader lock.
+ * If the thread already has exclusive lock it will stay exclusive.
+ */
+static void lock_loader_shared(void)
+{
+    ULONG recursion_count = inc_recursion_count();
+
+    TRACE("recursion_count %u, locked_exclusive %d.\n", recursion_count, locked_exclusive);
+
     if (!recursion_count && !RtlDllShutdownInProgress())
-        RtlAcquireSRWLockExclusive( &loader_srw_lock );
+        RtlAcquireSRWLockShared( &loader_srw_lock );
 }
 
 /*************************************************************************
@@ -269,8 +294,14 @@ static void unlock_loader(void)
 
     assert( recursion_count != ~0u );
 
-    if (!recursion_count)
+    if (recursion_count) return;
+
+    if (locked_exclusive)
+    {
+        locked_exclusive = FALSE;
         RtlReleaseSRWLockExclusive( &loader_srw_lock );
+    }
+    else RtlReleaseSRWLockShared( &loader_srw_lock );
 }
 
 #define RTL_UNLOAD_EVENT_TRACE_NUMBER 64
@@ -3283,7 +3314,7 @@ NTSTATUS WINAPI LdrQueryProcessModuleInformation(RTL_PROCESS_MODULES *smi,
 
     smi->ModulesCount = 0;
 
-    lock_loader_exclusive();
+    lock_loader_shared();
     mark = &NtCurrentTeb()->Peb->LdrData->InLoadOrderModuleList;
     for (entry = mark->Flink; entry != mark; entry = entry->Flink)
     {
-- 
2.31.1