Reimplementation of FunctionStub, to avoid rewriting potentially executing co...

chrome_frame/function_stub_unittest.cc

Index: chrome_frame/function_stub_unittest.cc
===================================================================
--- chrome_frame/function_stub_unittest.cc	(revision 41717)
+++ chrome_frame/function_stub_unittest.cc	(working copy)
@@ -1,65 +1,207 @@
-// Copyright (c) 2009 The Chromium Authors. All rights reserved.
+// Copyright (c) 2010 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
+
 #include "chrome_frame/function_stub.h"
 #include "testing/gtest/include/gtest/gtest.h"
+#include "testing/gmock/include/gmock/gmock.h"
 
-#define NO_INLINE __declspec(noinline)
-
 namespace {
 
-typedef int (__stdcall* FooPrototype)();
+// Test subclass to expose extra stuff.
+class TestFunctionStub: public FunctionStub {
+ public:
+  static void Init(TestFunctionStub* stub) {
+    stub->FunctionStub::Init(&stub->stub_);
+  }
 
-NO_INLINE int __stdcall Foo() {
-  return 1;
+  // Expose the offset to our signature_ field.
+  static const size_t kSignatureOffset;
+
+  void set_signature(HMODULE signature) { signature_ = signature; }
+};
+
+const size_t TestFunctionStub::kSignatureOffset =
+    FIELD_OFFSET(TestFunctionStub, signature_);
+
+class FunctionStubTest: public testing::Test {
+ public:
+  FunctionStubTest() : stub_(NULL) {
+  }
+
+  virtual void SetUp() {
+    SYSTEM_INFO sys_info;
+    ::GetSystemInfo(&sys_info);
+
+    // Playpen size is a system page.
+    playpen_size_ = sys_info.dwPageSize;
+
+    // Reserve two pages.
+    playpen_ = reinterpret_cast<uint8*>(
+        ::VirtualAlloc(NULL,
+                       2 * playpen_size_,
+                       MEM_RESERVE,
+                       PAGE_EXECUTE_READWRITE));
+    ASSERT_TRUE(playpen_ != NULL);
+
+    // And commit the first one.
+    ASSERT_TRUE(::VirtualAlloc(playpen_,
+                               playpen_size_,
+                               MEM_COMMIT,
+                               PAGE_EXECUTE_READWRITE));
+  }
+
+  virtual void TearDown() {
+    if (stub_ != NULL) {
+      EXPECT_TRUE(FunctionStub::Destroy(stub_));
+    }
+
+    if (playpen_ != NULL) {
+      EXPECT_TRUE(::VirtualFree(playpen_, 0, MEM_RELEASE));
+    }
+  }
+
+ protected:
+  typedef uintptr_t (CALLBACK *FuncPtr0)();
+  typedef uintptr_t (CALLBACK *FuncPtr1)(uintptr_t arg);
+
+  MOCK_METHOD0(Foo0, uintptr_t());
+  MOCK_METHOD1(Foo1, uintptr_t(uintptr_t));
+  MOCK_METHOD0(Bar0, uintptr_t());
+  MOCK_METHOD1(Bar1, uintptr_t(uintptr_t));
+
+  static uintptr_t CALLBACK FooCallback0(FunctionStubTest* test) {
+    return test->Foo0();
+  }
+  static uintptr_t CALLBACK FooCallback1(FunctionStubTest* test, uintptr_t arg) {
+    return test->Foo1(arg);
+  }
+  static uintptr_t CALLBACK BarCallback0(FunctionStubTest* test) {
+    return test->Foo0();
+  }
+  static uintptr_t CALLBACK BarCallback1(FunctionStubTest* test, uintptr_t arg) {
+    return test->Foo1(arg);
+  }
+
+  // If a stub is allocated during testing, assigning it here
+  // will deallocate it at the end of test.
+  FunctionStub* stub_;
+
+  // playpen_[0 .. playpen_size_ - 1] is committed, writable memory.
+  // playpen_[playpen_size_] is uncommitted, defined memory.
+  uint8* playpen_;
+  size_t playpen_size_;
+};
+
+const uintptr_t kDivertedRetVal = 0x42;
+const uintptr_t kFooRetVal = 0xCAFEBABE;
+const uintptr_t kFooArg = 0xF0F0F0F0;
+
+uintptr_t CALLBACK Foo() {
+  return kFooRetVal;
 }
 
-NO_INLINE int __stdcall PatchedFoo(FooPrototype original) {
-  return original() + 1;
+uintptr_t CALLBACK FooDivert(uintptr_t arg) {
+  return kFooRetVal;
 }
 
-}  // end namespace
+}  // namespace
 
-TEST(PatchTests, FunctionStub) {
-  EXPECT_EQ(Foo(), 1);
-  // Create a function stub that calls PatchedFoo and supplies it with
-  // a pointer to Foo.
-  FunctionStub* stub = FunctionStub::Create(reinterpret_cast<uintptr_t>(&Foo),
-                                            &PatchedFoo);
-  EXPECT_TRUE(stub != NULL);
-  // Call the stub as it were Foo().  The call should get forwarded to Foo().
-  FooPrototype patch = reinterpret_cast<FooPrototype>(stub->code());
-  EXPECT_EQ(patch(), 2);
-  // Now neutralize the stub so that it calls Foo() directly without touching
-  // PatchedFoo().
-  // stub->BypassStub(&Foo);
-  stub->BypassStub(reinterpret_cast<void*>(stub->argument()));
-  EXPECT_EQ(patch(), 1);
-  // We're done with the stub.
-  FunctionStub::Destroy(stub);
+TEST_F(FunctionStubTest, Accessors) {
+  uintptr_t argument = reinterpret_cast<uintptr_t>(this);
+  uintptr_t dest_fn = reinterpret_cast<uintptr_t>(FooDivert);
+  stub_ = FunctionStub::Create(argument, FooDivert);
+
+  EXPECT_FALSE(stub_->is_bypassed());
+  EXPECT_TRUE(stub_->is_valid());
+  EXPECT_TRUE(stub_->code() != NULL);
+
+  // Check that the stub code is executable.
+  MEMORY_BASIC_INFORMATION info = {};
+  EXPECT_NE(0, ::VirtualQuery(stub_->code(), &info, sizeof(info)));
+  const DWORD kExecutableMask = PAGE_EXECUTE | PAGE_EXECUTE_READ |
+      PAGE_EXECUTE_READWRITE | PAGE_EXECUTE_WRITECOPY;
+  EXPECT_NE(0, info.Protect & kExecutableMask);
+
+  EXPECT_EQ(argument, stub_->argument());
+  EXPECT_TRUE(stub_->bypass_address() != NULL);
+  EXPECT_EQ(dest_fn, stub_->destination_function());
 }
 
-// Basic tests to check the validity of a stub.
-TEST(PatchTests, FunctionStubCompare) {
-  EXPECT_EQ(Foo(), 1);
+TEST_F(FunctionStubTest, ZeroArgumentStub) {
+  stub_ = FunctionStub::Create(reinterpret_cast<uintptr_t>(this),
+                               &FunctionStubTest::FooCallback0);
 
-  // Detect the absence of a stub
-  FunctionStub* stub = reinterpret_cast<FunctionStub*>(&Foo);
-  EXPECT_FALSE(stub->is_valid());
+  FuncPtr0 func = reinterpret_cast<FuncPtr0>(stub_->code());
+  EXPECT_CALL(*this, Foo0())
+      .WillOnce(testing::Return(kDivertedRetVal));
 
-  stub = FunctionStub::Create(reinterpret_cast<uintptr_t>(&Foo), &PatchedFoo);
-  EXPECT_TRUE(stub != NULL);
-  EXPECT_TRUE(stub->is_valid());
+  EXPECT_EQ(kDivertedRetVal, func());
+}
 
-  FooPrototype patch = reinterpret_cast<FooPrototype>(stub->code());
-  EXPECT_EQ(patch(), 2);
+TEST_F(FunctionStubTest, OneArgumentStub) {
+  stub_ = FunctionStub::Create(reinterpret_cast<uintptr_t>(this),
+                               &FunctionStubTest::FooCallback1);
 
-  // See if we can get the correct absolute pointer to the hook function
-  // back from the stub.
-  EXPECT_EQ(stub->absolute_target(), reinterpret_cast<uintptr_t>(&PatchedFoo));
+  FuncPtr1 func = reinterpret_cast<FuncPtr1>(stub_->code());
+  EXPECT_CALL(*this, Foo1(kFooArg))
+      .WillOnce(testing::Return(kDivertedRetVal));
 
-  // Also verify that the argument being passed to the hook function is indeed
-  // the pointer to the original function (again, absolute not relative).
-  EXPECT_EQ(stub->argument(), reinterpret_cast<uintptr_t>(&Foo));
+  EXPECT_EQ(kDivertedRetVal, func(kFooArg));
 }
+
+TEST_F(FunctionStubTest, Bypass) {
+  stub_ = FunctionStub::Create(reinterpret_cast<uintptr_t>(this),
+                               &FunctionStubTest::FooCallback0);
+
+  FuncPtr0 func = reinterpret_cast<FuncPtr0>(stub_->code());
+  EXPECT_CALL(*this, Foo0())
+      .WillOnce(testing::Return(kDivertedRetVal));
+
+  // This will call through to foo.
+  EXPECT_EQ(kDivertedRetVal, func());
+
+  // Now bypass to Foo().
+  stub_->BypassStub(Foo);
+  EXPECT_TRUE(stub_->is_bypassed());
+  EXPECT_FALSE(stub_->is_valid());
+
+  // We should not call through anymore.
+  EXPECT_CALL(*this, Foo0())
+      .Times(0);
+
+  EXPECT_EQ(kFooRetVal, func());
+}
+
+TEST_F(FunctionStubTest, FromCode) {
+  // We should get NULL and no crash from reserved memory.
+  EXPECT_EQ(NULL, FunctionStub::FromCode(playpen_ + playpen_size_));
+
+  // Create a FunctionStub pointer whose signature_
+  // field hangs just off the playpen.
+  TestFunctionStub* stub =
+      reinterpret_cast<TestFunctionStub*>(playpen_ + playpen_size_ -
+                                          TestFunctionStub::kSignatureOffset);
+  TestFunctionStub::Init(stub);
+  EXPECT_EQ(NULL, FunctionStub::FromCode(stub));
+
+  // Create a stub in committed memory.
+  stub = reinterpret_cast<TestFunctionStub*>(playpen_);
+  TestFunctionStub::Init(stub);
+  // Signature is NULL, which won't do.
+  EXPECT_EQ(NULL, FunctionStub::FromCode(stub));
+
+  const DWORD kFlags = GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS |
+      GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT;
+
+  HMODULE my_module = NULL;
+  EXPECT_TRUE(::GetModuleHandleEx(kFlags,
+      reinterpret_cast<const wchar_t*>(&kDivertedRetVal),
+      &my_module));
+
+  // Set our module as signature.
+  stub->set_signature(my_module);
+  EXPECT_EQ(stub, FunctionStub::FromCode(stub));
+}
+