OLD | NEW |
| (Empty) |
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 | |
OLD | NEW |