Cache ntdll exports in ELF

chrome_elf/chrome_elf_main.cc

Index: chrome_elf/chrome_elf_main.cc
diff --git a/chrome_elf/chrome_elf_main.cc b/chrome_elf/chrome_elf_main.cc
index c6715ba3763b4acb450a946d0ccf0f6871641f2e..465c338d8417dadbf1f70742cebd466edf309e25 100644
--- a/chrome_elf/chrome_elf_main.cc
+++ b/chrome_elf/chrome_elf_main.cc
@@ -6,11 +6,56 @@
 
 #include "chrome_elf/chrome_elf_main.h"
 
+FunctionLookupTable g_ntdll_lookup = FunctionLookupTable();

[robertshield, 2013/12/03 03:26:26]

I don't think you need the extra assignment here, just: 

FunctionLookupTable g_ntdll_lookup;

should be fine.

[Cait (Slow), 2013/12/03 19:59:40]

Done.

+
 void InitChromeElf() {
   // This method is a no-op which may be called to force a load-time dependency
   // on chrome_elf.dll.
 }
 
 BOOL APIENTRY DllMain(HMODULE module, DWORD reason, LPVOID reserved) {
+  if (reason == DLL_PROCESS_ATTACH)
+    Init();
   return TRUE;
 }
+
+void Init() {
+  HMODULE ntdll_handle = ::GetModuleHandle(L"ntdll.dll");
+
+  // To find the Export Address Table address, we start from the DOS header.
+  // The module handle is actually the base address where the header is.

[robertshield, 2013/12/03 03:26:26]

s/base address where the header is/address of the header/

+  IMAGE_DOS_HEADER* dos_header =
+      reinterpret_cast<IMAGE_DOS_HEADER*>(ntdll_handle);
+  // The e_lfanew is an offset from the DOS header to the NT header. It should
+  // never be 0.
+  IMAGE_NT_HEADERS* nt_headers = reinterpret_cast<IMAGE_NT_HEADERS*>(
+      ntdll_handle + dos_header->e_lfanew / sizeof(uintptr_t));
+  // For modules that have an import address table, its offset from the
+  // DOS header is stored in the second data directory's VirtualAddress.
+  if (!nt_headers->OptionalHeader.DataDirectory[0].VirtualAddress)
+    return;
+
+  PBYTE base_addr = reinterpret_cast<PBYTE>(ntdll_handle);

[robertshield, 2013/12/03 03:26:26]

nit: use BYTE* rather than PBYTE.

[Cait (Slow), 2013/12/03 19:59:40]

Done.

+
+  IMAGE_DATA_DIRECTORY* exports_data_dir =
+      &nt_headers->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT];
+
+  IMAGE_EXPORT_DIRECTORY* exports = reinterpret_cast<IMAGE_EXPORT_DIRECTORY*>(
+      base_addr + exports_data_dir->VirtualAddress);
+
+  WORD* ordinals = reinterpret_cast<WORD*>(
+      base_addr + exports->AddressOfNameOrdinals);
+  DWORD* names = reinterpret_cast<DWORD*>(
+      base_addr + exports->AddressOfNames);
+  DWORD* funcs = reinterpret_cast<DWORD*>(
+      base_addr + exports->AddressOfFunctions);
+  int num_entries = exports->NumberOfNames;
+
+  for (int i = 0; i < num_entries; i++) {
+    char* name = reinterpret_cast<char*>(base_addr + names[i]);
+    WORD ord =  ordinals[i];
+    DWORD func = funcs[ord];
+    FARPROC func_addr = reinterpret_cast<FARPROC>(func + base_addr);
+    g_ntdll_lookup[std::string(name)] = func_addr;
+  }
+}