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.
5 #include "base/win/iat_patch_function.h"
7 #include "base/logging.h"
8 #include "base/win/pe_image.h"
15 struct InterceptFunctionInformation {
16 bool finished_operation;
17 const char* imported_from_module;
18 const char* function_name;
21 IMAGE_THUNK_DATA** iat_thunk;
25 void* GetIATFunction(IMAGE_THUNK_DATA* iat_thunk) {
26 if (NULL == iat_thunk) {
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.
36 IMAGE_THUNK_DATA thunk;
40 iat_function.thunk = *iat_thunk;
41 return iat_function.pointer;
43 // Change the page protection (of code pages) to writable and copy
44 // the data at the specified location
47 // old_code Target location to copy
49 // length Number of bytes to copy
51 // Returns: Windows error code (winerror.h). NO_ERROR if successful
52 DWORD ModifyCode(void* old_code, void* new_code, int length) {
53 if ((NULL == old_code) || (NULL == new_code) || (0 == length)) {
55 return ERROR_INVALID_PARAMETER;
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,
64 &old_page_protection)) {
67 CopyMemory(old_code, new_code, length);
69 // Restore the old page protection.
70 error = ERROR_SUCCESS;
71 VirtualProtect(old_code,
74 &old_page_protection);
76 error = GetLastError();
83 bool 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);
89 if (NULL == intercept_information) {
96 if ((0 == lstrcmpiA(module, intercept_information->imported_from_module)) &&
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);
104 if (NULL != intercept_information->iat_thunk) {
105 *(intercept_information->iat_thunk) = iat;
109 COMPILE_ASSERT(sizeof(iat->u1.Function) ==
110 sizeof(intercept_information->new_function), unknown_IAT_thunk_format);
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));
118 // Terminate further enumeration.
119 intercept_information->finished_operation = true;
126 // Helper to intercept a function in an import table of a specific
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.
138 // Returns: Returns NO_ERROR on success or Windows error code
139 // as defined in winerror.h
140 DWORD InterceptImportedFunction(HMODULE module_handle,
141 const char* imported_from_module,
142 const char* function_name, void* new_function,
144 IMAGE_THUNK_DATA** iat_thunk) {
145 if ((NULL == module_handle) || (NULL == imported_from_module) ||
146 (NULL == function_name) || (NULL == new_function)) {
148 return ERROR_INVALID_PARAMETER;
151 base::win::PEImage target_image(module_handle);
152 if (!target_image.VerifyMagic()) {
154 return ERROR_INVALID_PARAMETER;
157 InterceptFunctionInformation intercept_information = {
159 imported_from_module,
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);
174 return intercept_information.return_code;
177 // Restore intercepted IAT entry with the original function.
180 // intercept_function Interceptor function
181 // original_function Receives the original function pointer
183 // Returns: Returns NO_ERROR on success or Windows error code
184 // as defined in winerror.h
185 DWORD 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)) {
191 return ERROR_INVALID_PARAMETER;
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.
198 return ERROR_INVALID_FUNCTION;
201 return ModifyCode(&(iat_thunk->u1.Function),
203 sizeof(original_function));
208 IATPatchFunction::IATPatchFunction()
209 : module_handle_(NULL),
210 original_function_(NULL),
212 intercept_function_(NULL) {
215 IATPatchFunction::~IATPatchFunction() {
216 if (NULL != intercept_function_) {
217 DWORD error = Unpatch();
218 DCHECK_EQ(static_cast<DWORD>(NO_ERROR), error);
222 DWORD 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_);
230 HMODULE module_handle = LoadLibraryW(module);
232 if (module_handle == NULL) {
234 return GetLastError();
237 DWORD error = InterceptImportedFunction(module_handle,
238 imported_from_module,
244 if (NO_ERROR == error) {
245 DCHECK_NE(original_function_, intercept_function_);
246 module_handle_ = module_handle;
247 intercept_function_ = new_function;
249 FreeLibrary(module_handle);
255 DWORD IATPatchFunction::Unpatch() {
256 DWORD error = RestoreImportedFunction(intercept_function_,
259 DCHECK_EQ(static_cast<DWORD>(NO_ERROR), error);
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
268 FreeLibrary(module_handle_);
269 module_handle_ = NULL;
270 intercept_function_ = NULL;
271 original_function_ = NULL;