1 files changed, 73 insertions, 0 deletions
diff --git a/app/win/iat_patch_function.h b/app/win/iat_patch_function.h
new file mode 100644
index 0000000..37961dd
--- /dev/null
+++ b/app/win/iat_patch_function.h
@@ -0,0 +1,73 @@
+// 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.
+
+#ifndef APP_WIN_IAT_PATCH_FUNCTION_H_
+#define APP_WIN_IAT_PATCH_FUNCTION_H_
+#pragma once
+
+#include <windows.h>
+
+#include "base/basictypes.h"
+#include "base/pe_image.h"
+
+namespace app {
+namespace win {
+
+// A class that encapsulates Import Address Table patching helpers and restores
+// the original function in the destructor.
+//
+// It will intercept functions for a specific DLL imported from another DLL.
+// This is the case when, for example, we want to intercept
+// CertDuplicateCertificateContext function (exported from crypt32.dll) called
+// by wininet.dll.
+class IATPatchFunction {
+ public:
+  IATPatchFunction();
+  ~IATPatchFunction();
+
+  // Intercept a function in an import table of a specific
+  // module. Save the original function and the import
+  // table address. These values will be used later
+  // during Unpatch
+  //
+  // Arguments:
+  // module                 Module to be intercepted
+  // imported_from_module   Module that exports the 'function_name'
+  // function_name          Name of the API to be intercepted
+  //
+  // Returns: Windows error code (winerror.h). NO_ERROR if successful
+  //
+  // Note: Patching a function will make the IAT patch take some "ownership" on
+  // |module|.  It will LoadLibrary(module) to keep the DLL alive until a call
+  // to Unpatch(), which will call FreeLibrary() and allow the module to be
+  // unloaded.  The idea is to help prevent the DLL from going away while a
+  // patch is still active.
+  DWORD Patch(const wchar_t* module,
+              const char* imported_from_module,
+              const char* function_name,
+              void* new_function);
+
+  // Unpatch the IAT entry using internally saved original
+  // function.
+  //
+  // Returns: Windows error code (winerror.h). NO_ERROR if successful
+  DWORD Unpatch();
+
+  bool is_patched() const {
+    return (NULL != intercept_function_);
+  }
+
+ private:
+  HMODULE module_handle_;
+  void* intercept_function_;
+  void* original_function_;
+  IMAGE_THUNK_DATA* iat_thunk_;
+
+  DISALLOW_COPY_AND_ASSIGN(IATPatchFunction);
+};
+
+}  // namespace win
+}  // namespace app
+
+#endif  // APP_WIN_IAT_PATCH_FUNCTION_H_