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

chrome_frame/function_stub.cc

+// Copyright (c) 2009 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 <new>
+#include "base/lock.h"
+#include "base/logging.h"
+
+#ifndef _M_IX86
+#error Only x86 supported right now.
+#endif
+
+namespace {
+typedef enum AsmConstants {
+  POP_EAX = 0x58,
+  PUSH_IND = 0x35ff,
+  PUSH_EAX = 0x50,
+  JUMP_IND = 0x25ff,
+};
+
+// A quick and dirty wrapper class that allows us to defer allocating
+// the executable heap until first use, and to release it teardown.
+class ExecutableHeap {
+ public:
+  ExecutableHeap() : heap_(NULL) {
+  }
+
+  ~ExecutableHeap() {
+    if (heap_ != NULL) {
+      BOOL ret = ::HeapDestroy(heap_);
+      heap_ = NULL;
+    }
+  }
+
+  void* Allocate(size_t size) {
+    if (!heap_)
+      CreateHeap();
+
+    DCHECK(heap_);
+
+    return ::HeapAlloc(heap_, 0, size);
+  }
+
+  void Free(void* ptr) {
+    DCHECK(heap_ != NULL);
+    ::HeapFree(heap_, 0, ptr);
+  }
+
+  void CreateHeap() {
+    AutoLock lock(init_lock_);
+
+    if (heap_ == NULL)
+      heap_ = ::HeapCreate(HEAP_CREATE_ENABLE_EXECUTE, 0, 0);
+  }
+
+ private:
+  Lock init_lock_;
+  HANDLE heap_;
+};
+
+// Our executable heap instance, all stubs are allocated from here.
+ExecutableHeap heap_;
+
+}  // namespace
+
+extern "C" IMAGE_DOS_HEADER __ImageBase;
+
+bool FunctionStub::is_valid() const {
+  return signature_ == reinterpret_cast<HMODULE>(&__ImageBase) &&
+      !is_bypassed();
+}
+
+FunctionStub::FunctionStub(uintptr_t extra_argument, void* dest)
+    : signature_(reinterpret_cast<HMODULE>(&__ImageBase)),
+      argument_(extra_argument),
+      destination_function_(reinterpret_cast<uintptr_t>(dest)) {
+  bypass_address_ = reinterpret_cast<uintptr_t>(&stub_.pop_return_addr_);
+  Init(&stub_);
+}
+
+FunctionStub::~FunctionStub() {
+}
+
+void FunctionStub::Init(FunctionStubAsm* stub) {
+  DCHECK(stub != NULL);
+
+  stub->jump_to_bypass_ = JUMP_IND;
+  stub->bypass_target_addr_ = reinterpret_cast<uintptr_t>(&bypass_address_);
+  stub->pop_return_addr_ = POP_EAX;
+  stub->push_ = PUSH_IND;
+  stub->arg_addr_ = reinterpret_cast<uintptr_t>(&argument_);
+  stub->push_return_addr_ = PUSH_EAX;
+  stub->jump_to_target = JUMP_IND;
+  stub->target_addr_ = reinterpret_cast<uintptr_t>(&destination_function_);
+
+  // Flush the instruction cache for the newly written code.
+  BOOL ret = ::FlushInstructionCache(::GetCurrentProcess(),
+                                     stub,
+                                     sizeof(*stub));
+}
+
+void FunctionStub::BypassStub(void* new_target) {
+  set_bypass_address(reinterpret_cast<uintptr_t>(new_target));
+}
+
+FunctionStub* FunctionStub::Create(uintptr_t extra_argument, void* dest) {
+  DCHECK(dest);
+  FunctionStub* stub =
+      reinterpret_cast<FunctionStub*>(heap_.Allocate(sizeof(FunctionStub)));
+
+  if (stub != NULL)
+    new (stub) FunctionStub(extra_argument, dest);
+
+  return stub;
+}
+
+FunctionStub* FunctionStub::FromCode(void* address) {
+  // Address points to arbitrary code here, which may e.g.
+  // lie at the end of an executable segment, which in turn
+  // may terminate earlier than the last address we probe.
+  // We therefore execute under an SEH, so as not to crash
+  // on failed probes.
+  __try {
+    // Retrieve the candidata function stub.
+    FunctionStub* candidate = CONTAINING_RECORD(address, FunctionStub, stub_);
+    if (candidate->stub_.jump_to_bypass_ == JUMP_IND &&
+        candidate->signature_ == reinterpret_cast<HMODULE>(&__ImageBase)) {
+      return candidate;
+    }
+  } __except(EXCEPTION_EXECUTE_HANDLER) {
+  }
+
+  return NULL;
+}
+
+bool FunctionStub::Destroy(FunctionStub* stub) {
+  heap_.Free(stub);
+
+  return true;
+}