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