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/verifier.h"
+#include "chrome/browser/safe_browsing/module_integrity_verifier.h"
 
 #include <windows.h>
 #include <psapi.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 {
 
 // The number of bytes made to agree (between memory and disk) when the
 // relocations are accounted for.
 int bytes_corrected_by_reloc = 0;

[csharp, 2014/07/29 14:28:15]

This seems to be only used in VerifyModule, so it should probably be a local
variable there.

[krstnmnlsn, 2014/07/29 17:30:25]

It's used in the EnumRelocsCallback (hidden in some collapsed lines below).  I
actually debated squishing it into the cookie with the disk_peimage, but that
seemed excessively ugly.

 
 // This is true if the relocation table contains a reloc of type we don't
 // currently handle.
 bool unknown_reloc_type = false;

[csharp, 2014/07/29 14:28:15]

This probably should also be a local variable in VerifyModule that is passed to
EnumRelocsCallbacks

[krstnmnlsn, 2014/07/29 17:30:25]

Do you think it would be better to make some sort of struct and hand all these
things through the cookie (the cookie is the only way to pass inputs to
EnumRelocsCallbacks)?

 
-bool AddrIsInCodeSection(void* address, BYTE* code_addr, SIZE_T code_size) {
-  return (code_addr <= address && address <= code_addr + code_size);
+int CountBytesDiffInMemory(BYTE* disk_code_start,
+                           BYTE* memory_code_start,
+                           SIZE_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 GetCodeAddrsAndSize(const base::win::PEImage mem_peimage,
-                         base::win::PEImageAsData disk_peimage,
-                         BYTE*& mem_code_addr,
-                         BYTE*& disk_code_addr,
-                         SIZE_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<BYTE*>(
-      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<BYTE*>(disk_peimage.module()) + disk_code_offset;
-  return true;
-}
-
-int CountBytesDiffInPtr(intptr_t num_a, intptr_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;
+bool AddrIsInCodeSection(void* address, BYTE* code_addr, SIZE_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) {
   base::win::PEImageAsData disk_peimage =
       *reinterpret_cast<base::win::PEImageAsData*>(cookie);
   BYTE* mem_code_addr = NULL;
   BYTE* disk_code_addr = NULL;
   SIZE_T code_size = 0;
-  if (!GetCodeAddrsAndSize(
-          mem_peimage, disk_peimage, mem_code_addr, disk_code_addr, code_size))
+  if (!GetCodeAddrsAndSize(mem_peimage, disk_peimage,
+                           &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;
 
       BYTE* preferred_image_base = reinterpret_cast<BYTE*>(
           disk_peimage.GetNTHeaders()->OptionalHeader.ImageBase);
@@ skipping 17 matching lines @@
     default: {
       // TODO(krstnmnlsn): Find a reliable description of the behaviour of the
       // remaining types of relocation and handle them.
       unknown_reloc_type = true;
       break;
     }
   }
   return true;
 }
 
-int CountBytesDiffInMemory(BYTE* disk_code_start,
-                           BYTE* memory_code_start,
-                           SIZE_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;
-}
-
 }  // namespace
 
-ModuleState VerifyModule(wchar_t* module_name) {
-  HMODULE module_handle = GetModuleHandle(module_name);
-  if (module_handle == NULL)
+bool GetCodeAddrsAndSize(const base::win::PEImage& mem_peimage,
+                         const base::win::PEImageAsData& disk_peimage,
+                         BYTE** mem_code_addr,
+                         BYTE** disk_code_addr,
+                         SIZE_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<BYTE*>(
+      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<BYTE*>(disk_peimage.module()) + disk_code_offset;
+  return true;
+}
+
+int CountBytesDiffInPtr(intptr_t num_a, intptr_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) {
+  bytes_corrected_by_reloc = 0;
+  unknown_reloc_type = false;
+
+  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] = {0};
-  if (GetModuleFileNameEx(
-          GetCurrentProcess(), module_handle, module_path, MAX_PATH) == 0)
+  WCHAR module_path[MAX_PATH] = {};
+  DWORD length =
+      GetModuleFileName(module_handle, module_path, arraysize(module_path));
+  if (length == arraysize(module_path))
     return MODULE_STATE_UNKNOWN;
 
   base::MemoryMappedFile mapped_module;
   mapped_module.Initialize(base::FilePath(module_path));
   base::win::PEImageAsData disk_peimage(
       reinterpret_cast<HMODULE>(const_cast<uint8*>(mapped_module.data())));
 
   base::win::PEImage mem_peimage(module_handle);
   if (!mem_peimage.VerifyMagic() || !disk_peimage.VerifyMagic())
     return MODULE_STATE_UNKNOWN;
 
   BYTE* mem_code_addr = NULL;
   BYTE* disk_code_addr = NULL;
   SIZE_T code_size = 0;
-  if (!GetCodeAddrsAndSize(
-          mem_peimage, disk_peimage, mem_code_addr, disk_code_addr, code_size))
+  if (!GetCodeAddrsAndSize(mem_peimage,
+                           disk_peimage,
+                           &mem_code_addr,
+                           &disk_code_addr,
+                           &code_size))
     return MODULE_STATE_UNKNOWN;
 
   int num_bytes_different =
       CountBytesDiffInMemory(disk_code_addr, mem_code_addr, code_size);
-  mem_peimage.EnumRelocs(
-      (base::win::PEImage::EnumRelocsFunction) & EnumRelocsCallback,
-      &disk_peimage);
+  mem_peimage.EnumRelocs(EnumRelocsCallback, &disk_peimage);
   if (unknown_reloc_type)
     return MODULE_STATE_UNKNOWN;
 
-  if (num_bytes_different - bytes_corrected_by_reloc == 0)
+  if (num_bytes_different == bytes_corrected_by_reloc)
     return MODULE_STATE_UNMODIFIED;
   return MODULE_STATE_MODIFIED;
 }
 
 }  // namespace safe_browsing