blob: 4e04796f09a7422750d6036087d2b98a51c3e388 [file] [log] [blame]
[email protected]d999c7c2010-03-23 21:09:021// Copyright (c) 2010 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_frame/vtable_patch_manager.h"
6#include <unknwn.h>
[email protected]6d11ff82010-05-10 20:46:547#include "base/message_loop.h"
[email protected]34b99632011-01-01 01:01:068#include "base/threading/thread.h"
[email protected]d999c7c2010-03-23 21:09:029#include "base/scoped_handle.h"
10#include "gtest/gtest.h"
11#include "gmock/gmock.h"
12
13namespace {
14// GMock names we use.
15using testing::_;
16using testing::Return;
17
18class MockClassFactory : public IClassFactory {
19 public:
20 MOCK_METHOD2_WITH_CALLTYPE(__stdcall, QueryInterface,
21 HRESULT(REFIID riid, void **object));
22 MOCK_METHOD0_WITH_CALLTYPE(__stdcall, AddRef, ULONG());
23 MOCK_METHOD0_WITH_CALLTYPE(__stdcall, Release, ULONG());
24 MOCK_METHOD3_WITH_CALLTYPE(__stdcall, CreateInstance,
25 HRESULT (IUnknown *outer, REFIID riid, void **object));
26 MOCK_METHOD1_WITH_CALLTYPE(__stdcall, LockServer, HRESULT(BOOL lock));
27};
28
29// Retrieve the vtable for an interface.
30void* GetVtable(IUnknown* unk) {
31 return *reinterpret_cast<void**>(unk);
32}
33
34// Forward decl.
35extern vtable_patch::MethodPatchInfo IClassFactory_PatchInfo[];
36
37class VtablePatchManagerTest: public testing::Test {
38 public:
39 VtablePatchManagerTest() {
40 EXPECT_TRUE(current_ == NULL);
41 current_ = this;
42 }
43
44 ~VtablePatchManagerTest() {
45 EXPECT_TRUE(current_ == this);
46 current_ = NULL;
47 }
48
49 virtual void SetUp() {
50 // Make a backup of the test vtable and it's page protection settings.
51 void* vtable = GetVtable(&factory_);
52 MEMORY_BASIC_INFORMATION info;
53 ASSERT_TRUE(::VirtualQuery(vtable, &info, sizeof(info)));
54 vtable_protection_ = info.Protect;
55 memcpy(vtable_backup_, vtable, sizeof(vtable_backup_));
56 }
57
58 virtual void TearDown() {
59 // Unpatch to make sure we've restored state for subsequent test.
60 UnpatchInterfaceMethods(IClassFactory_PatchInfo);
61
62 // Restore the test vtable and its page protection settings.
63 void* vtable = GetVtable(&factory_);
64 DWORD old_protect = 0;
65 EXPECT_TRUE(::VirtualProtect(vtable, sizeof(vtable_backup_),
66 PAGE_EXECUTE_WRITECOPY, &old_protect));
67 memcpy(vtable, vtable_backup_, sizeof(vtable_backup_));
68 EXPECT_TRUE(::VirtualProtect(vtable, sizeof(vtable_backup_),
69 vtable_protection_, &old_protect));
70 }
71
72 typedef HRESULT (__stdcall* LockServerFun)(IClassFactory* self, BOOL lock);
73 MOCK_METHOD3(LockServerPatch,
74 HRESULT(LockServerFun old_fun, IClassFactory* self, BOOL lock));
75
76 static HRESULT STDMETHODCALLTYPE LockServerPatchCallback(
77 LockServerFun fun, IClassFactory* self, BOOL lock) {
78 EXPECT_TRUE(current_ != NULL);
79 if (current_ != NULL)
80 return current_->LockServerPatch(fun, self, lock);
81 else
82 return E_UNEXPECTED;
83 }
84
85 protected:
86 // Number of functions in the IClassFactory vtable.
87 static const size_t kFunctionCount = 5;
88
89 // Backup of the factory_ vtable as we found it at Setup.
90 PROC vtable_backup_[kFunctionCount];
91 // VirtualProtect flags on the factory_ vtable as we found it at Setup.
92 DWORD vtable_protection_;
93
94 // The mock factory class we patch.
95 MockClassFactory factory_;
96
97 // Current test running for routing the patch callback function.
98 static VtablePatchManagerTest* current_;
99};
100
101VtablePatchManagerTest* VtablePatchManagerTest::current_ = NULL;
102
103BEGIN_VTABLE_PATCHES(IClassFactory)
104 VTABLE_PATCH_ENTRY(4, &VtablePatchManagerTest::LockServerPatchCallback)
105END_VTABLE_PATCHES();
106
107} // namespace
108
109TEST_F(VtablePatchManagerTest, ReplacePointer) {
110 void* const kFunctionOriginal = reinterpret_cast<void*>(0xCAFEBABE);
111 void* const kFunctionFoo = reinterpret_cast<void*>(0xF0F0F0F0);
112 void* const kFunctionBar = reinterpret_cast<void*>(0xBABABABA);
113
114 using vtable_patch::internal::ReplaceFunctionPointer;
115 // Replacing a non-writable location should fail, but not crash.
116 EXPECT_FALSE(ReplaceFunctionPointer(NULL, kFunctionBar, kFunctionFoo));
117
118 void* foo_entry = kFunctionOriginal;
119 // Replacing with the wrong original function should
120 // fail and not change the entry.
121 EXPECT_FALSE(ReplaceFunctionPointer(&foo_entry, kFunctionBar, kFunctionFoo));
122 EXPECT_EQ(foo_entry, kFunctionOriginal);
123
124 // Replacing with the correct original should succeed.
125 EXPECT_TRUE(ReplaceFunctionPointer(&foo_entry,
126 kFunctionBar,
127 kFunctionOriginal));
128 EXPECT_EQ(foo_entry, kFunctionBar);
129}
130
131TEST_F(VtablePatchManagerTest, PatchInterfaceMethods) {
132 // Unpatched.
133 EXPECT_CALL(factory_, LockServer(TRUE))
134 .WillOnce(Return(E_FAIL));
135 EXPECT_EQ(E_FAIL, factory_.LockServer(TRUE));
136
137 EXPECT_HRESULT_SUCCEEDED(
138 PatchInterfaceMethods(&factory_, IClassFactory_PatchInfo));
139
140 EXPECT_NE(0, memcmp(GetVtable(&factory_),
141 vtable_backup_,
142 sizeof(vtable_backup_)));
143
144 // This should not be called while the patch is in effect.
145 EXPECT_CALL(factory_, LockServer(_))
146 .Times(0);
147
148 EXPECT_CALL(*this, LockServerPatch(testing::_, &factory_, TRUE))
149 .WillOnce(testing::Return(S_FALSE));
150
151 EXPECT_EQ(S_FALSE, factory_.LockServer(TRUE));
152}
153
154TEST_F(VtablePatchManagerTest, UnpatchInterfaceMethods) {
155 // Patch it.
156 EXPECT_HRESULT_SUCCEEDED(
157 PatchInterfaceMethods(&factory_, IClassFactory_PatchInfo));
158
159 EXPECT_NE(0, memcmp(GetVtable(&factory_),
160 vtable_backup_,
161 sizeof(vtable_backup_)));
162
163 // This should not be called while the patch is in effect.
164 EXPECT_CALL(factory_, LockServer(testing::_))
165 .Times(0);
166
167 EXPECT_CALL(*this, LockServerPatch(testing::_, &factory_, TRUE))
168 .WillOnce(testing::Return(S_FALSE));
169
170 EXPECT_EQ(S_FALSE, factory_.LockServer(TRUE));
171
172 // Now unpatch.
173 EXPECT_HRESULT_SUCCEEDED(
174 UnpatchInterfaceMethods(IClassFactory_PatchInfo));
175
176 // And check that the call comes through correctly.
177 EXPECT_CALL(factory_, LockServer(FALSE))
178 .WillOnce(testing::Return(E_FAIL));
179 EXPECT_EQ(E_FAIL, factory_.LockServer(FALSE));
180}
181
182TEST_F(VtablePatchManagerTest, DoublePatch) {
183 // Patch it.
184 EXPECT_HRESULT_SUCCEEDED(
185 PatchInterfaceMethods(&factory_, IClassFactory_PatchInfo));
186
187 // Capture the VTable after patching.
188 PROC vtable[kFunctionCount];
189 memcpy(vtable, GetVtable(&factory_), sizeof(vtable));
190
191 // Patch it again, this should be idempotent.
192 EXPECT_HRESULT_SUCCEEDED(
193 PatchInterfaceMethods(&factory_, IClassFactory_PatchInfo));
194
195 // Should not have changed the VTable on second call.
196 EXPECT_EQ(0, memcmp(vtable, GetVtable(&factory_), sizeof(vtable)));
197}
198
199namespace vtable_patch {
200// Expose internal implementation detail, purely for testing.
201extern Lock patch_lock_;
202
203} // namespace vtable_patch
204
205TEST_F(VtablePatchManagerTest, ThreadSafePatching) {
206 // It's difficult to test for threadsafe patching, but as a close proxy,
207 // test for no patching happening from a background thread while the patch
208 // lock is held.
209 base::Thread background("Background Test Thread");
210
211 EXPECT_TRUE(background.Start());
212 ScopedHandle event(::CreateEvent(NULL, TRUE, FALSE, NULL));
213
214 // Grab the patch lock.
215 vtable_patch::patch_lock_.Acquire();
216
217 // Instruct the background thread to patch factory_.
218 background.message_loop()->PostTask(FROM_HERE,
219 NewRunnableFunction(&vtable_patch::PatchInterfaceMethods,
220 &factory_,
221 &IClassFactory_PatchInfo[0]));
222
223 // And subsequently to signal the event. Neither of these actions should
224 // occur until we've released the patch lock.
225 background.message_loop()->PostTask(FROM_HERE,
226 NewRunnableFunction(::SetEvent, event.Get()));
227
228 // Wait for a little while, to give the background thread time to process.
229 // We expect this wait to time out, as the background thread should end up
230 // blocking on the patch lock.
231 EXPECT_EQ(WAIT_TIMEOUT, ::WaitForSingleObject(event.Get(), 50));
232
233 // Verify that patching did not take place yet.
234 EXPECT_CALL(factory_, LockServer(TRUE))
235 .WillOnce(Return(S_FALSE));
236 EXPECT_EQ(S_FALSE, factory_.LockServer(TRUE));
237
238 // Release the lock and wait on the event again to ensure
239 // the patching has taken place now.
240 vtable_patch::patch_lock_.Release();
241 EXPECT_EQ(WAIT_OBJECT_0, ::WaitForSingleObject(event.Get(), INFINITE));
242
243 // We should not get called here anymore.
244 EXPECT_CALL(factory_, LockServer(TRUE))
245 .Times(0);
246
247 // But should be diverted here.
248 EXPECT_CALL(*this, LockServerPatch(_, &factory_, TRUE))
249 .WillOnce(Return(S_FALSE));
250 EXPECT_EQ(S_FALSE, factory_.LockServer(TRUE));
251
252 // Same deal for unpatching.
253 ::ResetEvent(event.Get());
254
255 // Grab the patch lock.
256 vtable_patch::patch_lock_.Acquire();
257
258 // Instruct the background thread to unpatch.
259 background.message_loop()->PostTask(FROM_HERE,
260 NewRunnableFunction(&vtable_patch::UnpatchInterfaceMethods,
261 &IClassFactory_PatchInfo[0]));
262
263 // And subsequently to signal the event. Neither of these actions should
264 // occur until we've released the patch lock.
265 background.message_loop()->PostTask(FROM_HERE,
266 NewRunnableFunction(::SetEvent, event.Get()));
267
268 // Wait for a little while, to give the background thread time to process.
269 // We expect this wait to time out, as the background thread should end up
270 // blocking on the patch lock.
271 EXPECT_EQ(WAIT_TIMEOUT, ::WaitForSingleObject(event.Get(), 50));
272
273 // We should still be patched.
274 EXPECT_CALL(factory_, LockServer(TRUE))
275 .Times(0);
276 EXPECT_CALL(*this, LockServerPatch(_, &factory_, TRUE))
277 .WillOnce(Return(S_FALSE));
278 EXPECT_EQ(S_FALSE, factory_.LockServer(TRUE));
279
280 // Release the patch lock and wait on the event.
281 vtable_patch::patch_lock_.Release();
282 EXPECT_EQ(WAIT_OBJECT_0, ::WaitForSingleObject(event.Get(), INFINITE));
283
284 // Verify that unpatching took place.
285 EXPECT_CALL(factory_, LockServer(TRUE))
286 .WillOnce(Return(S_FALSE));
287 EXPECT_EQ(S_FALSE, factory_.LockServer(TRUE));
288}