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#include "base/win/iat_patch_function.h"
6
7#include "base/logging.h"
8#include "base/win/pe_image.h"
9
10namespace base {
11namespace win {
12
13namespace {
14
15struct InterceptFunctionInformation {
16  bool finished_operation;
17  const char* imported_from_module;
18  const char* function_name;
19  void* new_function;
20  void** old_function;
21  IMAGE_THUNK_DATA** iat_thunk;
22  DWORD return_code;
23};
24
25void* GetIATFunction(IMAGE_THUNK_DATA* iat_thunk) {
26  if (NULL == iat_thunk) {
27    NOTREACHED();
28    return NULL;
29  }
30
31  // Works around the 64 bit portability warning:
32  // The Function member inside IMAGE_THUNK_DATA is really a pointer
33  // to the IAT function. IMAGE_THUNK_DATA correctly maps to IMAGE_THUNK_DATA32
34  // or IMAGE_THUNK_DATA64 for correct pointer size.
35  union FunctionThunk {
36    IMAGE_THUNK_DATA thunk;
37    void* pointer;
38  } iat_function;
39
40  iat_function.thunk = *iat_thunk;
41  return iat_function.pointer;
42}
43// Change the page protection (of code pages) to writable and copy
44// the data at the specified location
45//
46// Arguments:
47// old_code               Target location to copy
48// new_code               Source
49// length                 Number of bytes to copy
50//
51// Returns: Windows error code (winerror.h). NO_ERROR if successful
52DWORD ModifyCode(void* old_code, void* new_code, int length) {
53  if ((NULL == old_code) || (NULL == new_code) || (0 == length)) {
54    NOTREACHED();
55    return ERROR_INVALID_PARAMETER;
56  }
57
58  // Change the page protection so that we can write.
59  DWORD error = NO_ERROR;
60  DWORD old_page_protection = 0;
61  if (VirtualProtect(old_code,
62                     length,
63                     PAGE_READWRITE,
64                     &old_page_protection)) {
65
66    // Write the data.
67    CopyMemory(old_code, new_code, length);
68
69    // Restore the old page protection.
70    error = ERROR_SUCCESS;
71    VirtualProtect(old_code,
72                  length,
73                  old_page_protection,
74                  &old_page_protection);
75  } else {
76    error = GetLastError();
77    NOTREACHED();
78  }
79
80  return error;
81}
82
83bool InterceptEnumCallback(const base::win::PEImage& image, const char* module,
84                           DWORD ordinal, const char* name, DWORD hint,
85                           IMAGE_THUNK_DATA* iat, void* cookie) {
86  InterceptFunctionInformation* intercept_information =
87    reinterpret_cast<InterceptFunctionInformation*>(cookie);
88
89  if (NULL == intercept_information) {
90    NOTREACHED();
91    return false;
92  }
93
94  DCHECK(module);
95
96  if ((0 == lstrcmpiA(module, intercept_information->imported_from_module)) &&
97     (NULL != name) &&
98     (0 == lstrcmpiA(name, intercept_information->function_name))) {
99    // Save the old pointer.
100    if (NULL != intercept_information->old_function) {
101      *(intercept_information->old_function) = GetIATFunction(iat);
102    }
103
104    if (NULL != intercept_information->iat_thunk) {
105      *(intercept_information->iat_thunk) = iat;
106    }
107
108    // portability check
109    COMPILE_ASSERT(sizeof(iat->u1.Function) ==
110      sizeof(intercept_information->new_function), unknown_IAT_thunk_format);
111
112    // Patch the function.
113    intercept_information->return_code =
114      ModifyCode(&(iat->u1.Function),
115                 &(intercept_information->new_function),
116                 sizeof(intercept_information->new_function));
117
118    // Terminate further enumeration.
119    intercept_information->finished_operation = true;
120    return false;
121  }
122
123  return true;
124}
125
126// Helper to intercept a function in an import table of a specific
127// module.
128//
129// Arguments:
130// module_handle          Module to be intercepted
131// imported_from_module   Module that exports the symbol
132// function_name          Name of the API to be intercepted
133// new_function           Interceptor function
134// old_function           Receives the original function pointer
135// iat_thunk              Receives pointer to IAT_THUNK_DATA
136//                        for the API from the import table.
137//
138// Returns: Returns NO_ERROR on success or Windows error code
139//          as defined in winerror.h
140DWORD InterceptImportedFunction(HMODULE module_handle,
141                                const char* imported_from_module,
142                                const char* function_name, void* new_function,
143                                void** old_function,
144                                IMAGE_THUNK_DATA** iat_thunk) {
145  if ((NULL == module_handle) || (NULL == imported_from_module) ||
146     (NULL == function_name) || (NULL == new_function)) {
147    NOTREACHED();
148    return ERROR_INVALID_PARAMETER;
149  }
150
151  base::win::PEImage target_image(module_handle);
152  if (!target_image.VerifyMagic()) {
153    NOTREACHED();
154    return ERROR_INVALID_PARAMETER;
155  }
156
157  InterceptFunctionInformation intercept_information = {
158    false,
159    imported_from_module,
160    function_name,
161    new_function,
162    old_function,
163    iat_thunk,
164    ERROR_GEN_FAILURE};
165
166  // First go through the IAT. If we don't find the import we are looking
167  // for in IAT, search delay import table.
168  target_image.EnumAllImports(InterceptEnumCallback, &intercept_information);
169  if (!intercept_information.finished_operation) {
170    target_image.EnumAllDelayImports(InterceptEnumCallback,
171                                     &intercept_information);
172  }
173
174  return intercept_information.return_code;
175}
176
177// Restore intercepted IAT entry with the original function.
178//
179// Arguments:
180// intercept_function     Interceptor function
181// original_function      Receives the original function pointer
182//
183// Returns: Returns NO_ERROR on success or Windows error code
184//          as defined in winerror.h
185DWORD RestoreImportedFunction(void* intercept_function,
186                              void* original_function,
187                              IMAGE_THUNK_DATA* iat_thunk) {
188  if ((NULL == intercept_function) || (NULL == original_function) ||
189      (NULL == iat_thunk)) {
190    NOTREACHED();
191    return ERROR_INVALID_PARAMETER;
192  }
193
194  if (GetIATFunction(iat_thunk) != intercept_function) {
195    // Check if someone else has intercepted on top of us.
196    // We cannot unpatch in this case, just raise a red flag.
197    NOTREACHED();
198    return ERROR_INVALID_FUNCTION;
199  }
200
201  return ModifyCode(&(iat_thunk->u1.Function),
202                    &original_function,
203                    sizeof(original_function));
204}
205
206}  // namespace
207
208IATPatchFunction::IATPatchFunction()
209    : module_handle_(NULL),
210      original_function_(NULL),
211      iat_thunk_(NULL),
212      intercept_function_(NULL) {
213}
214
215IATPatchFunction::~IATPatchFunction() {
216  if (NULL != intercept_function_) {
217    DWORD error = Unpatch();
218    DCHECK_EQ(static_cast<DWORD>(NO_ERROR), error);
219  }
220}
221
222DWORD IATPatchFunction::Patch(const wchar_t* module,
223                              const char* imported_from_module,
224                              const char* function_name,
225                              void* new_function) {
226  DCHECK_EQ(static_cast<void*>(NULL), original_function_);
227  DCHECK_EQ(static_cast<IMAGE_THUNK_DATA*>(NULL), iat_thunk_);
228  DCHECK_EQ(static_cast<void*>(NULL), intercept_function_);
229
230  HMODULE module_handle = LoadLibraryW(module);
231
232  if (module_handle == NULL) {
233    NOTREACHED();
234    return GetLastError();
235  }
236
237  DWORD error = InterceptImportedFunction(module_handle,
238                                          imported_from_module,
239                                          function_name,
240                                          new_function,
241                                          &original_function_,
242                                          &iat_thunk_);
243
244  if (NO_ERROR == error) {
245    DCHECK_NE(original_function_, intercept_function_);
246    module_handle_ = module_handle;
247    intercept_function_ = new_function;
248  } else {
249    FreeLibrary(module_handle);
250  }
251
252  return error;
253}
254
255DWORD IATPatchFunction::Unpatch() {
256  DWORD error = RestoreImportedFunction(intercept_function_,
257                                        original_function_,
258                                        iat_thunk_);
259  DCHECK_EQ(static_cast<DWORD>(NO_ERROR), error);
260
261  // Hands off the intercept if we fail to unpatch.
262  // If IATPatchFunction::Unpatch fails during RestoreImportedFunction
263  // it means that we cannot safely unpatch the import address table
264  // patch. In this case its better to be hands off the intercept as
265  // trying to unpatch again in the destructor of IATPatchFunction is
266  // not going to be any safer
267  if (module_handle_)
268    FreeLibrary(module_handle_);
269  module_handle_ = NULL;
270  intercept_function_ = NULL;
271  original_function_ = NULL;
272  iat_thunk_ = NULL;
273
274  return error;
275}
276
277}  // namespace win
278}  // namespace base
279