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

chrome/browser/safe_browsing/module_integrity_verifier_unittest.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 <windows.h>
+#include <psapi.h>
+
+#include "base/files/file_path.h"
+#include "base/files/memory_mapped_file.h"
+#include "base/native_library.h"
+#include "base/path_service.h"
+#include "base/scoped_native_library.h"
+#include "base/win/pe_image.h"
+#include "testing/gtest/include/gtest/gtest.h"
+
+namespace safe_browsing {
+
+const wchar_t kTestDllName[] = L"verifier_test_dll.dll";
+
+class SafeBrowsingModuleVerifierTest : public testing::Test {
+ protected:
+  SafeBrowsingModuleVerifierTest() {}
+  virtual ~SafeBrowsingModuleVerifierTest() {}
+
+  base::ScopedNativeLibrary mem_dll_handle_;
+  base::MemoryMappedFile disk_dll_handle_;
+  scoped_ptr<base::win::PEImageAsData> disk_peimage_ptr_;
+  scoped_ptr<base::win::PEImage> mem_peimage_ptr_;
+
+  void SetUpTestDllAndPEImages() {
+    LoadModule();
+    HMODULE mem_handle;
+    HMODULE disk_handle;
+    GetMemModuleHandle(&mem_handle);
+    GetDiskModuleHandle(&disk_handle);
+
+    disk_peimage_ptr_.reset(new base::win::PEImageAsData(disk_handle));
+    ASSERT_TRUE(disk_peimage_ptr_->VerifyMagic());
+    mem_peimage_ptr_.reset(new base::win::PEImage(mem_handle));
+    ASSERT_TRUE(mem_peimage_ptr_->VerifyMagic());
+  }
+
+  void LoadModule() {
+    base::FilePath current_dir;
+    ASSERT_TRUE(PathService::Get(base::DIR_EXE, &current_dir));
+    mem_dll_handle_.Reset(
+        LoadNativeLibrary(current_dir.Append(kTestDllName), NULL));
+    ASSERT_TRUE(mem_dll_handle_.is_valid());
+  }
+
+  void GetMemModuleHandle(HMODULE* mem_handle) {
+    *mem_handle = GetModuleHandle(kTestDllName);
+    ASSERT_TRUE(NULL != *mem_handle);
+  }
+
+  void GetDiskModuleHandle(HMODULE* disk_handle) {
+    // Use the module handle to find the it on disk, then load as a file.
+    HMODULE module_handle = GetModuleHandle(kTestDllName);
+    ASSERT_TRUE(NULL != module_handle);
+
+    WCHAR module_path[MAX_PATH] = {};
+    DWORD length =
+        GetModuleFileName(module_handle, module_path, arraysize(module_path));
+    ASSERT_NE(length, arraysize(module_path));
+
+    ASSERT_TRUE(disk_dll_handle_.Initialize(base::FilePath(module_path)));
+    *disk_handle =
+        reinterpret_cast<HMODULE>(const_cast<uint8*>(disk_dll_handle_.data()));
+  }
+
+ private:
+  DISALLOW_COPY_AND_ASSIGN(SafeBrowsingModuleVerifierTest);
+};
+
+TEST_F(SafeBrowsingModuleVerifierTest, CountBytesDiffInPtr) {
+  // Construct test pointers and try with CountBytesDiffInPtr.
+  // The first Bytes differ.
+  intptr_t num_a = 0xFF;
+  intptr_t num_b = 0x00;
+
+  // Any inbetween Bytes are identical.
+  int num_bytes_to_add = sizeof(num_a) - 2;
+  for (int i = 0; i < num_bytes_to_add; ++i) {
+    num_a <<= 8;
+    num_b <<= 8;
+    num_a += 0xFF;
+    num_b += 0xFF;
+  }
+
+  // The last Bytes differ.
+  num_a <<= 8;
+  num_b <<= 8;
+  num_a += 0x0F;
+  num_b += 0xFF;
+
+  EXPECT_EQ(2, CountBytesDiffInPtr(num_a, num_b));
+  EXPECT_EQ(2, CountBytesDiffInPtr(num_b, num_a));
+  EXPECT_EQ(0, CountBytesDiffInPtr(num_a, num_a));
+}
+
+TEST_F(SafeBrowsingModuleVerifierTest, VerifyModuleUnmodified) {
+  // Call VerifyModule before the module has been loaded, should fail.
+  EXPECT_EQ(MODULE_STATE_UNKNOWN, VerifyModule(kTestDllName));
+
+  // On loading, the module should be identical (up to relocations) in memory as
+  // on disk.
+  SetUpTestDllAndPEImages();
+  EXPECT_EQ(MODULE_STATE_UNMODIFIED, VerifyModule(kTestDllName));
+}
+
+TEST_F(SafeBrowsingModuleVerifierTest, VerifyModuleModified) {
+  // Confirm the module is identical in memory as on disk before we begin.
+  SetUpTestDllAndPEImages();
+  EXPECT_EQ(MODULE_STATE_UNMODIFIED, VerifyModule(kTestDllName));
+
+  BYTE* mem_code_addr = NULL;
+  BYTE* disk_code_addr = NULL;
+  SIZE_T code_size = 0;
+  EXPECT_TRUE(GetCodeAddrsAndSize(*mem_peimage_ptr_,
+                                  *disk_peimage_ptr_,
+                                  &mem_code_addr,
+                                  &disk_code_addr,
+                                  &code_size));
+
+  // Edit the first byte of the code section of the module.
+  BYTE new_val = (*mem_code_addr) + 1;
+  SIZE_T bytes_written = 0;
+  WriteProcessMemory(GetCurrentProcess(),
+                     mem_code_addr,
+                     reinterpret_cast<void*>(&new_val),
+                     1,
+                     &bytes_written);
+  EXPECT_EQ(1, bytes_written);
+
+  // VerifyModule should detect the change.
+  EXPECT_EQ(MODULE_STATE_MODIFIED, VerifyModule(kTestDllName));
+}
+
+}  // namespace safe_browsing