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

sandbox/win/src/win_utils.cc

 // Copyright (c) 2011 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 "sandbox/win/src/win_utils.h"
 
+#include <psapi.h>
 #include <stddef.h>
+#include <stdint.h>
 
 #include <map>
 #include <memory>
+#include <vector>
 
 #include "base/macros.h"
+#include "base/numerics/safe_math.h"
+#include "base/strings/string16.h"
 #include "base/strings/string_util.h"
 #include "base/win/pe_image.h"
 #include "sandbox/win/src/internal_types.h"
 #include "sandbox/win/src/nt_internals.h"
 #include "sandbox/win/src/sandbox_nt_util.h"
 
 namespace {
 
 // Holds the information about a known registry key.
 struct KnownReservedKey {
@@ skipping 74 matching lines @@
 
 const wchar_t kNTDotPrefix[] = L"\\\\.\\";
 const size_t kNTDotPrefixLen = arraysize(kNTDotPrefix) - 1;
 
 // Removes "\\\\.\\" from the path.
 void RemoveImpliedDevice(base::string16* path) {
   if (0 == path->compare(0, kNTDotPrefixLen, kNTDotPrefix))
     *path = path->substr(kNTDotPrefixLen);
 }
 
+// Get the native path to the process.
+bool GetProcessPath(HANDLE process, base::string16* path) {
+  wchar_t process_name[MAX_PATH];
+  DWORD size = MAX_PATH;
+  if (::QueryFullProcessImageNameW(process, PROCESS_NAME_NATIVE, process_name,
+                                   &size)) {
+    *path = process_name;
+    return true;
+  }
+  // Process name is potentially greater than MAX_PATH, try larger max size.
+  std::vector<wchar_t> process_name_buffer(SHRT_MAX);
+  size = SHRT_MAX;
+  if (::QueryFullProcessImageNameW(process, PROCESS_NAME_NATIVE,
+                                   &process_name_buffer[0], &size)) {
+    *path = &process_name_buffer[0];
+    return true;
+  }
+  return false;
+}
+
+// Get the native path for a mapped file.
+bool GetImageFilePath(HANDLE process,
+                      void* base_address,
+                      base::string16* path) {
+  wchar_t mapped_path[MAX_PATH];
+  if (::GetMappedFileNameW(process, base_address, mapped_path, MAX_PATH)) {
+    *path = mapped_path;
+    return true;
+  }
+  // Image name is potentially greater than MAX_PATH, try larger max size.
+  std::vector<wchar_t> mapped_path_buffer(SHRT_MAX);
+  if (::GetMappedFileNameW(process, base_address, &mapped_path_buffer[0],
+                           SHRT_MAX)) {
+    *path = &mapped_path_buffer[0];
+    return true;
+  }
+  return false;
+}
+
 }  // namespace
 
 namespace sandbox {
 
 // Returns true if the provided path points to a pipe.
 bool IsPipe(const base::string16& path) {
   size_t start = 0;
   if (0 == path.compare(0, sandbox::kNTPrefixLen, sandbox::kNTPrefix))
     start = sandbox::kNTPrefixLen;
 
@@ skipping 286 matching lines @@
 
   return ok;
 }
 
 DWORD GetLastErrorFromNtStatus(NTSTATUS status) {
   RtlNtStatusToDosErrorFunction NtStatusToDosError = nullptr;
   ResolveNTFunctionPtr("RtlNtStatusToDosError", &NtStatusToDosError);
   return NtStatusToDosError(status);
 }
 
+// This function walks the virtual memory map using VirtualQueryEx to find
+// the main executable's image section. We attempt to find the first image
+// section which matches the path returned for the process.  This shouldn't
+// be a major performance problem because a new process has a very limited
+// amount of memory allocated so the majority of the valid range should be
+// skipped immediately. However if it turns out to be the case it could be
+// optimized in the specific case of the process being the same as the
+// current process, which due to ASLR rules the image load address will almost
+// always match the current process's load address.
+void* GetProcessBaseAddress(HANDLE process) {
+  MEMORY_BASIC_INFORMATION mem_info = {};
+  // Start 64KiB above zero page.
+  void* current = reinterpret_cast<void*>(0x10000);
+  base::string16 process_path;
+
+  if (!GetProcessPath(process, &process_path))
+    return nullptr;
+
+  // Walk the virtual memory mappings trying to find image sections.
+  // VirtualQueryEx will return false if it encounters a location outside of
+  // the user memory range.
+  while (::VirtualQueryEx(process, current, &mem_info, sizeof(mem_info))) {
+    base::string16 image_path;
+    if (mem_info.Type == MEM_IMAGE &&
+        GetImageFilePath(process, mem_info.BaseAddress, &image_path) &&
+        EqualPath(process_path, image_path)) {
+      return mem_info.BaseAddress;
+    }
+    // VirtualQueryEx should fail before overflow, but just in case we'll check
+    // to prevent an infinite loop.
+    base::CheckedNumeric<uintptr_t> next_base =
+        reinterpret_cast<uintptr_t>(mem_info.BaseAddress);
+    next_base += mem_info.RegionSize;
+    if (!next_base.IsValid())
+      return nullptr;
+    current = reinterpret_cast<void*>(next_base.ValueOrDie());
+  }
+
+  return nullptr;
+}
+
 };  // namespace sandbox
 
 void ResolveNTFunctionPtr(const char* name, void* ptr) {
   static volatile HMODULE ntdll = NULL;
 
   if (!ntdll) {
     HMODULE ntdll_local = ::GetModuleHandle(sandbox::kNtdllName);
     // Use PEImage to sanity-check that we have a valid ntdll handle.
     base::win::PEImage ntdll_peimage(ntdll_local);
     CHECK_NT(ntdll_peimage.VerifyMagic());
     // Race-safe way to set static ntdll.
     ::InterlockedCompareExchangePointer(
         reinterpret_cast<PVOID volatile*>(&ntdll), ntdll_local, NULL);
-
   }
 
   CHECK_NT(ntdll);
   FARPROC* function_ptr = reinterpret_cast<FARPROC*>(ptr);
   *function_ptr = ::GetProcAddress(ntdll, name);
   CHECK_NT(*function_ptr);
 }