From 97e2af1f73b6985bd9c91fce18d0de1a1d95dc85 Mon Sep 17 00:00:00 2001
From: Nikolay Sivov <nsivov@codeweavers.com>
Date: Fri, 16 May 2014 09:03:27 +0400
Subject: [PATCH] ntdll: Implement non-blocking mode for LdrLockLoaderLock().

---
 dlls/ntdll/loader.c    | 26 ++++++++++++++++----
 dlls/ntdll/tests/rtl.c | 56 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 77 insertions(+), 5 deletions(-)

diff --git a/dlls/ntdll/loader.c b/dlls/ntdll/loader.c
index d458a31471a..64712901bf0 100644
--- a/dlls/ntdll/loader.c
+++ b/dlls/ntdll/loader.c
@@ -1337,17 +1337,33 @@ NTSTATUS WINAPI LdrFindEntryForAddress(const void* addr, PLDR_MODULE* pmod)
 /******************************************************************
  *		LdrLockLoaderLock  (NTDLL.@)
  *
- * Note: flags are not implemented.
+ * Note: some flags are not implemented.
  * Flag 0x01 is used to raise exceptions on errors.
- * Flag 0x02 is used to avoid waiting on the section (does RtlTryEnterCriticalSection instead).
  */
 NTSTATUS WINAPI LdrLockLoaderLock( ULONG flags, ULONG *result, ULONG *magic )
 {
-    if (flags) FIXME( "flags %x not supported\n", flags );
+    if (flags & ~0x2) FIXME( "flags %x not supported\n", flags );
 
-    if (result) *result = 1;
+    if (result) *result = 0;
+    if (magic) *magic = 0;
+    if (flags & ~0x3) return STATUS_INVALID_PARAMETER_1;
+    if (!result && (flags & 0x2)) return STATUS_INVALID_PARAMETER_2;
     if (!magic) return STATUS_INVALID_PARAMETER_3;
-    RtlEnterCriticalSection( &loader_section );
+
+    if (flags & 0x2)
+    {
+        if (!RtlTryEnterCriticalSection( &loader_section ))
+        {
+            *result = 2;
+            return STATUS_SUCCESS;
+        }
+        *result = 1;
+    }
+    else
+    {
+        RtlEnterCriticalSection( &loader_section );
+        if (result) *result = 1;
+    }
     *magic = GetCurrentThreadId();
     return STATUS_SUCCESS;
 }
diff --git a/dlls/ntdll/tests/rtl.c b/dlls/ntdll/tests/rtl.c
index 6eb6de3cb04..72cee050843 100644
--- a/dlls/ntdll/tests/rtl.c
+++ b/dlls/ntdll/tests/rtl.c
@@ -90,6 +90,8 @@ static CHAR *    (WINAPI *pRtlIpv4AddressToStringA)(const IN_ADDR *, LPSTR);
 static NTSTATUS  (WINAPI *pRtlIpv4AddressToStringExA)(const IN_ADDR *, USHORT, LPSTR, PULONG);
 static NTSTATUS  (WINAPI *pRtlIpv4StringToAddressA)(PCSTR, BOOLEAN, PCSTR *, IN_ADDR *);
 static NTSTATUS  (WINAPI *pLdrAddRefDll)(ULONG, HMODULE);
+static NTSTATUS  (WINAPI *pLdrLockLoaderLock)(ULONG, ULONG*, ULONG*);
+static NTSTATUS  (WINAPI *pLdrUnlockLoaderLock)(ULONG, ULONG);
 
 static HMODULE hkernel32 = 0;
 static BOOL      (WINAPI *pIsWow64Process)(HANDLE, PBOOL);
@@ -135,6 +137,8 @@ static void InitFunctionPtrs(void)
         pRtlIpv4AddressToStringExA = (void *)GetProcAddress(hntdll, "RtlIpv4AddressToStringExA");
         pRtlIpv4StringToAddressA = (void *)GetProcAddress(hntdll, "RtlIpv4StringToAddressA");
         pLdrAddRefDll = (void *)GetProcAddress(hntdll, "LdrAddRefDll");
+        pLdrLockLoaderLock = (void *)GetProcAddress(hntdll, "LdrLockLoaderLock");
+        pLdrUnlockLoaderLock = (void *)GetProcAddress(hntdll, "LdrUnlockLoaderLock");
     }
     hkernel32 = LoadLibraryA("kernel32.dll");
     ok(hkernel32 != 0, "LoadLibrary failed\n");
@@ -1543,6 +1547,57 @@ static void test_LdrAddRefDll(void)
     ok(mod2 != NULL, "got %p\n", mod2);
 }
 
+static void test_LdrLockLoaderLock(void)
+{
+    ULONG result, magic;
+    NTSTATUS status;
+
+    if (!pLdrLockLoaderLock)
+    {
+        win_skip("LdrLockLoaderLock() is not available\n");
+        return;
+    }
+
+    /* invalid flags */
+    result = 10;
+    magic = 0xdeadbeef;
+    status = pLdrLockLoaderLock(0x10, &result, &magic);
+    ok(status == STATUS_INVALID_PARAMETER_1, "got 0x%08x\n", status);
+    ok(result == 0, "got %d\n", result);
+    ok(magic == 0, "got 0x%08x\n", magic);
+
+    magic = 0xdeadbeef;
+    status = pLdrLockLoaderLock(0x10, NULL, &magic);
+    ok(status == STATUS_INVALID_PARAMETER_1, "got 0x%08x\n", status);
+    ok(magic == 0, "got 0x%08x\n", magic);
+
+    result = 10;
+    status = pLdrLockLoaderLock(0x10, &result, NULL);
+    ok(status == STATUS_INVALID_PARAMETER_1, "got 0x%08x\n", status);
+    ok(result == 0, "got %d\n", result);
+
+    /* non-blocking mode, result is null */
+    magic = 0xdeadbeef;
+    status = pLdrLockLoaderLock(0x2, NULL, &magic);
+    ok(status == STATUS_INVALID_PARAMETER_2, "got 0x%08x\n", status);
+    ok(magic == 0, "got 0x%08x\n", magic);
+
+    /* magic pointer is null */
+    result = 10;
+    status = pLdrLockLoaderLock(0, &result, NULL);
+    ok(status == STATUS_INVALID_PARAMETER_3, "got 0x%08x\n", status);
+    ok(result == 0, "got %d\n", result);
+
+    /* lock in non-blocking mode */
+    result = 0;
+    magic = 0;
+    status = pLdrLockLoaderLock(0x2, &result, &magic);
+    ok(status == STATUS_SUCCESS, "got 0x%08x\n", status);
+    ok(result == 1, "got %d\n", result);
+    ok(magic != 0, "got 0x%08x\n", magic);
+    pLdrUnlockLoaderLock(0, magic);
+}
+
 START_TEST(rtl)
 {
     InitFunctionPtrs();
@@ -1568,4 +1623,5 @@ START_TEST(rtl)
     test_RtlIpv4AddressToStringEx();
     test_RtlIpv4StringToAddress();
     test_LdrAddRefDll();
+    test_LdrLockLoaderLock();
 }