chrome_frame/vtable_patch_manager_unittest.cc

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