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
13namespace base {
14namespace 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.
23class 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
63 private:
64  HMODULE module_handle_;
65  void* intercept_function_;
66  void* original_function_;
67  IMAGE_THUNK_DATA* iat_thunk_;
68
69  DISALLOW_COPY_AND_ASSIGN(IATPatchFunction);
70};
71
72BASE_EXPORT DWORD ModifyCode(void* old_code, void* new_code, int length);
73
74}  // namespace win
75}  // namespace base
76
77#endif  // BASE_WIN_IAT_PATCH_FUNCTION_H_
78