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