From 146ff49161fb35095ae9ac6f4c40281cd8f748f4 Mon Sep 17 00:00:00 2001
From: Nikolay Sivov <nsivov@codeweavers.com>
Date: Mon, 14 Oct 2013 08:07:58 +0400
Subject: [PATCH] ntdll: Support pinning module refcount with LdrAddRefDll().

---
 dlls/ntdll/loader.c    |  7 ++++--
 dlls/ntdll/tests/rtl.c | 52 ++++++++++++++++++++++++++++++++++++++++++
 include/winternl.h     |  3 +++
 3 files changed, 60 insertions(+), 2 deletions(-)

diff --git a/dlls/ntdll/loader.c b/dlls/ntdll/loader.c
index b1c7600df0b..3a28b414ea4 100644
--- a/dlls/ntdll/loader.c
+++ b/dlls/ntdll/loader.c
@@ -2124,13 +2124,16 @@ NTSTATUS WINAPI LdrAddRefDll( ULONG flags, HMODULE module )
     NTSTATUS ret = STATUS_SUCCESS;
     WINE_MODREF *wm;
 
-    if (flags) FIXME( "%p flags %x not implemented\n", module, flags );
+    if (flags & ~LDR_ADDREF_DLL_PIN) FIXME( "%p flags %x not implemented\n", module, flags );
 
     RtlEnterCriticalSection( &loader_section );
 
     if ((wm = get_modref( module )))
     {
-        if (wm->ldr.LoadCount != -1) wm->ldr.LoadCount++;
+        if (flags & LDR_ADDREF_DLL_PIN)
+            wm->ldr.LoadCount = -1;
+        else
+            if (wm->ldr.LoadCount != -1) wm->ldr.LoadCount++;
         TRACE( "(%s) ldr.LoadCount: %d\n", debugstr_w(wm->ldr.BaseDllName.Buffer), wm->ldr.LoadCount );
     }
     else ret = STATUS_INVALID_PARAMETER;
diff --git a/dlls/ntdll/tests/rtl.c b/dlls/ntdll/tests/rtl.c
index 38909a31d1f..54e3db58f0f 100644
--- a/dlls/ntdll/tests/rtl.c
+++ b/dlls/ntdll/tests/rtl.c
@@ -89,6 +89,7 @@ static IMAGE_BASE_RELOCATION *(WINAPI *pLdrProcessRelocationBlock)(void*,UINT,US
 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 HMODULE hkernel32 = 0;
 static BOOL      (WINAPI *pIsWow64Process)(HANDLE, PBOOL);
@@ -133,6 +134,7 @@ static void InitFunctionPtrs(void)
         pRtlIpv4AddressToStringA = (void *)GetProcAddress(hntdll, "RtlIpv4AddressToStringA");
         pRtlIpv4AddressToStringExA = (void *)GetProcAddress(hntdll, "RtlIpv4AddressToStringExA");
         pRtlIpv4StringToAddressA = (void *)GetProcAddress(hntdll, "RtlIpv4StringToAddressA");
+        pLdrAddRefDll = (void *)GetProcAddress(hntdll, "LdrAddRefDll");
     }
     hkernel32 = LoadLibraryA("kernel32.dll");
     ok(hkernel32 != 0, "LoadLibrary failed\n");
@@ -1486,6 +1488,55 @@ static void test_RtlIpv4StringToAddress(void)
     }
 }
 
+static void test_LdrAddRefDll(void)
+{
+    HMODULE mod, mod2;
+    NTSTATUS status;
+    BOOL ret;
+
+    mod = LoadLibraryA("comctl32.dll");
+    ok(mod != NULL, "got %p\n", mod);
+    ret = FreeLibrary(mod);
+    ok(ret, "got %d\n", ret);
+
+    mod2 = GetModuleHandleA("comctl32.dll");
+    ok(mod2 == NULL, "got %p\n", mod2);
+
+    /* load, addref and release 2 times */
+    mod = LoadLibraryA("comctl32.dll");
+    ok(mod != NULL, "got %p\n", mod);
+    status = pLdrAddRefDll(0, mod);
+    ok(status == STATUS_SUCCESS, "got 0x%08x\n", status);
+    ret = FreeLibrary(mod);
+    ok(ret, "got %d\n", ret);
+
+    mod2 = GetModuleHandleA("comctl32.dll");
+    ok(mod2 != NULL, "got %p\n", mod2);
+    ret = FreeLibrary(mod);
+    ok(ret, "got %d\n", ret);
+
+    mod2 = GetModuleHandleA("comctl32.dll");
+    ok(mod2 == NULL, "got %p\n", mod2);
+
+    /* pin refcount */
+    mod = LoadLibraryA("comctl32.dll");
+    ok(mod != NULL, "got %p\n", mod);
+    status = pLdrAddRefDll(LDR_ADDREF_DLL_PIN, mod);
+    ok(status == STATUS_SUCCESS, "got 0x%08x\n", status);
+
+    ret = FreeLibrary(mod);
+    ok(ret, "got %d\n", ret);
+    ret = FreeLibrary(mod);
+    ok(ret, "got %d\n", ret);
+    ret = FreeLibrary(mod);
+    ok(ret, "got %d\n", ret);
+    ret = FreeLibrary(mod);
+    ok(ret, "got %d\n", ret);
+
+    mod2 = GetModuleHandleA("comctl32.dll");
+    ok(mod2 != NULL, "got %p\n", mod2);
+}
+
 START_TEST(rtl)
 {
     InitFunctionPtrs();
@@ -1510,4 +1561,5 @@ START_TEST(rtl)
     test_RtlIpv4AddressToString();
     test_RtlIpv4AddressToStringEx();
     test_RtlIpv4StringToAddress();
+    test_LdrAddRefDll();
 }
diff --git a/include/winternl.h b/include/winternl.h
index c8794503ff5..d93b37aca0f 100644
--- a/include/winternl.h
+++ b/include/winternl.h
@@ -1970,6 +1970,9 @@ typedef struct _LDR_MODULE
 #define LDR_DONT_RESOLVE_REFS           0x40000000
 #define LDR_WINE_INTERNAL               0x80000000
 
+/* flag for LdrAddRefDll */
+#define LDR_ADDREF_DLL_PIN              0x00000001
+
 /* FIXME: to be checked */
 #define MAXIMUM_FILENAME_LENGTH 256