[PATCH v5 3/6] ntdll/tests: Add NtAllocateVirtualMemory tests for zero_bits behavior

[Rémi Bernon]

The zero_bits parameter doesn't behave as expected, and some 64bit code
use it to allocate memory in the lower 32bit address space.

The expected full behaviour is:

* zero_bits == 0: no constraint on address range
* 0 < zero_bits <= 15: returned address should have as many upper bits
                       set to 0, starting at bit 31. In 64bit mode,
                       upper 32bits should all be 0 as well.
* 15 < zero_bits <= 31: unsure, but probably same as zero_bits == 15.
* zero_bits > 31: (64bit/WoW64 only) zero_bits behaves as a bitmask, as
                  if it was set to the number of leading 0 in the
                  bitmask, works in the whole 64bit range.

Signed-off-by: Rémi Bernon 
---
 dlls/ntdll/tests/virtual.c | 60 ++++++++++++++++++++++++++++++++++----
 1 file changed, 55 insertions(+), 5 deletions(-)

diff --git a/dlls/ntdll/tests/virtual.c b/dlls/ntdll/tests/virtual.c
index 14eafe5cd89..d536b4a81af 100644
--- a/dlls/ntdll/tests/virtual.c
+++ b/dlls/ntdll/tests/virtual.c
@@ -13,12 +13,15 @@

static NTSTATUS (WINAPI *pNtAllocateVirtualMemory)(HANDLE, PVOID *, ULONG, SIZE_T *, ULONG, ULONG);
 static NTSTATUS (WINAPI *pNtFreeVirtualMemory)(HANDLE, PVOID *, SIZE_T *, ULONG);
+static BOOL     (WINAPI *pIsWow64Process)(HANDLE, PBOOL);

static void test_AllocateVirtualMemory(void)
 {
     void *addr1, *addr2;
     NTSTATUS status;
     SIZE_T size;
+    ULONG zero_bits;
+    BOOL is_wow64;

/* simple allocation should success */
     size = 0x1000;
@@ -48,13 +51,15 @@ static void test_AllocateVirtualMemory(void)
         ok(status == STATUS_SUCCESS, "pNtFreeVirtualMemory return %08x, addr2: %p\n", status, addr2);
     }

-    /* 21 zero bits is valid */
+    /* 1 zero bits should zero 63-31 upper bits */
     size = 0x1000;
     addr2 = NULL;
-    status = pNtAllocateVirtualMemory(GetCurrentProcess(), &addr2, 21, &size,
-                                      MEM_RESERVE | MEM_COMMIT, PAGE_EXECUTE_READWRITE);
-    ok(status == STATUS_SUCCESS || status == STATUS_NO_MEMORY,
-       "NtAllocateVirtualMemory returned %08x\n", status);
+    zero_bits = 1;
+    status = pNtAllocateVirtualMemory(GetCurrentProcess(), &addr2, 1, &size,
+                                      MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
+    ok((status == STATUS_SUCCESS || broken(status == STATUS_INVALID_PARAMETER_3) /* winxp */) &&
+       ((UINT_PTR)addr2 >> (32 - zero_bits)) == 0,
+       "NtAllocateVirtualMemory returned %08x, addr2: %p\n", status, addr2);
     if (status == STATUS_SUCCESS)
     {
         size = 0;
@@ -62,6 +67,23 @@ static void test_AllocateVirtualMemory(void)
         ok(status == STATUS_SUCCESS, "pNtFreeVirtualMemory return %08x, addr2: %p\n", status, addr2);
     }

+    for (zero_bits = 2; zero_bits <= 21; zero_bits++)
+    {
+        size = 0x1000;
+        addr2 = NULL;
+        status = pNtAllocateVirtualMemory(GetCurrentProcess(), &addr2, zero_bits, &size,
+                                          MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
+        todo_wine_if(zero_bits >= 12)
+        ok((status == STATUS_SUCCESS || status == STATUS_NO_MEMORY) && ((UINT_PTR)addr2 >> (32 - zero_bits)) == 0,
+           "NtAllocateVirtualMemory with %d zero_bits returned %08x, addr2: %p\n", zero_bits, status, addr2);
+        if (status == STATUS_SUCCESS)
+        {
+            size = 0;
+            status = pNtFreeVirtualMemory(GetCurrentProcess(), &addr2, &size, MEM_RELEASE);
+            ok(status == STATUS_SUCCESS, "pNtFreeVirtualMemory return %08x, addr2: %p\n", status, addr2);
+        }
+    }
+
     /* 22 zero bits is invalid */
     size = 0x1000;
     addr2 = NULL;
@@ -75,6 +97,31 @@ static void test_AllocateVirtualMemory(void)
         ok(status == STATUS_SUCCESS, "pNtFreeVirtualMemory return %08x, addr2: %p\n", status, addr2);
     }

+    /* zero bits > 31 should be considered as bitmask on 64bit and WoW64 */
+    size = 0x1000;
+    addr2 = NULL;
+    zero_bits = 0x1fffffff;
+    status = pNtAllocateVirtualMemory(GetCurrentProcess(), &addr2, zero_bits, &size,
+                                      MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
+
+    if (sizeof(void *) == sizeof(int) && (!pIsWow64Process ||
+        !pIsWow64Process( GetCurrentProcess(), &is_wow64 ) || !is_wow64))
+    {
+        ok(status == STATUS_INVALID_PARAMETER_3, "NtAllocateVirtualMemory returned %08x\n", status);
+    }
+    else
+    {
+        todo_wine
+        ok(status == STATUS_SUCCESS && ((UINT_PTR)addr2 & ~zero_bits) == 0,
+           "NtAllocateVirtualMemory returned %08x, addr2: %p\n", status, addr2);
+        if (status == STATUS_SUCCESS)
+        {
+            size = 0;
+            status = pNtFreeVirtualMemory(GetCurrentProcess(), &addr2, &size, MEM_RELEASE);
+            ok(status == STATUS_SUCCESS, "pNtFreeVirtualMemory return %08x, addr2: %p\n", status, addr2);
+        }
+    }
+
     /* AT_ROUND_TO_PAGE flag is not supported for NtAllocateVirtualMemory */
     size = 0x1000;
     addr2 = (char *)addr1 + 0x1000;
@@ -90,11 +137,14 @@ START_TEST(virtual)
 {
     SYSTEM_INFO si;
     HMODULE hntdll;
+    HMODULE hkernel32;

hntdll = GetModuleHandleA("ntdll.dll");
+    hkernel32 = GetModuleHandleA("kernel32.dll");

pNtAllocateVirtualMemory = (void *)GetProcAddress( hntdll, "NtAllocateVirtualMemory" );
     pNtFreeVirtualMemory = (void *)GetProcAddress( hntdll, "NtFreeVirtualMemory" );
+    pIsWow64Process = (void *)GetProcAddress( hkernel32, "IsWow64Process" );

GetSystemInfo(&si);
     trace("system page size %#x\n", si.dwPageSize);