Add base to the repository.

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

1 files changed, 246 insertions, 0 deletions

diff --git a/base/iat_patch.cc b/base/iat_patch.cc
new file mode 100644
index 0000000..6c4fe1e
--- /dev/null
+++ b/base/iat_patch.cc
@@ -0,0 +1,246 @@
+// Copyright 2008, Google Inc.
+// All rights reserved.
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are
+// met:
+//
+//    * Redistributions of source code must retain the above copyright
+// notice, this list of conditions and the following disclaimer.
+//    * Redistributions in binary form must reproduce the above
+// copyright notice, this list of conditions and the following disclaimer
+// in the documentation and/or other materials provided with the
+// distribution.
+//    * Neither the name of Google Inc. nor the names of its
+// contributors may be used to endorse or promote products derived from
+// this software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+#include "base/iat_patch.h"
+#include "base/logging.h"
+
+namespace iat_patch {
+
+struct InterceptFunctionInformation {
+  bool finished_operation;
+  const char* imported_from_module;
+  const char* function_name;
+  void* new_function;
+  void** old_function;
+  IMAGE_THUNK_DATA** iat_thunk;
+  DWORD return_code;
+};
+
+static void* GetIATFunction(IMAGE_THUNK_DATA* iat_thunk) {
+  if (NULL == iat_thunk) {
+    NOTREACHED();
+    return NULL;
+  }
+
+  // Works around the 64 bit portability warning:
+  // The Function member inside IMAGE_THUNK_DATA is really a pointer
+  // to the IAT function. IMAGE_THUNK_DATA correctly maps to IMAGE_THUNK_DATA32
+  // or IMAGE_THUNK_DATA64 for correct pointer size.
+  union FunctionThunk {
+    IMAGE_THUNK_DATA thunk;
+    void* pointer;
+  } iat_function;
+
+  iat_function.thunk = *iat_thunk;
+  return iat_function.pointer;
+}
+
+static bool InterceptEnumCallback(const PEImage &image, const char* module,
+                                  DWORD ordinal, const char* name, DWORD hint,
+                                  IMAGE_THUNK_DATA* iat, void* cookie) {
+  InterceptFunctionInformation* intercept_information =
+    reinterpret_cast<InterceptFunctionInformation*>(cookie);
+
+  if (NULL == intercept_information) {
+    NOTREACHED();
+    return false;
+  }
+
+  DCHECK(module);
+
+  if ((0 == lstrcmpiA(module, intercept_information->imported_from_module)) &&
+     (NULL != name) &&
+     (0 == lstrcmpiA(name, intercept_information->function_name))) {
+    // Save the old pointer.
+    if (NULL != intercept_information->old_function) {
+      *(intercept_information->old_function) = GetIATFunction(iat);
+    }
+
+    if (NULL != intercept_information->iat_thunk) {
+      *(intercept_information->iat_thunk) = iat;
+    }
+
+    // portability check
+    COMPILE_ASSERT(sizeof(iat->u1.Function) ==
+      sizeof(intercept_information->new_function), unknown_IAT_thunk_format);
+
+    // Patch the function.
+    intercept_information->return_code =
+      ModifyCode(&(iat->u1.Function),
+                 &(intercept_information->new_function),
+                 sizeof(intercept_information->new_function));
+
+    // Terminate further enumeration.
+    intercept_information->finished_operation = true;
+    return false;
+  }
+
+  return true;
+}
+
+DWORD InterceptImportedFunction(HMODULE module_handle,
+                                const char* imported_from_module,
+                                const char* function_name, void* new_function,
+                                void** old_function,
+                                IMAGE_THUNK_DATA** iat_thunk) {
+  if ((NULL == module_handle) || (NULL == imported_from_module) ||
+     (NULL == function_name) || (NULL == new_function)) {
+    NOTREACHED();
+    return ERROR_INVALID_PARAMETER;
+  }
+
+  PEImage target_image(module_handle);
+  if (!target_image.VerifyMagic()) {
+    NOTREACHED();
+    return ERROR_INVALID_PARAMETER;
+  }
+
+  InterceptFunctionInformation intercept_information = {
+    false,
+    imported_from_module,
+    function_name,
+    new_function,
+    old_function,
+    iat_thunk,
+    ERROR_GEN_FAILURE};
+
+  // First go through the IAT. If we don't find the import we are looking
+  // for in IAT, search delay import table.
+  target_image.EnumAllImports(InterceptEnumCallback, &intercept_information);
+  if (!intercept_information.finished_operation) {
+    target_image.EnumAllDelayImports(InterceptEnumCallback,
+                                     &intercept_information);
+  }
+
+  return intercept_information.return_code;
+}
+
+DWORD RestoreImportedFunction(void* intercept_function,
+                              void* original_function,
+                              IMAGE_THUNK_DATA* iat_thunk) {
+  if ((NULL == intercept_function) || (NULL == original_function) ||
+      (NULL == iat_thunk)) {
+    NOTREACHED();
+    return ERROR_INVALID_PARAMETER;
+  }
+
+  if (GetIATFunction(iat_thunk) != intercept_function) {
+    // Check if someone else has intercepted on top of us.
+    // We cannot unpatch in this case, just raise a red flag.
+    NOTREACHED();
+    return ERROR_INVALID_FUNCTION;
+  }
+
+  return ModifyCode(&(iat_thunk->u1.Function),
+                    &original_function,
+                    sizeof(original_function));
+}
+
+DWORD ModifyCode(void* old_code, void* new_code, int length) {
+  if ((NULL == old_code) || (NULL == new_code) || (0 == length)) {
+    NOTREACHED();
+    return ERROR_INVALID_PARAMETER;
+  }
+
+  // Change the page protection so that we can write.
+  DWORD error = NO_ERROR;
+  DWORD old_page_protection = 0;
+  if (VirtualProtect(old_code,
+                     length,
+                     PAGE_READWRITE,
+                     &old_page_protection)) {
+
+    // Write the data.
+    CopyMemory(old_code, new_code, length);
+
+    // Restore the old page protection.
+    error = ERROR_SUCCESS;
+    VirtualProtect(old_code,
+                  length,
+                  old_page_protection,
+                  &old_page_protection);
+  } else {
+    error = GetLastError();
+    NOTREACHED();
+  }
+
+  return error;
+}
+
+IATPatchFunction::IATPatchFunction()
+    : original_function_(NULL),
+      iat_thunk_(NULL),
+      intercept_function_(NULL) {
+}
+
+IATPatchFunction::~IATPatchFunction() {
+  if (NULL != intercept_function_) {
+    DWORD error = Unpatch();
+    DCHECK_EQ(NO_ERROR, error);
+  }
+}
+
+DWORD IATPatchFunction::Patch(HMODULE module_handle,
+                              const char* imported_from_module,
+                              const char* function_name,
+                              void* new_function) {
+  DCHECK_EQ(static_cast<void*>(NULL), original_function_);
+  DCHECK_EQ(static_cast<IMAGE_THUNK_DATA*>(NULL), iat_thunk_);
+  DCHECK_EQ(static_cast<void*>(NULL), intercept_function_);
+
+  DWORD error = InterceptImportedFunction(module_handle,
+                                          imported_from_module,
+                                          function_name,
+                                          new_function,
+                                          &original_function_,
+                                          &iat_thunk_);
+
+  if (NO_ERROR == error) {
+    DCHECK_NE(original_function_, intercept_function_);
+    intercept_function_ = new_function;
+  }
+
+  return error;
+}
+
+DWORD IATPatchFunction::Unpatch() {
+  DWORD error = RestoreImportedFunction(intercept_function_,
+                                        original_function_,
+                                        iat_thunk_);
+
+  if (NO_ERROR == error) {
+    intercept_function_ = NULL;
+    original_function_ = NULL;
+    iat_thunk_ = NULL;
+  }
+
+  return error;
+}
+
+}  // namespace iat_patch