license.bot | f003cfe | 2008-08-24 09:55:55 +0900 | [diff] [blame^] | 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. |
initial.commit | 3f4a732 | 2008-07-27 06:49:38 +0900 | [diff] [blame] | 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 | : original_function_(NULL), |
| 173 | iat_thunk_(NULL), |
| 174 | intercept_function_(NULL) { |
| 175 | } |
| 176 | |
| 177 | IATPatchFunction::~IATPatchFunction() { |
| 178 | if (NULL != intercept_function_) { |
| 179 | DWORD error = Unpatch(); |
| 180 | DCHECK_EQ(NO_ERROR, error); |
| 181 | } |
| 182 | } |
| 183 | |
| 184 | DWORD IATPatchFunction::Patch(HMODULE module_handle, |
| 185 | const char* imported_from_module, |
| 186 | const char* function_name, |
| 187 | void* new_function) { |
| 188 | DCHECK_EQ(static_cast<void*>(NULL), original_function_); |
| 189 | DCHECK_EQ(static_cast<IMAGE_THUNK_DATA*>(NULL), iat_thunk_); |
| 190 | DCHECK_EQ(static_cast<void*>(NULL), intercept_function_); |
| 191 | |
| 192 | DWORD error = InterceptImportedFunction(module_handle, |
| 193 | imported_from_module, |
| 194 | function_name, |
| 195 | new_function, |
| 196 | &original_function_, |
| 197 | &iat_thunk_); |
| 198 | |
| 199 | if (NO_ERROR == error) { |
| 200 | DCHECK_NE(original_function_, intercept_function_); |
| 201 | intercept_function_ = new_function; |
| 202 | } |
| 203 | |
| 204 | return error; |
| 205 | } |
| 206 | |
| 207 | DWORD IATPatchFunction::Unpatch() { |
| 208 | DWORD error = RestoreImportedFunction(intercept_function_, |
| 209 | original_function_, |
| 210 | iat_thunk_); |
| 211 | |
| 212 | if (NO_ERROR == error) { |
| 213 | intercept_function_ = NULL; |
| 214 | original_function_ = NULL; |
| 215 | iat_thunk_ = NULL; |
| 216 | } |
| 217 | |
| 218 | return error; |
| 219 | } |
| 220 | |
| 221 | } // namespace iat_patch |
license.bot | f003cfe | 2008-08-24 09:55:55 +0900 | [diff] [blame^] | 222 | |