Implement new method to get process base address in Windows sandbox.

sandbox/win/src/target_process.cc

Index: sandbox/win/src/target_process.cc
diff --git a/sandbox/win/src/target_process.cc b/sandbox/win/src/target_process.cc
index 7fddffb4c41d0672a7539562ab41a93cd14ec6a4..72e2780c8c3fff520c93704b4ca4b3c66e1a4283 100644
--- a/sandbox/win/src/target_process.cc
+++ b/sandbox/win/src/target_process.cc
@@ -12,7 +12,6 @@
 
 #include "base/macros.h"
 #include "base/memory/free_deleter.h"
-#include "base/win/pe_image.h"
 #include "base/win/startup_information.h"
 #include "base/win/windows_version.h"
 #include "sandbox/win/src/crosscall_client.h"
@@ -50,26 +49,6 @@ SANDBOX_INTERCEPT HANDLE g_shared_section;
 SANDBOX_INTERCEPT size_t g_shared_IPC_size;
 SANDBOX_INTERCEPT size_t g_shared_policy_size;
 
-// Returns the address of the main exe module in memory taking in account
-// address space layout randomization.
-void* GetBaseAddress(const wchar_t* exe_name, void* entry_point) {
-  HMODULE exe = ::LoadLibrary(exe_name);
-  if (NULL == exe)
-    return exe;
-
-  base::win::PEImage pe(exe);
-  if (!pe.VerifyMagic()) {
-    ::FreeLibrary(exe);
-    return exe;
-  }
-  PIMAGE_NT_HEADERS nt_header = pe.GetNTHeaders();
-  char* base = reinterpret_cast<char*>(entry_point) -
-    nt_header->OptionalHeader.AddressOfEntryPoint;
-
-  ::FreeLibrary(exe);
-  return base;
-}
-
 TargetProcess::TargetProcess(base::win::ScopedHandle initial_token,
                              base::win::ScopedHandle lockdown_token,
                              HANDLE job,
@@ -180,31 +159,20 @@ ResultCode TargetProcess::Create(
     initial_token_.Close();
   }
 
-  CONTEXT context;
-  context.ContextFlags = CONTEXT_ALL;
-  if (!::GetThreadContext(process_info.thread_handle(), &context)) {
-    *win_error = ::GetLastError();
-    ::TerminateProcess(process_info.process_handle(), 0);
-    return SBOX_ERROR_GET_THREAD_CONTEXT;
-  }
-
-#if defined(_WIN64)
-  void* entry_point = reinterpret_cast<void*>(context.Rcx);
-#else
-#pragma warning(push)
-#pragma warning(disable: 4312)
-  // This cast generates a warning because it is 32 bit specific.
-  void* entry_point = reinterpret_cast<void*>(context.Eax);
-#pragma warning(pop)
-#endif  // _WIN64
-
   if (!target_info->DuplicateFrom(process_info)) {
     *win_error = ::GetLastError();  // This may or may not be correct.
     ::TerminateProcess(process_info.process_handle(), 0);
     return SBOX_ERROR_DUPLICATE_TARGET_INFO;
   }
 
-  base_address_ = GetBaseAddress(exe_path, entry_point);
+  base_address_ = GetProcessBaseAddress(process_info.process_handle());
+  DCHECK(base_address_);
+  if (!base_address_) {
+    *win_error = ::GetLastError();
+    ::TerminateProcess(process_info.process_handle(), 0);
+    return SBOX_ERROR_CANNOT_FIND_BASE_ADDRESS;
+  }
+
   sandbox_process_info_.Set(process_info.Take());
   return SBOX_ALL_OK;
 }