Changes to module_integrity_verifier.cc that allow the function VerifyModule

chrome/browser/safe_browsing/module_integrity_verifier.cc

Index: chrome/browser/safe_browsing/module_integrity_verifier.cc
diff --git a/chrome/browser/safe_browsing/module_integrity_verifier.cc b/chrome/browser/safe_browsing/module_integrity_verifier.cc
index 82b9aca1c0dc193117db9641c95224b45415d15f..906e97a9e3805d094906cfc8e77a645fc6e11882 100644
--- a/chrome/browser/safe_browsing/module_integrity_verifier.cc
+++ b/chrome/browser/safe_browsing/module_integrity_verifier.cc
@@ -4,6 +4,8 @@
 
 #include "chrome/browser/safe_browsing/module_integrity_verifier.h"
 
+#include <set>
+
 #include "base/files/file_path.h"
 #include "base/files/memory_mapped_file.h"
 #include "base/scoped_native_library.h"
@@ -13,41 +15,118 @@ namespace safe_browsing {
 
 namespace {
 
-struct RelocEnumerationState {
-  explicit RelocEnumerationState(HMODULE hModule);
-  ~RelocEnumerationState();
+struct Export {
+  Export(void* addr, std::string name);
+  ~Export();
+
+  bool operator<(const Export& other) const;
+
+  void* addr;
+  std::string name;
+};
+
+Export::Export(void* addr, std::string name) : addr(addr), name(name) {
+}
+
+Export::~Export() {
+}
+
+bool Export::operator<(const Export& other) const {
+  return addr < other.addr;
+}
+
+struct ModuleVerificationState {
+  explicit ModuleVerificationState(HMODULE hModule);
+  ~ModuleVerificationState();
 
   base::win::PEImageAsData disk_peimage;
 
-  // The number of bytes made equivalent between memory and disk due to
-  // relocations.
-  int bytes_corrected_by_reloc;
+  // The module's preferred base address minus the base address it actually
+  // loaded at.
+  uintptr_t image_base_delta;
+
+  // The location of the disk_peimage module's code section minus that of the
+  // mem_peimage module's code section.
+  uintptr_t code_section_delta;
+
+  // The bytes corrected by relocs.
+  base::hash_set<uintptr_t> reloc_addr;
 
   // 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);
+  DISALLOW_COPY_AND_ASSIGN(ModuleVerificationState);
 };
 
-RelocEnumerationState::RelocEnumerationState(HMODULE hModule)
+ModuleVerificationState::ModuleVerificationState(HMODULE hModule)
     : disk_peimage(hModule),
-      bytes_corrected_by_reloc(0),
+      image_base_delta(0),
+      code_section_delta(0),
+      reloc_addr(),
       unknown_reloc_type(false) {
 }
 
-RelocEnumerationState::~RelocEnumerationState() {
+ModuleVerificationState::~ModuleVerificationState() {
+}
+
+bool ByteAccountedForByReloc(uint8_t* byte_addr,
+                             ModuleVerificationState* state) {
+  return ((state->reloc_addr.count(reinterpret_cast<uintptr_t>(byte_addr))) >
+          0);
+}
+
+// Checks each byte in the module's code section again the corresponding byte on
+// disk.  Returns a list of the functions who may have been modified.
+int ExamineBytesDiffInMemory(uint8_t* disk_code_start,
+                             uint8_t* mem_code_start,
+                             uint32_t code_size,
+                             std::vector<Export> exports,
+                             ModuleVerificationState* state,
+                             std::set<std::string>* modified_exports) {
+  int bytes_different = 0;
+  std::vector<Export>::iterator export_it = exports.begin();
+
+  for (uint32_t i = 0; i < code_size; ++i) {
+    if (*(disk_code_start + i) != *(mem_code_start + i) &&
+        !ByteAccountedForByReloc(mem_code_start + i, state)) {
+      // We get the largest export address still smaller than |addr|.  It is
+      // possible that |addr| belongs to some nonexported function located
+      // between this export and the following one.
+      Export addr =
+          Export(reinterpret_cast<void*>(mem_code_start + i), std::string());
+      std::vector<Export>::iterator modified_export_it =
+          std::upper_bound(export_it, exports.end(), addr);
+
+      if (modified_export_it != exports.begin())
+        modified_exports->insert((modified_export_it - 1)->name);
+      ++bytes_different;
+
+      // No later byte can belong to an earlier export.
+      export_it = modified_export_it;
+    }
+  }
+  return bytes_different;
 }
 
-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;
+// Adds to |reloc_addr| the bytes of the pointer at |address| that are corrected
+// by adding |image_base_delta|.
+void AddBytesCorrectedByReloc(ModuleVerificationState* state,
+                              uintptr_t address) {
+  uintptr_t orig_value = *reinterpret_cast<uintptr_t*>(address);
+  uintptr_t fixed_mem_value = orig_value + state->image_base_delta;
+  uintptr_t disk_value =
+      *reinterpret_cast<uintptr_t*>(address + state->code_section_delta);
+
+  for (int i = 0; i < sizeof(address); ++i) {
+    if ((orig_value & 0xFF) != (disk_value & 0xFF) &&
+        (fixed_mem_value & 0xFF) == (disk_value & 0xFF)) {
+      state->reloc_addr.insert(address + i);
+    }
+    orig_value >>= 8;
+    fixed_mem_value >>= 8;
+    disk_value >>= 8;
   }
-  return counter;
 }
 
 bool AddrIsInCodeSection(void* address,
@@ -60,15 +139,14 @@ 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;
+  ModuleVerificationState* state =
+      reinterpret_cast<ModuleVerificationState*>(cookie);
+
   uint8_t* mem_code_addr = NULL;
   uint8_t* disk_code_addr = NULL;
   uint32_t code_size = 0;
   if (!GetCodeAddrsAndSize(mem_peimage,
-                           *disk_peimage_ptr,
+                           state->disk_peimage,
                            &mem_code_addr,
                            &disk_code_addr,
                            &code_size))
@@ -80,21 +158,7 @@ bool EnumRelocsCallback(const base::win::PEImage& mem_peimage,
 
   switch (type) {
     case IMAGE_REL_BASED_HIGHLOW: {
-      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;
-      }
+      AddBytesCorrectedByReloc(state, reinterpret_cast<uintptr_t>(address));
       break;
     }
     case IMAGE_REL_BASED_ABSOLUTE:
@@ -104,13 +168,27 @@ bool EnumRelocsCallback(const base::win::PEImage& mem_peimage,
     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;
+      state->unknown_reloc_type = true;
       break;
     }
   }
   return true;
 }
 
+bool EnumExportsCallback(const base::win::PEImage& mem_peimage,
+                         DWORD ordinal,
+                         DWORD hint,
+                         LPCSTR name,
+                         PVOID function_addr,
+                         LPCSTR forward,
+                         PVOID cookie) {
+  std::vector<Export>* exports = reinterpret_cast<std::vector<Export>*>(cookie);
+  if (name) {
+    exports->push_back(Export(function_addr, std::string(name)));
+  }
+  return true;
+}
+
 }  // namespace
 
 bool GetCodeAddrsAndSize(const base::win::PEImage& mem_peimage,
@@ -140,18 +218,9 @@ bool GetCodeAddrsAndSize(const base::win::PEImage& mem_peimage,
   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) {
+ModuleState VerifyModule(const wchar_t* module_name,
+                         std::set<std::string>* modified_exports) {
+  // Get module handle, load a copy from disk as data and create PEImages.
   HMODULE module_handle = NULL;
   if (!GetModuleHandleEx(0, module_name, &module_handle))
     return MODULE_STATE_UNKNOWN;
@@ -166,31 +235,52 @@ ModuleState VerifyModule(const wchar_t* module_name) {
   base::MemoryMappedFile mapped_module;
   if (!mapped_module.Initialize(base::FilePath(module_path)))
     return MODULE_STATE_UNKNOWN;
-  RelocEnumerationState reloc_enum_state(
+  ModuleVerificationState 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())
+  if (!mem_peimage.VerifyMagic() || !state.disk_peimage.VerifyMagic())
     return MODULE_STATE_UNKNOWN;
 
+  // Get the list of exports.
+  std::vector<Export> exports;
+  mem_peimage.EnumExports(EnumExportsCallback, &exports);
+  std::sort(exports.begin(), exports.end());
+
+  // Get the addresses of the code sections then calculate |code_section_delta|
+  // and |image_base_delta|.
   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,
+                           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)
+  state.code_section_delta = disk_code_addr - mem_code_addr;
+
+  uint8_t* preferred_image_base = reinterpret_cast<uint8_t*>(
+      state.disk_peimage.GetNTHeaders()->OptionalHeader.ImageBase);
+  state.image_base_delta =
+      preferred_image_base - reinterpret_cast<uint8_t*>(mem_peimage.module());
+
+  // Get the relocations.
+  mem_peimage.EnumRelocs(EnumRelocsCallback, &state);
+  if (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)
+  // Count the modified bytes (after accounting for relocs) and get the set of
+  // modified functions.
+  int num_bytes_different = ExamineBytesDiffInMemory(disk_code_addr,
+                                                     mem_code_addr,
+                                                     code_size,
+                                                     exports,
+                                                     &state,
+                                                     modified_exports);
+
+  if (num_bytes_different == 0)
     return MODULE_STATE_UNMODIFIED;
   return MODULE_STATE_MODIFIED;
 }