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