Home | History | Annotate | Download | only in chrome_frame 
      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 
    210