base/win/iat_patch_function.h

   1 // Copyright (c) 2011 The Chromium Authors. All rights reserved.
   2 // Use of this source code is governed by a BSD-style license that can be
   3 // found in the LICENSE file.
   4
   5 #ifndef BASE_WIN_IAT_PATCH_FUNCTION_H_
   6 #define BASE_WIN_IAT_PATCH_FUNCTION_H_
   7
   8 #include <windows.h>
   9
  10 #include "base/base_export.h"
  11 #include "base/basictypes.h"
  12
  13 namespace base {
  14 namespace win {
  15
  16 // A class that encapsulates Import Address Table patching helpers and restores
  17 // the original function in the destructor.
  18 //
  19 // It will intercept functions for a specific DLL imported from another DLL.
  20 // This is the case when, for example, we want to intercept
  21 // CertDuplicateCertificateContext function (exported from crypt32.dll) called
  22 // by wininet.dll.
  23 class BASE_EXPORT IATPatchFunction {
  24  public:
  25   IATPatchFunction();
  26   ~IATPatchFunction();
  27
  28   // Intercept a function in an import table of a specific
  29   // module. Save the original function and the import
  30   // table address. These values will be used later
  31   // during Unpatch
  32   //
  33   // Arguments:
  34   // module                 Module to be intercepted
  35   // imported_from_module   Module that exports the 'function_name'
  36   // function_name          Name of the API to be intercepted
  37   //
  38   // Returns: Windows error code (winerror.h). NO_ERROR if successful
  39   //
  40   // Note: Patching a function will make the IAT patch take some "ownership" on
  41   // |module|.  It will LoadLibrary(module) to keep the DLL alive until a call
  42   // to Unpatch(), which will call FreeLibrary() and allow the module to be
  43   // unloaded.  The idea is to help prevent the DLL from going away while a
  44   // patch is still active.
  45   DWORD Patch(const wchar_t* module,
  46               const char* imported_from_module,
  47               const char* function_name,
  48               void* new_function);
  49
  50   // Same as Patch(), but uses a handle to a |module| instead of the DLL name.
  51   DWORD PatchFromModule(HMODULE module,
  52                         const char* imported_from_module,
  53                         const char* function_name,
  54                         void* new_function);
  55
  56   // Unpatch the IAT entry using internally saved original
  57   // function.
  58   //
  59   // Returns: Windows error code (winerror.h). NO_ERROR if successful
  60   DWORD Unpatch();
  61
  62   bool is_patched() const {
  63     return (NULL != intercept_function_);
  64   }
  65
  66   void* original_function() const;
  67
  68
  69  private:
  70   HMODULE module_handle_;
  71   void* intercept_function_;
  72   void* original_function_;
  73   IMAGE_THUNK_DATA* iat_thunk_;
  74
  75   DISALLOW_COPY_AND_ASSIGN(IATPatchFunction);
  76 };
  77
  78 BASE_EXPORT DWORD ModifyCode(void* old_code, void* new_code, int length);
  79
  80 }  // namespace win
  81 }  // namespace base
  82
  83 #endif  // BASE_WIN_IAT_PATCH_FUNCTION_H_