Cache ntdll proc addresses in ELF

BUG=

Committed: https://src.chromium.org/viewvc/chrome?view=rev&revision=238539

Review URL: https://codereview.chromium.org/85403005

git-svn-id: svn://svn.chromium.org/chrome/trunk/src@239437 0039d316-1c4b-4281-b951-d872f2087c98

6 files changed, 145 insertions, 0 deletions

diff --git a/chrome_elf/chrome_elf.gyp b/chrome_elf/chrome_elf.gyp
index a8f15bd..cf0a0e9 100644
--- a/chrome_elf/chrome_elf.gyp
+++ b/chrome_elf/chrome_elf.gyp
@@ -21,6 +21,9 @@
         'chrome_elf_main.cc',
         'chrome_elf_main.h',
       ],
+      'dependencies': [
+        'chrome_elf_lib',
+      ],
       'msvs_settings': {
         'VCLinkerTool': {
           'BaseAddress': '0x01c20000',
@@ -29,5 +32,32 @@
         },
       },
     },
+    {
+      'target_name': 'chrome_elf_unittests',
+      'type': 'executable',
+      'sources': [
+        'ntdll_cache_unittest.cc',
+      ],
+      'include_dirs': [
+        '..',
+      ],
+      'dependencies': [
+        'chrome_elf_lib',
+        '<(DEPTH)/base/base.gyp:run_all_unittests',
+        '<(DEPTH)/testing/gtest.gyp:gtest',
+      ],
+    },
+    {
+      'target_name': 'chrome_elf_lib',
+      'type': 'static_library',
+      'include_dirs': [
+        '..',
+      ],
+      'sources': [
+        'chrome_elf_types.h',
+        'ntdll_cache.cc',
+        'ntdll_cache.h',
+      ],
+    },
   ],
 }
diff --git a/chrome_elf/chrome_elf_main.cc b/chrome_elf/chrome_elf_main.cc
index c6715ba..4291430 100644
--- a/chrome_elf/chrome_elf_main.cc
+++ b/chrome_elf/chrome_elf_main.cc
@@ -6,11 +6,15 @@
 
 #include "chrome_elf/chrome_elf_main.h"
 
+#include "chrome_elf/ntdll_cache.h"
+
 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)
+    InitCache();
   return TRUE;
 }
diff --git a/chrome_elf/chrome_elf_types.h b/chrome_elf/chrome_elf_types.h
new file mode 100644
index 0000000..9cad485
--- /dev/null
+++ b/chrome_elf/chrome_elf_types.h
@@ -0,0 +1,13 @@
+// Copyright 2013 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef CHROME_ELF_CHROME_ELF_TYPES_H_
+#define CHROME_ELF_CHROME_ELF_TYPES_H_
+
+#include <map>
+#include <string>
+
+typedef std::map<std::string, void*> FunctionLookupTable;
+
+#endif  // CHROME_ELF_CHROME_ELF_TYPES_H_
diff --git a/chrome_elf/ntdll_cache.cc b/chrome_elf/ntdll_cache.cc
new file mode 100644
index 0000000..e550442
--- /dev/null
+++ b/chrome_elf/ntdll_cache.cc
@@ -0,0 +1,51 @@
+// Copyright 2013 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include <stdint.h>
+#include <windows.h>
+
+#include "chrome_elf/ntdll_cache.h"
+
+FunctionLookupTable g_ntdll_lookup;
+
+void InitCache() {
+  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 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(uint32_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;
+
+  BYTE* base_addr = reinterpret_cast<BYTE*>(ntdll_handle);
+
+  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;
+  }
+}
diff --git a/chrome_elf/ntdll_cache.h b/chrome_elf/ntdll_cache.h
new file mode 100644
index 0000000..4608cf19
--- /dev/null
+++ b/chrome_elf/ntdll_cache.h
@@ -0,0 +1,15 @@
+// Copyright 2013 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef CHROME_ELF_NTDLL_CACHE_H_
+#define CHROME_ELF_NTDLL_CACHE_H_
+
+#include "chrome_elf/chrome_elf_types.h"
+
+// Caches the addresses of all functions exported by ntdll in  |g_ntdll_lookup|.
+void InitCache();
+
+extern FunctionLookupTable g_ntdll_lookup;
+
+#endif  // CHROME_ELF_NTDLL_CACHE_H_
diff --git a/chrome_elf/ntdll_cache_unittest.cc b/chrome_elf/ntdll_cache_unittest.cc
new file mode 100644
index 0000000..3ecc086
--- /dev/null
+++ b/chrome_elf/ntdll_cache_unittest.cc
@@ -0,0 +1,32 @@
+// Copyright 2013 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include <windows.h>
+
+#include "base/basictypes.h"
+#include "base/compiler_specific.h"
+#include "chrome_elf/ntdll_cache.h"
+#include "testing/gtest/include/gtest/gtest.h"
+
+namespace {
+
+class NTDLLCacheTest : public testing::Test {
+ protected:
+  virtual void SetUp() OVERRIDE {
+    InitCache();
+  }
+
+};
+
+TEST_F(NTDLLCacheTest, NtDLLCacheSanityCheck) {
+  HMODULE ntdll_handle = ::GetModuleHandle(L"ntdll.dll");
+  // Grab a couple random entries from the cache and make sure they match the
+  // addresses exported by ntdll.
+  EXPECT_EQ(::GetProcAddress(ntdll_handle, "A_SHAFinal"),
+            g_ntdll_lookup["A_SHAFinal"]);
+  EXPECT_EQ(::GetProcAddress(ntdll_handle, "ZwTraceControl"),
+            g_ntdll_lookup["ZwTraceControl"]);
+}
+
+}  // namespace