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
10namespace {
11
12// Test subclass to expose extra stuff.
13class 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
25const size_t TestFunctionStub::kSignatureOffset =
26    FIELD_OFFSET(TestFunctionStub, signature_);
27
28class 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
99const uintptr_t kDivertedRetVal = 0x42;
100const uintptr_t kFooRetVal = 0xCAFEBABE;
101const uintptr_t kFooArg = 0xF0F0F0F0;
102
103uintptr_t CALLBACK Foo() {
104  return kFooRetVal;
105}
106
107uintptr_t CALLBACK FooDivert(uintptr_t arg) {
108  return kFooRetVal;
109}
110
111}  // namespace
112
113TEST_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
134TEST_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
145TEST_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
156TEST_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
179TEST_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