| 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 |