chrome_frame/function_stub_unittest.cc

   1 // 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
   6 #include "chrome_frame/function_stub.h"
   7 #include "testing/gtest/include/gtest/gtest.h"
   8 #include "testing/gmock/include/gmock/gmock.h"
   9
  10 namespace {
  11
  12 // Test subclass to expose extra stuff.
  13 class TestFunctionStub: public FunctionStub {
  14  public:
  15   static void Init(TestFunctionStub* stub) {
  16     stub->FunctionStub::Init(&stub->stub_);
  17   }
  18
  19   // Expose the offset to our signature_ field.
  20   static const size_t kSignatureOffset;
  21
  22   void set_signature(HMODULE signature) { signature_ = signature; }
  23 };
  24
  25 const size_t TestFunctionStub::kSignatureOffset =
  26     FIELD_OFFSET(TestFunctionStub, signature_);
  27
  28 class FunctionStubTest: public testing::Test {
  29  public:
  30   FunctionStubTest() : stub_(NULL) {
  31   }
  32
  33   virtual void SetUp() {
  34     SYSTEM_INFO sys_info;
  35     ::GetSystemInfo(&sys_info);
  36
  37     // Playpen size is a system page.
  38     playpen_size_ = sys_info.dwPageSize;
  39
  40     // Reserve two pages.
  41     playpen_ = reinterpret_cast<uint8*>(
  42         ::VirtualAlloc(NULL,
  43                        2 * playpen_size_,
  44                        MEM_RESERVE,
  45                        PAGE_EXECUTE_READWRITE));
  46     ASSERT_TRUE(playpen_ != NULL);
  47
  48     // And commit the first one.
  49     ASSERT_TRUE(::VirtualAlloc(playpen_,
  50                                playpen_size_,
  51                                MEM_COMMIT,
  52                                PAGE_EXECUTE_READWRITE));
  53   }
  54
  55   virtual void TearDown() {
  56     if (stub_ != NULL) {
  57       EXPECT_TRUE(FunctionStub::Destroy(stub_));
  58     }
  59
  60     if (playpen_ != NULL) {
  61       EXPECT_TRUE(::VirtualFree(playpen_, 0, MEM_RELEASE));
  62     }
  63   }
  64
  65  protected:
  66   typedef uintptr_t (CALLBACK *FuncPtr0)();
  67   typedef uintptr_t (CALLBACK *FuncPtr1)(uintptr_t arg);
  68
  69   MOCK_METHOD0(Foo0, uintptr_t());
  70   MOCK_METHOD1(Foo1, uintptr_t(uintptr_t));
  71   MOCK_METHOD0(Bar0, uintptr_t());
  72   MOCK_METHOD1(Bar1, uintptr_t(uintptr_t));
  73
  74   static uintptr_t CALLBACK FooCallback0(FunctionStubTest* test) {
  75     return test->Foo0();
  76   }
  77   static uintptr_t CALLBACK FooCallback1(FunctionStubTest* test,
  78                                          uintptr_t arg) {
  79     return test->Foo1(arg);
  80   }
  81   static uintptr_t CALLBACK BarCallback0(FunctionStubTest* test) {
  82     return test->Foo0();
  83   }
  84   static uintptr_t CALLBACK BarCallback1(FunctionStubTest* test,
  85                                          uintptr_t arg) {
  86     return test->Foo1(arg);
  87   }
  88
  89   // If a stub is allocated during testing, assigning it here
  90   // will deallocate it at the end of test.
  91   FunctionStub* stub_;
  92
  93   // playpen_[0 .. playpen_size_ - 1] is committed, writable memory.
  94   // playpen_[playpen_size_] is uncommitted, defined memory.
  95   uint8* playpen_;
  96   size_t playpen_size_;
  97 };
  98
  99 const uintptr_t kDivertedRetVal = 0x42;
 100 const uintptr_t kFooRetVal = 0xCAFEBABE;
 101 const uintptr_t kFooArg = 0xF0F0F0F0;
 102
 103 uintptr_t CALLBACK Foo() {
 104   return kFooRetVal;
 105 }
 106
 107 uintptr_t CALLBACK FooDivert(uintptr_t arg) {
 108   return kFooRetVal;
 109 }
 110
 111 }  // namespace
 112
 113 TEST_F(FunctionStubTest, Accessors) {
 114   uintptr_t argument = reinterpret_cast<uintptr_t>(this);
 115   uintptr_t dest_fn = reinterpret_cast<uintptr_t>(FooDivert);
 116   stub_ = FunctionStub::Create(argument, FooDivert);
 117
 118   EXPECT_FALSE(stub_->is_bypassed());
 119   EXPECT_TRUE(stub_->is_valid());
 120   EXPECT_TRUE(stub_->code() != NULL);
 121
 122   // Check that the stub code is executable.
 123   MEMORY_BASIC_INFORMATION info = {};
 124   EXPECT_NE(0u, ::VirtualQuery(stub_->code(), &info, sizeof(info)));
 125   const DWORD kExecutableMask = PAGE_EXECUTE | PAGE_EXECUTE_READ |
 126       PAGE_EXECUTE_READWRITE | PAGE_EXECUTE_WRITECOPY;
 127   EXPECT_NE(0u, info.Protect & kExecutableMask);
 128
 129   EXPECT_EQ(argument, stub_->argument());
 130   EXPECT_TRUE(stub_->bypass_address() != NULL);
 131   EXPECT_EQ(dest_fn, stub_->destination_function());
 132 }
 133
 134 TEST_F(FunctionStubTest, ZeroArgumentStub) {
 135   stub_ = FunctionStub::Create(reinterpret_cast<uintptr_t>(this),
 136                                &FunctionStubTest::FooCallback0);
 137
 138   FuncPtr0 func = reinterpret_cast<FuncPtr0>(stub_->code());
 139   EXPECT_CALL(*this, Foo0())
 140       .WillOnce(testing::Return(kDivertedRetVal));
 141
 142   EXPECT_EQ(kDivertedRetVal, func());
 143 }
 144
 145 TEST_F(FunctionStubTest, OneArgumentStub) {
 146   stub_ = FunctionStub::Create(reinterpret_cast<uintptr_t>(this),
 147                                &FunctionStubTest::FooCallback1);
 148
 149   FuncPtr1 func = reinterpret_cast<FuncPtr1>(stub_->code());
 150   EXPECT_CALL(*this, Foo1(kFooArg))
 151       .WillOnce(testing::Return(kDivertedRetVal));
 152
 153   EXPECT_EQ(kDivertedRetVal, func(kFooArg));
 154 }
 155
 156 TEST_F(FunctionStubTest, Bypass) {
 157   stub_ = FunctionStub::Create(reinterpret_cast<uintptr_t>(this),
 158                                &FunctionStubTest::FooCallback0);
 159
 160   FuncPtr0 func = reinterpret_cast<FuncPtr0>(stub_->code());
 161   EXPECT_CALL(*this, Foo0())
 162       .WillOnce(testing::Return(kDivertedRetVal));
 163
 164   // This will call through to foo.
 165   EXPECT_EQ(kDivertedRetVal, func());
 166
 167   // Now bypass to Foo().
 168   stub_->BypassStub(Foo);
 169   EXPECT_TRUE(stub_->is_bypassed());
 170   EXPECT_FALSE(stub_->is_valid());
 171
 172   // We should not call through anymore.
 173   EXPECT_CALL(*this, Foo0())
 174       .Times(0);
 175
 176   EXPECT_EQ(kFooRetVal, func());
 177 }
 178
 179 TEST_F(FunctionStubTest, FromCode) {
 180   // We should get NULL and no crash from reserved memory.
 181   EXPECT_EQ(NULL, FunctionStub::FromCode(playpen_ + playpen_size_));
 182
 183   // Create a FunctionStub pointer whose signature_
 184   // field hangs just off the playpen.
 185   TestFunctionStub* stub =
 186       reinterpret_cast<TestFunctionStub*>(playpen_ + playpen_size_ -
 187                                           TestFunctionStub::kSignatureOffset);
 188   TestFunctionStub::Init(stub);
 189   EXPECT_EQ(NULL, FunctionStub::FromCode(stub));
 190
 191   // Create a stub in committed memory.
 192   stub = reinterpret_cast<TestFunctionStub*>(playpen_);
 193   TestFunctionStub::Init(stub);
 194   // Signature is NULL, which won't do.
 195   EXPECT_EQ(NULL, FunctionStub::FromCode(stub));
 196
 197   const DWORD kFlags = GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS |
 198       GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT;
 199
 200   HMODULE my_module = NULL;
 201   EXPECT_TRUE(::GetModuleHandleEx(kFlags,
 202       reinterpret_cast<const wchar_t*>(&kDivertedRetVal),
 203       &my_module));
 204
 205   // Set our module as signature.
 206   stub->set_signature(my_module);
 207   EXPECT_EQ(stub, FunctionStub::FromCode(stub));
 208 }
 209