1 : // Copyright 2012 Google Inc. All Rights Reserved.
2 : //
3 : // Licensed under the Apache License, Version 2.0 (the "License");
4 : // you may not use this file except in compliance with the License.
5 : // You may obtain a copy of the License at
6 : //
7 : // http://www.apache.org/licenses/LICENSE-2.0
8 : //
9 : // Unless required by applicable law or agreed to in writing, software
10 : // distributed under the License is distributed on an "AS IS" BASIS,
11 : // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 : // See the License for the specific language governing permissions and
13 : // limitations under the License.
14 :
15 : #include "syzygy/agent/asan/iat_patcher.h"
16 :
17 : #include <vector>
18 :
19 : #include "base/bind.h"
20 : #include "base/files/file_path.h"
21 : #include "base/win/pe_image.h"
22 : #include "gmock/gmock.h"
23 : #include "gtest/gtest.h"
24 : #include "syzygy/agent/asan/constants.h"
25 : #include "syzygy/core/unittest_util.h"
26 :
27 : namespace agent {
28 : namespace asan {
29 :
30 : namespace {
31 :
32 : class LenientIATPatcherTest : public testing::Test {
33 : public:
34 : using ImportTable = std::vector<FunctionPointer>;
35 :
36 E : LenientIATPatcherTest() : test_dll_(nullptr) {
37 E : }
38 :
39 E : void SetUp() override {
40 : base::FilePath path =
41 E : testing::GetExeRelativePath(L"test_dll.dll");
42 E : test_dll_ = ::LoadLibrary(path.value().c_str());
43 E : ASSERT_NE(nullptr, test_dll_);
44 E : }
45 :
46 E : void TearDown() override {
47 E : if (test_dll_ != nullptr) {
48 E : ::FreeLibrary(test_dll_);
49 E : test_dll_ = nullptr;
50 : }
51 E : }
52 :
53 E : ImportTable GetIAT(HMODULE module) {
54 E : base::win::PEImage image(module);
55 E : ImportTable ret;
56 :
57 E : image.EnumAllImports(OnImport, &ret);
58 :
59 E : return ret;
60 E : }
61 :
62 E : DWORD GetIATPageProtection(HMODULE module) {
63 E : base::win::PEImage image(module);
64 :
65 : const void* iat =
66 E : image.GetImageDirectoryEntryAddr(IMAGE_DIRECTORY_ENTRY_IAT);
67 :
68 E : MEMORY_BASIC_INFORMATION memory_info {};
69 E : EXPECT_TRUE(::VirtualQuery(iat, &memory_info, sizeof(memory_info)));
70 :
71 E : return memory_info.Protect;
72 E : }
73 :
74 E : void ReprotectPage(void* page, DWORD old_prot) {
75 E : DWORD prot = 0;
76 E : ASSERT_TRUE(::VirtualProtect(
77 : page, agent::asan::GetPageSize(), old_prot, &prot));
78 E : }
79 :
80 E : MOCK_METHOD2(OnUnprotect, void(void*, DWORD));
81 :
82 : protected:
83 : static bool OnImport(const base::win::PEImage &image, LPCSTR module,
84 : DWORD ordinal, LPCSTR name, DWORD hint,
85 E : PIMAGE_THUNK_DATA iat, PVOID cookie) {
86 E : ImportTable* imports = reinterpret_cast<ImportTable*>(cookie);
87 E : imports->push_back(reinterpret_cast<FunctionPointer>(iat->u1.Function));
88 :
89 E : return true;
90 E : }
91 :
92 : HMODULE test_dll_;
93 : };
94 : using IATPatcherTest = testing::StrictMock<LenientIATPatcherTest>;
95 :
96 E : static void PatchDestination() {
97 E : }
98 :
99 : } // namespace
100 :
101 E : TEST_F(IATPatcherTest, PatchIATForModule) {
102 : // Capture the IAT of the test module before patching.
103 E : ImportTable iat_before = GetIAT(test_dll_);
104 E : DWORD prot_before = GetIATPageProtection(test_dll_);
105 :
106 : const DWORD kWritableMask =
107 : PAGE_READWRITE | PAGE_WRITECOPY |
108 E : PAGE_EXECUTE_READWRITE | PAGE_EXECUTE_WRITECOPY;
109 :
110 : // We expect the IAT not to be writable.
111 E : ASSERT_EQ(0, prot_before & kWritableMask);
112 :
113 : // None of the imports should point to the dummy destination.
114 E : for (auto fn : iat_before)
115 E : ASSERT_TRUE(PatchDestination != fn);
116 :
117 : // Construct a patch map to patch the named export_dll imports to a dummy
118 : // function.
119 E : IATPatchMap patches;
120 E : patches["function1"] = PatchDestination;
121 E : patches["function3"] = PatchDestination;
122 :
123 : // Patch'er up!
124 E : ASSERT_EQ(PATCH_SUCCEEDED, PatchIATForModule(test_dll_, patches));
125 :
126 : // Make sure the IAT page protections have been reset.
127 E : ASSERT_EQ(prot_before, GetIATPageProtection(test_dll_));
128 :
129 : // Capture the IAT of the test module after patching and verify that the
130 : // expected number of functions got redirected to the dummy destination.
131 E : ImportTable iat_after = GetIAT(test_dll_);
132 E : size_t patched = 0;
133 E : for (auto func : iat_after) {
134 E : if (func == &PatchDestination)
135 E : ++patched;
136 E : }
137 :
138 E : ASSERT_EQ(2, patched);
139 E : }
140 :
141 E : TEST_F(IATPatcherTest, FailsWithAccessViolation) {
142 : // Construct a patch map to patch the named export_dll imports to a dummy
143 : // function.
144 E : IATPatchMap patches;
145 E : patches["function1"] = PatchDestination;
146 E : patches["function3"] = PatchDestination;
147 :
148 : // Create a callback to the mock.
149 : ScopedPageProtections::OnUnprotectCallback on_unprotect =
150 E : base::Bind(&IATPatcherTest::OnUnprotect, base::Unretained(this));
151 :
152 : // Expect a single call to the function to unprotect the IAT. In that call
153 : // reprotect the page.
154 E : EXPECT_CALL(*this, OnUnprotect(testing::_, testing::_)).WillOnce(
155 : testing::Invoke(this, &IATPatcherTest::ReprotectPage));
156 :
157 : // Expect the patching to fail with an access violation, and expect the IAT
158 : // to remain unchanged.
159 E : ImportTable iat_before = GetIAT(test_dll_);
160 E : auto result = PatchIATForModule(test_dll_, patches, on_unprotect);
161 E : ASSERT_NE(0u, PATCH_FAILED_ACCESS_VIOLATION & result);
162 E : ImportTable iat_after = GetIAT(test_dll_);
163 E : EXPECT_EQ(iat_before, iat_after);
164 E : }
165 :
166 : } // namespace asan
167 : } // namespace agent
|