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   // Unpatch the IAT entry using internally saved original
  51   // function.
  52   //
  53   // Returns: Windows error code (winerror.h). NO_ERROR if successful
  54   DWORD Unpatch();
  55
  56   bool is_patched() const {
  57     return (NULL != intercept_function_);
  58   }
  59
  60   void* original_function() const;
  61
  62  private:
  63   HMODULE module_handle_;
  64   void* intercept_function_;
  65   void* original_function_;
  66   IMAGE_THUNK_DATA* iat_thunk_;
  67
  68   DISALLOW_COPY_AND_ASSIGN(IATPatchFunction);
  69 };
  70
  71 }  // namespace win
  72 }  // namespace base
  73
  74 #endif  // BASE_WIN_IAT_PATCH_FUNCTION_H_