[PATCH v2] ntdll/tests: Add tests for RtlEncode(System)Pointer.

[Myah Caron]

Signed-off-by: Myah Caron <qsniyg at protonmail.com>
---
v2:
 - Fix the algorithm according to the linked MSDN post
 - Don't require _rotr(64)

 dlls/ntdll/tests/rtl.c | 88 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 88 insertions(+)

diff --git a/dlls/ntdll/tests/rtl.c b/dlls/ntdll/tests/rtl.c
index 7a62670ea05..fef7103d677 100644
--- a/dlls/ntdll/tests/rtl.c
+++ b/dlls/ntdll/tests/rtl.c
@@ -82,6 +82,11 @@ static NTSTATUS  (WINAPI *pRtlInitializeCriticalSectionEx)(CRITICAL_SECTION *, U
 static NTSTATUS  (WINAPI *pLdrEnumerateLoadedModules)(void *, void *, void *);
 static NTSTATUS  (WINAPI *pLdrRegisterDllNotification)(ULONG, PLDR_DLL_NOTIFICATION_FUNCTION, void *, void **);
 static NTSTATUS  (WINAPI *pLdrUnregisterDllNotification)(void *);
+static VOID*     (WINAPI *pRtlEncodePointer)(void *);
+static VOID*     (WINAPI *pRtlDecodePointer)(void *);
+static VOID*     (WINAPI *pRtlEncodeSystemPointer)(void *);
+static VOID*     (WINAPI *pRtlDecodeSystemPointer)(void *);
+static NTSTATUS  (WINAPI *pNtQueryInformationProcess)(HANDLE, PROCESSINFOCLASS, PVOID, ULONG, PULONG);

 static HMODULE hkernel32 = 0;
 static BOOL      (WINAPI *pIsWow64Process)(HANDLE, PBOOL);
@@ -122,6 +127,11 @@ static void InitFunctionPtrs(void)
         pLdrEnumerateLoadedModules = (void *)GetProcAddress(hntdll, "LdrEnumerateLoadedModules");
         pLdrRegisterDllNotification = (void *)GetProcAddress(hntdll, "LdrRegisterDllNotification");
         pLdrUnregisterDllNotification = (void *)GetProcAddress(hntdll, "LdrUnregisterDllNotification");
+        pRtlEncodePointer = (void *)GetProcAddress(hntdll, "RtlEncodePointer");
+        pRtlDecodePointer = (void *)GetProcAddress(hntdll, "RtlDecodePointer");
+        pRtlEncodeSystemPointer = (void *)GetProcAddress(hntdll, "RtlEncodeSystemPointer");
+        pRtlDecodeSystemPointer = (void *)GetProcAddress(hntdll, "RtlDecodeSystemPointer");
+        pNtQueryInformationProcess = (void *)GetProcAddress(hntdll, "NtQueryInformationProcess");
     }
     hkernel32 = LoadLibraryA("kernel32.dll");
     ok(hkernel32 != 0, "LoadLibrary failed\n");
@@ -3483,6 +3493,83 @@ static void test_LdrRegisterDllNotification(void)
     pLdrUnregisterDllNotification(cookie);
 }

+static DWORD_PTR rotr_ptr( DWORD_PTR num, int shift )
+{
+#ifdef _WIN64
+#define ROT_BITS 64
+#else
+#define ROT_BITS 32
+#endif
+
+    shift &= ROT_BITS - 1;
+    return (num >> shift) | (num << (ROT_BITS-shift));
+
+#undef ROT_BITS
+}
+
+static PVOID our_encode_pointer( PVOID ptr, ULONG cookie )
+{
+    DWORD_PTR ptrval = (DWORD_PTR) ptr;
+    DWORD_PTR cookieval = (DWORD_PTR) cookie;
+
+    /* http://blogs.msdn.com/b/michael_howard/archive/2006/08/16/702707.aspx */
+
+    ptrval = (ptrval ^ cookieval);
+    return (PVOID)rotr_ptr(ptrval, cookie);
+}
+
+static void test_RtlEncodePointer(void)
+{
+    void *addr = (void*)0xdeadbeef0123abcd;
+    void *encoded, *encoded1, *decoded;
+    void *our_encoded;
+    ULONG process_cookie = 0xdeadbeef;
+    ULONG system_cookie = *(ULONG*)(0x7ffe0000 + 0x330); /* user_shared_data.Cookie */
+    NTSTATUS status;
+
+    if (!pRtlEncodePointer || !pRtlDecodePointer)
+    {
+        win_skip("RtlEn/DecodePointer not available\n");
+        return;
+    }
+
+    encoded = pRtlEncodePointer(addr);
+    ok(encoded != addr, "got %p\n", encoded);
+
+    decoded = pRtlDecodePointer(encoded);
+    ok(decoded == addr, "got %p\n", decoded);
+
+    if (!pNtQueryInformationProcess)
+    {
+        win_skip("NtQueryInformationProcess not available\n");
+    }
+    else
+    {
+        status = pNtQueryInformationProcess(GetCurrentProcess(), ProcessCookie, &process_cookie, sizeof(process_cookie), NULL);
+        ok(status == S_OK, "got %08x\n", status);
+        ok(process_cookie != 0xdeadbeef, "got %08x\n", process_cookie);
+
+        our_encoded = our_encode_pointer(addr, process_cookie);
+        todo_wine ok(encoded == our_encoded || broken(encoded == (void*)((uintptr_t)addr ^ (uintptr_t)process_cookie)), "got %p (ours %p)\n", encoded, our_encoded);
+    }
+
+    if (!pRtlEncodeSystemPointer || !pRtlDecodeSystemPointer)
+    {
+        win_skip("RtlEn/DecodeSystemPointer not available\n");
+        return;
+    }
+
+    encoded1 = pRtlEncodeSystemPointer(addr);
+    ok(encoded != addr, "got %p\n", encoded);
+    todo_wine ok(encoded1 != encoded, "got %p\n", encoded1);
+
+    decoded = pRtlDecodeSystemPointer(encoded1);
+    ok(decoded == addr, "got %p\n", decoded);
+
+    our_encoded = our_encode_pointer(addr, system_cookie);
+    todo_wine ok(encoded1 == our_encoded || broken(encoded1 == (void*)((uintptr_t)addr ^ (uintptr_t)system_cookie)), "got %p (ours %p)\n", encoded1, our_encoded);
+}
+
 START_TEST(rtl)
 {
     InitFunctionPtrs();
@@ -3523,4 +3610,5 @@ START_TEST(rtl)
     test_LdrEnumerateLoadedModules();
     test_RtlMakeSelfRelativeSD();
     test_LdrRegisterDllNotification();
+    test_RtlEncodePointer();
 }
--
2.28.0