1// Copyright 2014 The Chromium Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5#include "chrome/browser/safe_browsing/incident_reporting/module_integrity_verifier_win.h"
6
7#include "base/files/file_path.h"
8#include "base/files/memory_mapped_file.h"
9#include "base/native_library.h"
10#include "base/path_service.h"
11#include "base/scoped_native_library.h"
12#include "base/win/pe_image.h"
13#include "chrome/browser/safe_browsing/incident_reporting/module_integrity_unittest_util_win.h"
14#include "testing/gtest/include/gtest/gtest.h"
15
16namespace safe_browsing {
17
18class SafeBrowsingModuleVerifierWinTest : public testing::Test {
19 protected:
20  void SetUpTestDllAndPEImages() {
21    LoadModule();
22    HMODULE mem_handle;
23    GetMemModuleHandle(&mem_handle);
24    mem_peimage_ptr_.reset(new base::win::PEImage(mem_handle));
25    ASSERT_TRUE(mem_peimage_ptr_->VerifyMagic());
26
27    LoadDLLAsFile();
28    HMODULE disk_handle;
29    GetDiskModuleHandle(&disk_handle);
30    disk_peimage_ptr_.reset(new base::win::PEImageAsData(disk_handle));
31    ASSERT_TRUE(disk_peimage_ptr_->VerifyMagic());
32  }
33
34  void LoadModule() {
35    HMODULE mem_dll_handle =
36        LoadNativeLibrary(base::FilePath(kTestDllNames[0]), NULL);
37    ASSERT_NE(static_cast<HMODULE>(NULL), mem_dll_handle)
38        << "GLE=" << GetLastError();
39    mem_dll_handle_.Reset(mem_dll_handle);
40    ASSERT_TRUE(mem_dll_handle_.is_valid());
41  }
42
43  void GetMemModuleHandle(HMODULE* mem_handle) {
44    *mem_handle = GetModuleHandle(kTestDllNames[0]);
45    ASSERT_NE(static_cast<HMODULE>(NULL), *mem_handle);
46  }
47
48  void LoadDLLAsFile() {
49    // Use the module handle to find the it on disk, then load as a file.
50    HMODULE module_handle;
51    GetMemModuleHandle(&module_handle);
52
53    WCHAR module_path[MAX_PATH] = {};
54    DWORD length =
55        GetModuleFileName(module_handle, module_path, arraysize(module_path));
56    ASSERT_NE(arraysize(module_path), length);
57    ASSERT_TRUE(disk_dll_handle_.Initialize(base::FilePath(module_path)));
58  }
59
60  void GetDiskModuleHandle(HMODULE* disk_handle) {
61    *disk_handle =
62        reinterpret_cast<HMODULE>(const_cast<uint8*>(disk_dll_handle_.data()));
63  }
64
65  // Edits the first byte of the single function exported by the test dll.
66  void EditExport() {
67    HMODULE mem_handle;
68    GetMemModuleHandle(&mem_handle);
69    uint8_t* export_addr =
70        reinterpret_cast<uint8_t*>(GetProcAddress(mem_handle, kTestExportName));
71    EXPECT_NE(reinterpret_cast<uint8_t*>(NULL), export_addr);
72
73    // Edit the first byte of the function.
74    uint8_t new_val = (*export_addr) + 1;
75    SIZE_T bytes_written = 0;
76    WriteProcessMemory(GetCurrentProcess(),
77                       export_addr,
78                       reinterpret_cast<void*>(&new_val),
79                       1,
80                       &bytes_written);
81    EXPECT_EQ(1, bytes_written);
82  }
83
84  base::ScopedNativeLibrary mem_dll_handle_;
85  base::MemoryMappedFile disk_dll_handle_;
86  scoped_ptr<base::win::PEImageAsData> disk_peimage_ptr_;
87  scoped_ptr<base::win::PEImage> mem_peimage_ptr_;
88};
89
90TEST_F(SafeBrowsingModuleVerifierWinTest, VerifyModuleUnmodified) {
91  std::set<std::string> modified_exports;
92  // Call VerifyModule before the module has been loaded, should fail.
93  EXPECT_EQ(MODULE_STATE_UNKNOWN,
94            VerifyModule(kTestDllNames[0], &modified_exports));
95  EXPECT_EQ(0, modified_exports.size());
96
97  // On loading, the module should be identical (up to relocations) in memory as
98  // on disk.
99  SetUpTestDllAndPEImages();
100  EXPECT_EQ(MODULE_STATE_UNMODIFIED,
101            VerifyModule(kTestDllNames[0], &modified_exports));
102  EXPECT_EQ(0, modified_exports.size());
103}
104
105TEST_F(SafeBrowsingModuleVerifierWinTest, VerifyModuleModified) {
106  std::set<std::string> modified_exports;
107  // Confirm the module is identical in memory as on disk before we begin.
108  SetUpTestDllAndPEImages();
109  EXPECT_EQ(MODULE_STATE_UNMODIFIED,
110            VerifyModule(kTestDllNames[0], &modified_exports));
111
112  uint8_t* mem_code_addr = NULL;
113  uint8_t* disk_code_addr = NULL;
114  uint32_t code_size = 0;
115  EXPECT_TRUE(GetCodeAddrsAndSize(*mem_peimage_ptr_,
116                                  *disk_peimage_ptr_,
117                                  &mem_code_addr,
118                                  &disk_code_addr,
119                                  &code_size));
120
121  // Edit the first byte of the code section of the module (this may be before
122  // the address of any export).
123  uint8_t new_val = (*mem_code_addr) + 1;
124  SIZE_T bytes_written = 0;
125  WriteProcessMemory(GetCurrentProcess(),
126                     mem_code_addr,
127                     reinterpret_cast<void*>(&new_val),
128                     1,
129                     &bytes_written);
130  EXPECT_EQ(1, bytes_written);
131
132  // VerifyModule should detect the change.
133  EXPECT_EQ(MODULE_STATE_MODIFIED,
134            VerifyModule(kTestDllNames[0], &modified_exports));
135}
136
137TEST_F(SafeBrowsingModuleVerifierWinTest, VerifyModuleExportModified) {
138  std::set<std::string> modified_exports;
139  // Confirm the module is identical in memory as on disk before we begin.
140  SetUpTestDllAndPEImages();
141  EXPECT_EQ(MODULE_STATE_UNMODIFIED,
142            VerifyModule(kTestDllNames[0], &modified_exports));
143  modified_exports.clear();
144
145  // Edit the exported function, VerifyModule should now return the function
146  // name in modified_exports.
147  EditExport();
148  EXPECT_EQ(MODULE_STATE_MODIFIED,
149            VerifyModule(kTestDllNames[0], &modified_exports));
150  EXPECT_EQ(1, modified_exports.size());
151  EXPECT_EQ(0, std::string(kTestExportName).compare(*modified_exports.begin()));
152}
153
154}  // namespace safe_browsing
155