ntdll: Make the ntdll exports and relocations processing more generic.

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

dlls/ntdll/unix/loader.c

@@ -710,20 +710,6 @@ static NTSTATUS map_so_dll( const IMAGE_NT_HEADERS *nt_descr, HMODULE module )
     return STATUS_SUCCESS;
 }
 
-static const IMAGE_EXPORT_DIRECTORY *get_export_dir( HMODULE module )
-{
-    const IMAGE_DOS_HEADER *dos = (const IMAGE_DOS_HEADER *)module;
-    const IMAGE_NT_HEADERS *nt;
-    DWORD addr;
-
-    if (dos->e_magic != IMAGE_DOS_SIGNATURE) return NULL;
-    nt = (IMAGE_NT_HEADERS *)((const BYTE *)dos + dos->e_lfanew);
-    if (nt->Signature != IMAGE_NT_SIGNATURE) return NULL;
-    addr = nt->OptionalHeader.DataDirectory[IMAGE_FILE_EXPORT_DIRECTORY].VirtualAddress;
-    if (!addr) return NULL;
-    return (IMAGE_EXPORT_DIRECTORY *)((BYTE *)module + addr);
-}
-
 static ULONG_PTR find_ordinal_export( HMODULE module, const IMAGE_EXPORT_DIRECTORY *exports, DWORD ordinal )
 {
     const DWORD *functions = (const DWORD *)((BYTE *)module + exports->AddressOfFunctions);
@@ -771,19 +757,29 @@ static inline void *get_rva( void *module, ULONG_PTR addr )
     return (BYTE *)module + addr;
 }
 
+static const void *get_module_data_dir( HMODULE module, ULONG dir, ULONG *size )
+{
+    const IMAGE_NT_HEADERS *nt = get_rva( module, ((IMAGE_DOS_HEADER *)module)->e_lfanew );
+    const IMAGE_DATA_DIRECTORY *data;
+
+    if (nt->OptionalHeader.Magic == IMAGE_NT_OPTIONAL_HDR64_MAGIC)
+        data = &((const IMAGE_NT_HEADERS64 *)nt)->OptionalHeader.DataDirectory[dir];
+    else if (nt->OptionalHeader.Magic == IMAGE_NT_OPTIONAL_HDR32_MAGIC)
+        data = &((const IMAGE_NT_HEADERS32 *)nt)->OptionalHeader.DataDirectory[dir];
+    else
+        return NULL;
+    if (!data->VirtualAddress || !data->Size) return NULL;
+    if (size) *size = data->Size;
+    return get_rva( module, data->VirtualAddress );
+}
+
 static NTSTATUS fixup_ntdll_imports( const char *name, HMODULE module )
 {
-    const IMAGE_NT_HEADERS *nt;
     const IMAGE_IMPORT_DESCRIPTOR *descr;
-    const IMAGE_DATA_DIRECTORY *dir;
     const IMAGE_THUNK_DATA *import_list;
     IMAGE_THUNK_DATA *thunk_list;
 
-    nt = get_rva( module, ((IMAGE_DOS_HEADER *)module)->e_lfanew );
+    if (!(descr = get_module_data_dir( module, IMAGE_FILE_IMPORT_DIRECTORY, NULL ))) return STATUS_SUCCESS;
-    dir = &nt->OptionalHeader.DataDirectory[IMAGE_FILE_IMPORT_DIRECTORY];
-    if (!dir->VirtualAddress || !dir->Size) return STATUS_SUCCESS;
-
-    descr = get_rva( module, dir->VirtualAddress );
     for (; descr->Name && descr->FirstThunk; descr++)
     {
         thunk_list = get_rva( module, descr->FirstThunk );
@@ -824,7 +820,7 @@ static void load_ntdll_functions( HMODULE module )
 {
     void **ptr;
 
-    ntdll_exports = get_export_dir( module );
+    ntdll_exports = get_module_data_dir( module, IMAGE_FILE_EXPORT_DIRECTORY, NULL );
     assert( ntdll_exports );
 
 #define GET_FUNC(name) \
@@ -852,8 +848,8 @@ static void load_ntdll_functions( HMODULE module )
 }
 
 /* reimplementation of LdrProcessRelocationBlock */
-static IMAGE_BASE_RELOCATION *process_relocation_block( void *module, IMAGE_BASE_RELOCATION *rel,
+static const IMAGE_BASE_RELOCATION *process_relocation_block( void *module, const IMAGE_BASE_RELOCATION *rel,
-                                                        INT_PTR delta )
+                                                              INT_PTR delta )
 {
     char *page = get_rva( module, rel->VirtualAddress );
     UINT count = (rel->SizeOfBlock - sizeof(*rel)) / sizeof(USHORT);
@@ -875,11 +871,9 @@ static IMAGE_BASE_RELOCATION *process_relocation_block( void *module, IMAGE_BASE
         case IMAGE_REL_BASED_HIGHLOW:
             *(int *)(page + offset) += delta;
             break;
-#ifdef _WIN64
         case IMAGE_REL_BASED_DIR64:
-            *(INT_PTR *)(page + offset) += delta;
+            *(INT64 *)(page + offset) += delta;
             break;
-#elif defined(__arm__)
         case IMAGE_REL_BASED_THUMB_MOV32:
         {
             DWORD *inst = (DWORD *)(page + offset);
@@ -897,7 +891,6 @@ static IMAGE_BASE_RELOCATION *process_relocation_block( void *module, IMAGE_BASE
                                                ((hi << 20) & 0x70000000) + ((hi << 16) & 0xff0000);
             break;
         }
-#endif
         default:
             FIXME("Unknown/unsupported relocation %x\n", *relocs);
             return NULL;
@@ -909,18 +902,15 @@ static IMAGE_BASE_RELOCATION *process_relocation_block( void *module, IMAGE_BASE
 
 static void relocate_ntdll( void *module )
 {
-    IMAGE_DOS_HEADER *dos = module;
+    const IMAGE_NT_HEADERS *nt = get_rva( module, ((IMAGE_DOS_HEADER *)module)->e_lfanew );
-    IMAGE_NT_HEADERS *nt = (IMAGE_NT_HEADERS *)((char *)dos + dos->e_lfanew);
+    const IMAGE_BASE_RELOCATION *rel, *end;
-    IMAGE_DATA_DIRECTORY *relocs = &nt->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC];
+    const IMAGE_SECTION_HEADER *sec;
-    IMAGE_BASE_RELOCATION *rel, *end;
+    ULONG protect_old[96], i, size;
-    IMAGE_SECTION_HEADER *sec;
-    ULONG protect_old[96], i;
     INT_PTR delta;
 
-    ERR( "ntdll could not be mapped at preferred address (%p/%p), expect trouble\n",
+    ERR( "ntdll could not be mapped at preferred address (%p), expect trouble\n", module );
-         module, (void *)nt->OptionalHeader.ImageBase );
 
-    if (!relocs->Size || !relocs->VirtualAddress) return;
+    if (!(rel = get_module_data_dir( module, IMAGE_DIRECTORY_ENTRY_BASERELOC, &size ))) return;
 
     sec = (IMAGE_SECTION_HEADER *)((char *)&nt->OptionalHeader + nt->FileHeader.SizeOfOptionalHeader);
     for (i = 0; i < nt->FileHeader.NumberOfSections; i++)
@@ -930,8 +920,7 @@ static void relocate_ntdll( void *module )
         NtProtectVirtualMemory( NtCurrentProcess(), &addr, &size, PAGE_READWRITE, &protect_old[i] );
     }
 
-    rel = get_rva( module, relocs->VirtualAddress );
+    end = (IMAGE_BASE_RELOCATION *)((const char *)rel + size);
-    end = get_rva( module, relocs->VirtualAddress + relocs->Size );
     delta = (char *)module - (char *)nt->OptionalHeader.ImageBase;
     while (rel && rel < end - 1 && rel->SizeOfBlock) rel = process_relocation_block( module, rel, delta );