Reimplementation of FunctionStub to avoid rewriting potentially executing code for a slight improvement in thread safety.

Make VTABLE patching treadsafe to the extent possible. As-is it's now safe against itself running on other threads at lease, as well as against other similar implementations, though the inherent VM operation race is resolved by retrying.

BUG=27415
TEST=Included unittests.

Review URL: http://codereview.chromium.org/992008

git-svn-id: svn://svn.chromium.org/chrome/trunk/src@42381 0039d316-1c4b-4281-b951-d872f2087c98

1 files changed, 207 insertions, 0 deletions

diff --git a/chrome_frame/function_stub_unittest.cc b/chrome_frame/function_stub_unittest.cc
new file mode 100644
index 0000000..4ad2c31
--- /dev/null
+++ b/chrome_frame/function_stub_unittest.cc
@@ -0,0 +1,207 @@
+// 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"
+
+namespace {
+
+// Test subclass to expose extra stuff.
+class TestFunctionStub: public FunctionStub {
+ public:
+  static void Init(TestFunctionStub* stub) {
+    stub->FunctionStub::Init(&stub->stub_);
+  }
+
+  // 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;
+}
+
+uintptr_t CALLBACK FooDivert(uintptr_t arg) {
+  return kFooRetVal;
+}
+
+}  // namespace
+
+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());
+}
+
+TEST_F(FunctionStubTest, ZeroArgumentStub) {
+  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));
+
+  EXPECT_EQ(kDivertedRetVal, func());
+}
+
+TEST_F(FunctionStubTest, OneArgumentStub) {
+  stub_ = FunctionStub::Create(reinterpret_cast<uintptr_t>(this),
+                               &FunctionStubTest::FooCallback1);
+
+  FuncPtr1 func = reinterpret_cast<FuncPtr1>(stub_->code());
+  EXPECT_CALL(*this, Foo1(kFooArg))
+      .WillOnce(testing::Return(kDivertedRetVal));
+
+  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));
+}
+