Adding the VerifyModule function (and helpers) to safe browsing.

chrome/browser/safe_browsing/module_integrity_verifier.cc

+// Copyright 2014 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 "chrome/browser/safe_browsing/module_integrity_verifier.h"
+
+#include "base/files/file_path.h"
+#include "base/files/memory_mapped_file.h"
+#include "base/scoped_native_library.h"
+#include "base/win/pe_image.h"
+
+namespace safe_browsing {
+
+namespace {
+
+struct RelocEnumerationState {
+  explicit RelocEnumerationState(HMODULE hModule);
+  ~RelocEnumerationState();
+
+  base::win::PEImageAsData disk_peimage;
+
+  // The number of bytes made to agree (between memory and disk) when the

[csharp, 2014/07/30 19:29:42]

This comment confused me a bit, maybe just say "The number of bytes made
equivalent between memory and disk due to relocations."

[krstnmnlsn, 2014/07/30 21:03:58]

Done.

+  // relocations are accounted for.
+  int bytes_corrected_by_reloc;
+
+  // Set true if the relocation table contains a reloc of type that we don't
+  // currently handle.
+  bool unknown_reloc_type;
+
+  DISALLOW_COPY_AND_ASSIGN(RelocEnumerationState);
+};
+
+RelocEnumerationState::RelocEnumerationState(HMODULE hModule)
+    : disk_peimage(hModule),
+      bytes_corrected_by_reloc(0),
+      unknown_reloc_type(false) {
+}
+
+RelocEnumerationState::~RelocEnumerationState() {
+}
+
+int CountBytesDiffInMemory(uint8_t* disk_code_start,
+                           uint8_t* memory_code_start,
+                           uint32_t code_size) {
+  int counter = 0;
+  for (int i = 0; i < static_cast<int>(code_size); ++i) {
+    if (*(disk_code_start + i) != *(memory_code_start + i))
+      ++counter;
+  }
+  return counter;
+}
+
+bool AddrIsInCodeSection(void* address,
+                         uint8_t* code_addr,
+                         uint32_t code_size) {
+  return (code_addr <= address && address < code_addr + code_size);
+}
+
+bool EnumRelocsCallback(const base::win::PEImage& mem_peimage,
+                        WORD type,
+                        void* address,
+                        void* cookie) {
+  RelocEnumerationState* reloc_enum_state =
+      reinterpret_cast<RelocEnumerationState*>(cookie);
+  const base::win::PEImageAsData* disk_peimage_ptr =
+      &reloc_enum_state->disk_peimage;
+  uint8_t* mem_code_addr = NULL;
+  uint8_t* disk_code_addr = NULL;
+  uint32_t code_size = 0;
+  if (!GetCodeAddrsAndSize(mem_peimage,
+                           *disk_peimage_ptr,
+                           &mem_code_addr,
+                           &disk_code_addr,
+                           &code_size))
+    return false;
+
+  switch (type) {
+    case IMAGE_REL_BASED_HIGHLOW: {
+      // If not in the code section return true to continue to the next reloc.
+      if (!AddrIsInCodeSection(address, mem_code_addr, code_size))
+        return true;

[csharp, 2014/07/30 19:29:42]

Does it make sense to pull this check out above the switch? (would the other
relocs also only care if the change was in code)

[krstnmnlsn, 2014/07/30 21:03:58]

They would care yes.  Thanks!

+
+      uint8_t* preferred_image_base = reinterpret_cast<uint8_t*>(
+          disk_peimage_ptr->GetNTHeaders()->OptionalHeader.ImageBase);
+      uintptr_t delta = preferred_image_base -
+                        reinterpret_cast<uint8_t*>(mem_peimage.module());
+      uint8_t* new_value = (*reinterpret_cast<uint8_t**>(address)) + delta;
+      int bytes_corrected =
+          CountBytesDiffInPtr(reinterpret_cast<uintptr_t>(new_value),
+                              *reinterpret_cast<uintptr_t*>(address));
+
+      // Check that the adding delta corrected the value to agree with disk.
+      uint8_t** disk_address = reinterpret_cast<uint8_t**>(
+          reinterpret_cast<uint8_t*>(address) - mem_code_addr + disk_code_addr);
+      if (new_value == *disk_address) {
+        reloc_enum_state->bytes_corrected_by_reloc += bytes_corrected;
+      }
+      break;
+    }
+    case IMAGE_REL_BASED_ABSOLUTE:
+      // Absolute type relocations are a noop, sometimes used to pad a section
+      // of relocations.
+      break;
+    default: {
+      // TODO(krstnmnlsn): Find a reliable description of the behaviour of the
+      // remaining types of relocation and handle them.
+      reloc_enum_state->unknown_reloc_type = true;
+      break;
+    }
+  }
+  return true;
+}
+
+}  // namespace
+
+bool GetCodeAddrsAndSize(const base::win::PEImage& mem_peimage,
+                         const base::win::PEImageAsData& disk_peimage,
+                         uint8_t** mem_code_addr,
+                         uint8_t** disk_code_addr,
+                         uint32_t* code_size) {
+  DWORD base_of_code = mem_peimage.GetNTHeaders()->OptionalHeader.BaseOfCode;
+
+  // Get the address and size of the code section in the loaded module image.
+  PIMAGE_SECTION_HEADER mem_code_header =
+      mem_peimage.GetImageSectionFromAddr(mem_peimage.RVAToAddr(base_of_code));
+  if (mem_code_header == NULL)
+    return false;
+  *mem_code_addr = reinterpret_cast<uint8_t*>(
+      mem_peimage.RVAToAddr(mem_code_header->VirtualAddress));
+  *code_size = mem_code_header->Misc.VirtualSize;
+
+  // Get the address of the code section in the module mapped as data from disk.
+  // Section size will be the same as in the loaded module.
+  DWORD disk_code_offset = 0;
+  if (!mem_peimage.ImageAddrToOnDiskOffset(
+          reinterpret_cast<void*>(*mem_code_addr), &disk_code_offset))
+    return false;
+  *disk_code_addr =
+      reinterpret_cast<uint8_t*>(disk_peimage.module()) + disk_code_offset;
+  return true;
+}
+
+int CountBytesDiffInPtr(uintptr_t num_a, uintptr_t num_b) {
+  int num_bytes = 0;
+  for (int i = 0; i < sizeof(num_a); ++i) {
+    if ((num_a & 0xFF) != (num_b & 0xFF))
+      ++num_bytes;
+    num_a >>= 8;
+    num_b >>= 8;
+  }
+  return num_bytes;
+}
+
+ModuleState VerifyModule(const wchar_t* module_name) {
+  HMODULE module_handle = NULL;
+  if (!GetModuleHandleEx(0, module_name, &module_handle))
+    return MODULE_STATE_UNKNOWN;
+  base::ScopedNativeLibrary native_library(module_handle);
+
+  WCHAR module_path[MAX_PATH] = {};
+  DWORD length =
+      GetModuleFileName(module_handle, module_path, arraysize(module_path));
+  if (!length || length == arraysize(module_path))
+    return MODULE_STATE_UNKNOWN;
+
+  base::MemoryMappedFile mapped_module;
+  if (!mapped_module.Initialize(base::FilePath(module_path)))
+    return MODULE_STATE_UNKNOWN;
+  RelocEnumerationState reloc_enum_state(
+      reinterpret_cast<HMODULE>(const_cast<uint8*>(mapped_module.data())));
+
+  base::win::PEImage mem_peimage(module_handle);
+  if (!mem_peimage.VerifyMagic() ||
+      !reloc_enum_state.disk_peimage.VerifyMagic())
+    return MODULE_STATE_UNKNOWN;
+
+  uint8_t* mem_code_addr = NULL;
+  uint8_t* disk_code_addr = NULL;
+  uint32_t code_size = 0;
+  if (!GetCodeAddrsAndSize(mem_peimage,
+                           reloc_enum_state.disk_peimage,
+                           &mem_code_addr,
+                           &disk_code_addr,
+                           &code_size))
+    return MODULE_STATE_UNKNOWN;
+
+  mem_peimage.EnumRelocs(EnumRelocsCallback, &reloc_enum_state);
+  if (reloc_enum_state.unknown_reloc_type)
+    return MODULE_STATE_UNKNOWN;
+
+  int num_bytes_different =
+      CountBytesDiffInMemory(disk_code_addr, mem_code_addr, code_size);
+  if (num_bytes_different == reloc_enum_state.bytes_corrected_by_reloc)
+    return MODULE_STATE_UNMODIFIED;
+  return MODULE_STATE_MODIFIED;
+}
+
+}  // namespace safe_browsing