initial.commit | 3f4a732 | 2008-07-27 06:49:38 +0900 | [diff] [blame^] | 1 | // Copyright 2008, Google Inc. |
| 2 | // All rights reserved. |
| 3 | // |
| 4 | // Redistribution and use in source and binary forms, with or without |
| 5 | // modification, are permitted provided that the following conditions are |
| 6 | // met: |
| 7 | // |
| 8 | // * Redistributions of source code must retain the above copyright |
| 9 | // notice, this list of conditions and the following disclaimer. |
| 10 | // * Redistributions in binary form must reproduce the above |
| 11 | // copyright notice, this list of conditions and the following disclaimer |
| 12 | // in the documentation and/or other materials provided with the |
| 13 | // distribution. |
| 14 | // * Neither the name of Google Inc. nor the names of its |
| 15 | // contributors may be used to endorse or promote products derived from |
| 16 | // this software without specific prior written permission. |
| 17 | // |
| 18 | // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS |
| 19 | // "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT |
| 20 | // LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR |
| 21 | // A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT |
| 22 | // OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, |
| 23 | // SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT |
| 24 | // LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, |
| 25 | // DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY |
| 26 | // THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT |
| 27 | // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE |
| 28 | // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. |
| 29 | |
| 30 | #include "base/iat_patch.h" |
| 31 | #include "base/logging.h" |
| 32 | |
| 33 | namespace iat_patch { |
| 34 | |
| 35 | struct InterceptFunctionInformation { |
| 36 | bool finished_operation; |
| 37 | const char* imported_from_module; |
| 38 | const char* function_name; |
| 39 | void* new_function; |
| 40 | void** old_function; |
| 41 | IMAGE_THUNK_DATA** iat_thunk; |
| 42 | DWORD return_code; |
| 43 | }; |
| 44 | |
| 45 | static void* GetIATFunction(IMAGE_THUNK_DATA* iat_thunk) { |
| 46 | if (NULL == iat_thunk) { |
| 47 | NOTREACHED(); |
| 48 | return NULL; |
| 49 | } |
| 50 | |
| 51 | // Works around the 64 bit portability warning: |
| 52 | // The Function member inside IMAGE_THUNK_DATA is really a pointer |
| 53 | // to the IAT function. IMAGE_THUNK_DATA correctly maps to IMAGE_THUNK_DATA32 |
| 54 | // or IMAGE_THUNK_DATA64 for correct pointer size. |
| 55 | union FunctionThunk { |
| 56 | IMAGE_THUNK_DATA thunk; |
| 57 | void* pointer; |
| 58 | } iat_function; |
| 59 | |
| 60 | iat_function.thunk = *iat_thunk; |
| 61 | return iat_function.pointer; |
| 62 | } |
| 63 | |
| 64 | static bool InterceptEnumCallback(const PEImage &image, const char* module, |
| 65 | DWORD ordinal, const char* name, DWORD hint, |
| 66 | IMAGE_THUNK_DATA* iat, void* cookie) { |
| 67 | InterceptFunctionInformation* intercept_information = |
| 68 | reinterpret_cast<InterceptFunctionInformation*>(cookie); |
| 69 | |
| 70 | if (NULL == intercept_information) { |
| 71 | NOTREACHED(); |
| 72 | return false; |
| 73 | } |
| 74 | |
| 75 | DCHECK(module); |
| 76 | |
| 77 | if ((0 == lstrcmpiA(module, intercept_information->imported_from_module)) && |
| 78 | (NULL != name) && |
| 79 | (0 == lstrcmpiA(name, intercept_information->function_name))) { |
| 80 | // Save the old pointer. |
| 81 | if (NULL != intercept_information->old_function) { |
| 82 | *(intercept_information->old_function) = GetIATFunction(iat); |
| 83 | } |
| 84 | |
| 85 | if (NULL != intercept_information->iat_thunk) { |
| 86 | *(intercept_information->iat_thunk) = iat; |
| 87 | } |
| 88 | |
| 89 | // portability check |
| 90 | COMPILE_ASSERT(sizeof(iat->u1.Function) == |
| 91 | sizeof(intercept_information->new_function), unknown_IAT_thunk_format); |
| 92 | |
| 93 | // Patch the function. |
| 94 | intercept_information->return_code = |
| 95 | ModifyCode(&(iat->u1.Function), |
| 96 | &(intercept_information->new_function), |
| 97 | sizeof(intercept_information->new_function)); |
| 98 | |
| 99 | // Terminate further enumeration. |
| 100 | intercept_information->finished_operation = true; |
| 101 | return false; |
| 102 | } |
| 103 | |
| 104 | return true; |
| 105 | } |
| 106 | |
| 107 | DWORD InterceptImportedFunction(HMODULE module_handle, |
| 108 | const char* imported_from_module, |
| 109 | const char* function_name, void* new_function, |
| 110 | void** old_function, |
| 111 | IMAGE_THUNK_DATA** iat_thunk) { |
| 112 | if ((NULL == module_handle) || (NULL == imported_from_module) || |
| 113 | (NULL == function_name) || (NULL == new_function)) { |
| 114 | NOTREACHED(); |
| 115 | return ERROR_INVALID_PARAMETER; |
| 116 | } |
| 117 | |
| 118 | PEImage target_image(module_handle); |
| 119 | if (!target_image.VerifyMagic()) { |
| 120 | NOTREACHED(); |
| 121 | return ERROR_INVALID_PARAMETER; |
| 122 | } |
| 123 | |
| 124 | InterceptFunctionInformation intercept_information = { |
| 125 | false, |
| 126 | imported_from_module, |
| 127 | function_name, |
| 128 | new_function, |
| 129 | old_function, |
| 130 | iat_thunk, |
| 131 | ERROR_GEN_FAILURE}; |
| 132 | |
| 133 | // First go through the IAT. If we don't find the import we are looking |
| 134 | // for in IAT, search delay import table. |
| 135 | target_image.EnumAllImports(InterceptEnumCallback, &intercept_information); |
| 136 | if (!intercept_information.finished_operation) { |
| 137 | target_image.EnumAllDelayImports(InterceptEnumCallback, |
| 138 | &intercept_information); |
| 139 | } |
| 140 | |
| 141 | return intercept_information.return_code; |
| 142 | } |
| 143 | |
| 144 | DWORD RestoreImportedFunction(void* intercept_function, |
| 145 | void* original_function, |
| 146 | IMAGE_THUNK_DATA* iat_thunk) { |
| 147 | if ((NULL == intercept_function) || (NULL == original_function) || |
| 148 | (NULL == iat_thunk)) { |
| 149 | NOTREACHED(); |
| 150 | return ERROR_INVALID_PARAMETER; |
| 151 | } |
| 152 | |
| 153 | if (GetIATFunction(iat_thunk) != intercept_function) { |
| 154 | // Check if someone else has intercepted on top of us. |
| 155 | // We cannot unpatch in this case, just raise a red flag. |
| 156 | NOTREACHED(); |
| 157 | return ERROR_INVALID_FUNCTION; |
| 158 | } |
| 159 | |
| 160 | return ModifyCode(&(iat_thunk->u1.Function), |
| 161 | &original_function, |
| 162 | sizeof(original_function)); |
| 163 | } |
| 164 | |
| 165 | DWORD ModifyCode(void* old_code, void* new_code, int length) { |
| 166 | if ((NULL == old_code) || (NULL == new_code) || (0 == length)) { |
| 167 | NOTREACHED(); |
| 168 | return ERROR_INVALID_PARAMETER; |
| 169 | } |
| 170 | |
| 171 | // Change the page protection so that we can write. |
| 172 | DWORD error = NO_ERROR; |
| 173 | DWORD old_page_protection = 0; |
| 174 | if (VirtualProtect(old_code, |
| 175 | length, |
| 176 | PAGE_READWRITE, |
| 177 | &old_page_protection)) { |
| 178 | |
| 179 | // Write the data. |
| 180 | CopyMemory(old_code, new_code, length); |
| 181 | |
| 182 | // Restore the old page protection. |
| 183 | error = ERROR_SUCCESS; |
| 184 | VirtualProtect(old_code, |
| 185 | length, |
| 186 | old_page_protection, |
| 187 | &old_page_protection); |
| 188 | } else { |
| 189 | error = GetLastError(); |
| 190 | NOTREACHED(); |
| 191 | } |
| 192 | |
| 193 | return error; |
| 194 | } |
| 195 | |
| 196 | IATPatchFunction::IATPatchFunction() |
| 197 | : original_function_(NULL), |
| 198 | iat_thunk_(NULL), |
| 199 | intercept_function_(NULL) { |
| 200 | } |
| 201 | |
| 202 | IATPatchFunction::~IATPatchFunction() { |
| 203 | if (NULL != intercept_function_) { |
| 204 | DWORD error = Unpatch(); |
| 205 | DCHECK_EQ(NO_ERROR, error); |
| 206 | } |
| 207 | } |
| 208 | |
| 209 | DWORD IATPatchFunction::Patch(HMODULE module_handle, |
| 210 | const char* imported_from_module, |
| 211 | const char* function_name, |
| 212 | void* new_function) { |
| 213 | DCHECK_EQ(static_cast<void*>(NULL), original_function_); |
| 214 | DCHECK_EQ(static_cast<IMAGE_THUNK_DATA*>(NULL), iat_thunk_); |
| 215 | DCHECK_EQ(static_cast<void*>(NULL), intercept_function_); |
| 216 | |
| 217 | DWORD error = InterceptImportedFunction(module_handle, |
| 218 | imported_from_module, |
| 219 | function_name, |
| 220 | new_function, |
| 221 | &original_function_, |
| 222 | &iat_thunk_); |
| 223 | |
| 224 | if (NO_ERROR == error) { |
| 225 | DCHECK_NE(original_function_, intercept_function_); |
| 226 | intercept_function_ = new_function; |
| 227 | } |
| 228 | |
| 229 | return error; |
| 230 | } |
| 231 | |
| 232 | DWORD IATPatchFunction::Unpatch() { |
| 233 | DWORD error = RestoreImportedFunction(intercept_function_, |
| 234 | original_function_, |
| 235 | iat_thunk_); |
| 236 | |
| 237 | if (NO_ERROR == error) { |
| 238 | intercept_function_ = NULL; |
| 239 | original_function_ = NULL; |
| 240 | iat_thunk_ = NULL; |
| 241 | } |
| 242 | |
| 243 | return error; |
| 244 | } |
| 245 | |
| 246 | } // namespace iat_patch |