ntdll/tests: Add tests for RtlWow64GetCurrentCpuArea().

Signed-off-by: Alexandre Julliard <julliard@winehq.org>

dlls/ntdll/tests/wow64.c

@@ -34,6 +34,7 @@ static NTSTATUS (WINAPI *pNtWow64WriteVirtualMemory64)(HANDLE,ULONG64,const void
 #endif
 
 static BOOL is_wow64;
+static void *code_mem;
 
 static void init(void)
 {
@@ -54,6 +55,8 @@ static void init(void)
     GET_PROC( NtWow64WriteVirtualMemory64 );
 #endif
 #undef GET_PROC
+
+    code_mem = VirtualAlloc( NULL, 65536, MEM_RESERVE | MEM_COMMIT, PAGE_EXECUTE_READWRITE );
 }
 
 static void test_process_architecture( HANDLE process, USHORT expect_machine, USHORT expect_native )
@@ -463,6 +466,52 @@ static void test_cpu_area(void)
 
 #else  /* _WIN64 */
 
+static const BYTE call_func64_code[] =
+{
+    0x58,                               /* pop %eax */
+    0x0e,                               /* push %cs */
+    0x50,                               /* push %eax */
+    0x6a, 0x33,                         /* push $0x33 */
+    0xe8, 0x00, 0x00, 0x00, 0x00,       /* call 1f */
+    0x83, 0x04, 0x24, 0x05,             /* 1: addl $0x5,(%esp) */
+    0xcb,                               /* lret */
+    /* in 64-bit mode: */
+    0x4c, 0x87, 0xf4,                   /* xchg %r14,%rsp */
+    0x55,                               /* push %rbp */
+    0x48, 0x89, 0xe5,                   /* mov %rsp,%rbp */
+    0x56,                               /* push %rsi */
+    0x57,                               /* push %rdi */
+    0x41, 0x8b, 0x4e, 0x10,             /* mov 0x10(%r14),%ecx */
+    0x41, 0x8b, 0x76, 0x14,             /* mov 0x14(%r14),%esi */
+    0x67, 0x8d, 0x04, 0xcd, 0, 0, 0, 0, /* lea 0x0(,%ecx,8),%eax */
+    0x83, 0xf8, 0x20,                   /* cmp $0x20,%eax */
+    0x7d, 0x05,                         /* jge 1f */
+    0xb8, 0x20, 0x00, 0x00, 0x00,       /* mov $0x20,%eax */
+    0x48, 0x29, 0xc4,                   /* 1: sub %rax,%rsp */
+    0x48, 0x83, 0xe4, 0xf0,             /* and $~15,%rsp */
+    0x48, 0x89, 0xe7,                   /* mov %rsp,%rdi */
+    0xf3, 0x48, 0xa5,                   /* rep movsq */
+    0x48, 0x8b, 0x0c, 0x24,             /* mov (%rsp),%rcx */
+    0x48, 0x8b, 0x54, 0x24, 0x08,       /* mov 0x8(%rsp),%rdx */
+    0x4c, 0x8b, 0x44, 0x24, 0x10,       /* mov 0x10(%rsp),%r8 */
+    0x4c, 0x8b, 0x4c, 0x24, 0x18,       /* mov 0x18(%rsp),%r9 */
+    0x41, 0xff, 0x56, 0x08,             /* callq *0x8(%r14) */
+    0x48, 0x8d, 0x65, 0xf0,             /* lea -0x10(%rbp),%rsp */
+    0x5f,                               /* pop %rdi */
+    0x5e,                               /* pop %rsi */
+    0x5d,                               /* pop %rbp */
+    0x4c, 0x87, 0xf4,                   /* xchg %r14,%rsp */
+    0xcb,                               /* lret */
+};
+
+static NTSTATUS call_func64( ULONG64 func64, int nb_args, ULONG64 *args )
+{
+    NTSTATUS (WINAPI *func)( ULONG64 func64, int nb_args, ULONG64 *args ) = code_mem;
+
+    memcpy( code_mem, call_func64_code, sizeof(call_func64_code) );
+    return func( func64, nb_args, args );
+}
+
 static ULONG64 main_module, ntdll_module, wow64_module, wow64cpu_module, wow64win_module;
 
 static void enum_modules64( void (*func)(ULONG64,const WCHAR *) )
@@ -514,6 +563,53 @@ done:
     NtClose( process );
 }
 
+static ULONG64 get_proc_address64( ULONG64 module, const char *name )
+{
+    IMAGE_DOS_HEADER dos;
+    IMAGE_NT_HEADERS64 nt;
+    IMAGE_EXPORT_DIRECTORY exports;
+    ULONG i, *names, *funcs;
+    USHORT *ordinals;
+    NTSTATUS status;
+    HANDLE process;
+    ULONG64 ret = 0;
+    char buffer[64];
+
+    if (!module) return 0;
+    process = OpenProcess( PROCESS_ALL_ACCESS, FALSE, GetCurrentProcessId() );
+    ok( process != 0, "failed to open current process %u\n", GetLastError() );
+    status = pNtWow64ReadVirtualMemory64( process, module, &dos, sizeof(dos), NULL );
+    ok( !status, "NtWow64ReadVirtualMemory64 failed %x\n", status );
+    status = pNtWow64ReadVirtualMemory64( process, module + dos.e_lfanew, &nt, sizeof(nt), NULL );
+    ok( !status, "NtWow64ReadVirtualMemory64 failed %x\n", status );
+    status = pNtWow64ReadVirtualMemory64( process, module + nt.OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress,
+                                          &exports, sizeof(exports), NULL );
+    ok( !status, "NtWow64ReadVirtualMemory64 failed %x\n", status );
+    names = calloc( exports.NumberOfNames, sizeof(*names) );
+    ordinals = calloc( exports.NumberOfNames, sizeof(*ordinals) );
+    funcs = calloc( exports.NumberOfFunctions, sizeof(*funcs) );
+    status = pNtWow64ReadVirtualMemory64( process, module + exports.AddressOfNames,
+                                          names, exports.NumberOfNames * sizeof(*names), NULL );
+    ok( !status, "NtWow64ReadVirtualMemory64 failed %x\n", status );
+    status = pNtWow64ReadVirtualMemory64( process, module + exports.AddressOfNameOrdinals,
+                                          ordinals, exports.NumberOfNames * sizeof(*ordinals), NULL );
+    ok( !status, "NtWow64ReadVirtualMemory64 failed %x\n", status );
+    status = pNtWow64ReadVirtualMemory64( process, module + exports.AddressOfFunctions,
+                                          funcs, exports.NumberOfFunctions * sizeof(*funcs), NULL );
+    ok( !status, "NtWow64ReadVirtualMemory64 failed %x\n", status );
+    for (i = 0; i < exports.NumberOfNames && !ret; i++)
+    {
+        status = pNtWow64ReadVirtualMemory64( process, module + names[i], buffer, sizeof(buffer), NULL );
+        ok( !status, "NtWow64ReadVirtualMemory64 failed %x\n", status );
+        if (!strcmp( buffer, name )) ret = module + funcs[ordinals[i]];
+    }
+    free( funcs );
+    free( ordinals );
+    free( names );
+    NtClose( process );
+    return ret;
+}
+
 static void check_module( ULONG64 base, const WCHAR *name )
 {
     if (base == (ULONG_PTR)GetModuleHandleW(0))
@@ -650,6 +746,35 @@ static void test_nt_wow64(void)
     NtClose( process );
 }
 
+static void test_cpu_area(void)
+{
+    TEB64 *teb64 = (TEB64 *)NtCurrentTeb()->GdiBatchCount;
+    ULONG64 ptr;
+    NTSTATUS status;
+
+    if (!is_wow64) return;
+    if (!ntdll_module) return;
+
+    if ((ptr = get_proc_address64( ntdll_module, "RtlWow64GetCurrentCpuArea" )))
+    {
+        USHORT machine = 0xdead;
+        ULONG64 context, context_ex;
+        ULONG64 args[] = { (ULONG_PTR)&machine, (ULONG_PTR)&context, (ULONG_PTR)&context_ex };
+
+        status = call_func64( ptr, ARRAY_SIZE(args), args );
+        ok( !status, "RtlWow64GetCpuAreaInfo failed %x\n", status );
+        ok( machine == IMAGE_FILE_MACHINE_I386, "wrong machine %x\n", machine );
+        ok( context == teb64->TlsSlots[WOW64_TLS_CPURESERVED] + 4, "wrong context %s / %s\n",
+            wine_dbgstr_longlong(context), wine_dbgstr_longlong(teb64->TlsSlots[WOW64_TLS_CPURESERVED]) );
+        ok( !context_ex, "got context_ex %s\n", wine_dbgstr_longlong(context_ex) );
+        args[0] = args[1] = args[2] = 0;
+        status = call_func64( ptr, ARRAY_SIZE(args), args );
+        ok( !status, "RtlWow64GetCpuAreaInfo failed %x\n", status );
+    }
+    else win_skip( "RtlWow64GetCpuAreaInfo not supported\n" );
+
+}
+
 #endif  /* _WIN64 */
 
 
@@ -658,10 +783,9 @@ START_TEST(wow64)
     init();
     test_query_architectures();
     test_peb_teb();
-#ifdef _WIN64
+#ifndef _WIN64
-    test_cpu_area();
-#else
     test_nt_wow64();
     test_modules();
 #endif
+    test_cpu_area();
 }