base/iat_patch.h

   1 // Copyright (c) 2010 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 // This file declares a helpers to intercept functions from a DLL.
   6 //
   7 // This set of functions are designed to intercept functions for a
   8 // specific DLL imported from another DLL. This is the case when,
   9 // for example, we want to intercept CertDuplicateCertificateContext
  10 // function (exported from crypt32.dll) called by wininet.dll.
  11
  12 #ifndef BASE_IAT_PATCH_H__
  13 #define BASE_IAT_PATCH_H__
  14
  15 #include <windows.h>
  16 #include "base/basictypes.h"
  17 #include "base/pe_image.h"
  18
  19 namespace iat_patch {
  20
  21 // Helper to intercept a function in an import table of a specific
  22 // module.
  23 //
  24 // Arguments:
  25 // module_handle          Module to be intercepted
  26 // imported_from_module   Module that exports the symbol
  27 // function_name          Name of the API to be intercepted
  28 // new_function           Interceptor function
  29 // old_function           Receives the original function pointer
  30 // iat_thunk              Receives pointer to IAT_THUNK_DATA
  31 //                        for the API from the import table.
  32 //
  33 // Returns: Returns NO_ERROR on success or Windows error code
  34 //          as defined in winerror.h
  35 //
  36 DWORD InterceptImportedFunction(HMODULE module_handle,
  37                                 const char* imported_from_module,
  38                                 const char* function_name,
  39                                 void* new_function,
  40                                 void** old_function,
  41                                 IMAGE_THUNK_DATA** iat_thunk);
  42
  43 // Restore intercepted IAT entry with the original function.
  44 //
  45 // Arguments:
  46 // intercept_function     Interceptor function
  47 // original_function      Receives the original function pointer
  48 //
  49 // Returns: Returns NO_ERROR on success or Windows error code
  50 //          as defined in winerror.h
  51 //
  52 DWORD RestoreImportedFunction(void* intercept_function,
  53                               void* original_function,
  54                               IMAGE_THUNK_DATA* iat_thunk);
  55
  56 // Change the page protection (of code pages) to writable and copy
  57 // the data at the specified location
  58 //
  59 // Arguments:
  60 // old_code               Target location to copy
  61 // new_code               Source
  62 // length                 Number of bytes to copy
  63 //
  64 // Returns: Windows error code (winerror.h). NO_ERROR if successful
  65 //
  66 DWORD ModifyCode(void* old_code,
  67                  void* new_code,
  68                  int length);
  69
  70 // A class that encapsulates IAT patching helpers and restores
  71 // the original function in the destructor.
  72 class IATPatchFunction {
  73  public:
  74   IATPatchFunction();
  75   ~IATPatchFunction();
  76
  77   // Intercept a function in an import table of a specific
  78   // module. Save the original function and the import
  79   // table address. These values will be used later
  80   // during Unpatch
  81   //
  82   // Arguments:
  83   // module                 Module to be intercepted
  84   // imported_from_module   Module that exports the 'function_name'
  85   // function_name          Name of the API to be intercepted
  86   //
  87   // Returns: Windows error code (winerror.h). NO_ERROR if successful
  88   //
  89   // Note: Patching a function will make the IAT patch take some "ownership" on
  90   // |module|.  It will LoadLibrary(module) to keep the DLL alive until a call
  91   // to Unpatch(), which will call FreeLibrary() and allow the module to be
  92   // unloaded.  The idea is to help prevent the DLL from going away while a
  93   // patch is still active.
  94   //
  95   DWORD Patch(const wchar_t* module,
  96               const char* imported_from_module,
  97               const char* function_name,
  98               void* new_function);
  99
 100   // Unpatch the IAT entry using internally saved original
 101   // function.
 102   //
 103   // Returns: Windows error code (winerror.h). NO_ERROR if successful
 104   //
 105   DWORD Unpatch();
 106
 107   bool is_patched() const {
 108     return (NULL != intercept_function_);
 109   }
 110
 111  private:
 112   HMODULE module_handle_;
 113   void* intercept_function_;
 114   void* original_function_;
 115   IMAGE_THUNK_DATA* iat_thunk_;
 116
 117   DISALLOW_COPY_AND_ASSIGN(IATPatchFunction);
 118 };
 119
 120 }  // namespace iat_patch
 121
 122 #endif  // BASE_IAT_PATCH_H__