base/win/iat_patch_function.cc

   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
  10 namespace base {
  11 namespace win {
  12
  13 namespace {
  14
  15 struct 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
  25 void* 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
  44 bool InterceptEnumCallback(const base::win::PEImage& image, const char* module,
  45                            DWORD ordinal, const char* name, DWORD hint,
  46                            IMAGE_THUNK_DATA* iat, void* cookie) {
  47   InterceptFunctionInformation* intercept_information =
  48     reinterpret_cast<InterceptFunctionInformation*>(cookie);
  49
  50   if (NULL == intercept_information) {
  51     NOTREACHED();
  52     return false;
  53   }
  54
  55   DCHECK(module);
  56
  57   if ((0 == lstrcmpiA(module, intercept_information->imported_from_module)) &&
  58      (NULL != name) &&
  59      (0 == lstrcmpiA(name, intercept_information->function_name))) {
  60     // Save the old pointer.
  61     if (NULL != intercept_information->old_function) {
  62       *(intercept_information->old_function) = GetIATFunction(iat);
  63     }
  64
  65     if (NULL != intercept_information->iat_thunk) {
  66       *(intercept_information->iat_thunk) = iat;
  67     }
  68
  69     // portability check
  70     COMPILE_ASSERT(sizeof(iat->u1.Function) ==
  71       sizeof(intercept_information->new_function), unknown_IAT_thunk_format);
  72
  73     // Patch the function.
  74     intercept_information->return_code =
  75       ModifyCode(&(iat->u1.Function),
  76                  &(intercept_information->new_function),
  77                  sizeof(intercept_information->new_function));
  78
  79     // Terminate further enumeration.
  80     intercept_information->finished_operation = true;
  81     return false;
  82   }
  83
  84   return true;
  85 }
  86
  87 // Helper to intercept a function in an import table of a specific
  88 // module.
  89 //
  90 // Arguments:
  91 // module_handle          Module to be intercepted
  92 // imported_from_module   Module that exports the symbol
  93 // function_name          Name of the API to be intercepted
  94 // new_function           Interceptor function
  95 // old_function           Receives the original function pointer
  96 // iat_thunk              Receives pointer to IAT_THUNK_DATA
  97 //                        for the API from the import table.
  98 //
  99 // Returns: Returns NO_ERROR on success or Windows error code
 100 //          as defined in winerror.h
 101 DWORD InterceptImportedFunction(HMODULE module_handle,
 102                                 const char* imported_from_module,
 103                                 const char* function_name, void* new_function,
 104                                 void** old_function,
 105                                 IMAGE_THUNK_DATA** iat_thunk) {
 106   if ((NULL == module_handle) || (NULL == imported_from_module) ||
 107      (NULL == function_name) || (NULL == new_function)) {
 108     NOTREACHED();
 109     return ERROR_INVALID_PARAMETER;
 110   }
 111
 112   base::win::PEImage target_image(module_handle);
 113   if (!target_image.VerifyMagic()) {
 114     NOTREACHED();
 115     return ERROR_INVALID_PARAMETER;
 116   }
 117
 118   InterceptFunctionInformation intercept_information = {
 119     false,
 120     imported_from_module,
 121     function_name,
 122     new_function,
 123     old_function,
 124     iat_thunk,
 125     ERROR_GEN_FAILURE};
 126
 127   // First go through the IAT. If we don't find the import we are looking
 128   // for in IAT, search delay import table.
 129   target_image.EnumAllImports(InterceptEnumCallback, &intercept_information);
 130   if (!intercept_information.finished_operation) {
 131     target_image.EnumAllDelayImports(InterceptEnumCallback,
 132                                      &intercept_information);
 133   }
 134
 135   return intercept_information.return_code;
 136 }
 137
 138 // Restore intercepted IAT entry with the original function.
 139 //
 140 // Arguments:
 141 // intercept_function     Interceptor function
 142 // original_function      Receives the original function pointer
 143 //
 144 // Returns: Returns NO_ERROR on success or Windows error code
 145 //          as defined in winerror.h
 146 DWORD RestoreImportedFunction(void* intercept_function,
 147                               void* original_function,
 148                               IMAGE_THUNK_DATA* iat_thunk) {
 149   if ((NULL == intercept_function) || (NULL == original_function) ||
 150       (NULL == iat_thunk)) {
 151     NOTREACHED();
 152     return ERROR_INVALID_PARAMETER;
 153   }
 154
 155   if (GetIATFunction(iat_thunk) != intercept_function) {
 156     // Check if someone else has intercepted on top of us.
 157     // We cannot unpatch in this case, just raise a red flag.
 158     NOTREACHED();
 159     return ERROR_INVALID_FUNCTION;
 160   }
 161
 162   return ModifyCode(&(iat_thunk->u1.Function),
 163                     &original_function,
 164                     sizeof(original_function));
 165 }
 166
 167 }  // namespace
 168
 169 // Change the page protection (of code pages) to writable and copy
 170 // the data at the specified location
 171 //
 172 // Arguments:
 173 // old_code               Target location to copy
 174 // new_code               Source
 175 // length                 Number of bytes to copy
 176 //
 177 // Returns: Windows error code (winerror.h). NO_ERROR if successful
 178 DWORD ModifyCode(void* old_code, void* new_code, int length) {
 179   if ((NULL == old_code) || (NULL == new_code) || (0 == length)) {
 180     NOTREACHED();
 181     return ERROR_INVALID_PARAMETER;
 182   }
 183
 184   // Change the page protection so that we can write.
 185   MEMORY_BASIC_INFORMATION memory_info;
 186   DWORD error = NO_ERROR;
 187   DWORD old_page_protection = 0;
 188
 189   if (!VirtualQuery(old_code, &memory_info, sizeof(memory_info))) {
 190     error = GetLastError();
 191     return error;
 192   }
 193
 194   DWORD is_executable = (PAGE_EXECUTE | PAGE_EXECUTE_READ |
 195                         PAGE_EXECUTE_READWRITE | PAGE_EXECUTE_WRITECOPY) &
 196                         memory_info.Protect;
 197
 198   if (VirtualProtect(old_code,
 199                      length,
 200                      is_executable ? PAGE_EXECUTE_READWRITE :
 201                                      PAGE_READWRITE,
 202                      &old_page_protection)) {
 203
 204     // Write the data.
 205     CopyMemory(old_code, new_code, length);
 206
 207     // Restore the old page protection.
 208     error = ERROR_SUCCESS;
 209     VirtualProtect(old_code,
 210                   length,
 211                   old_page_protection,
 212                   &old_page_protection);
 213   } else {
 214     error = GetLastError();
 215   }
 216
 217   return error;
 218 }
 219
 220 IATPatchFunction::IATPatchFunction()
 221     : module_handle_(NULL),
 222       original_function_(NULL),
 223       iat_thunk_(NULL),
 224       intercept_function_(NULL) {
 225 }
 226
 227 IATPatchFunction::~IATPatchFunction() {
 228   if (NULL != intercept_function_) {
 229     DWORD error = Unpatch();
 230     DCHECK_EQ(static_cast<DWORD>(NO_ERROR), error);
 231   }
 232 }
 233
 234 DWORD IATPatchFunction::Patch(const wchar_t* module,
 235                               const char* imported_from_module,
 236                               const char* function_name,
 237                               void* new_function) {
 238   HMODULE module_handle = LoadLibraryW(module);
 239   if (module_handle == NULL) {
 240     NOTREACHED();
 241     return GetLastError();
 242   }
 243
 244   DWORD error = PatchFromModule(module_handle, imported_from_module,
 245                                 function_name, new_function);
 246   if (NO_ERROR == error) {
 247     module_handle_ = module_handle;
 248   } else {
 249     FreeLibrary(module_handle);
 250   }
 251
 252   return error;
 253 }
 254
 255 DWORD IATPatchFunction::PatchFromModule(HMODULE module,
 256                                         const char* imported_from_module,
 257                                         const char* function_name,
 258                                         void* new_function) {
 259   DCHECK_EQ(static_cast<void*>(NULL), original_function_);
 260   DCHECK_EQ(static_cast<IMAGE_THUNK_DATA*>(NULL), iat_thunk_);
 261   DCHECK_EQ(static_cast<void*>(NULL), intercept_function_);
 262   DCHECK(module);
 263
 264   DWORD error = InterceptImportedFunction(module,
 265                                           imported_from_module,
 266                                           function_name,
 267                                           new_function,
 268                                           &original_function_,
 269                                           &iat_thunk_);
 270
 271   if (NO_ERROR == error) {
 272     DCHECK_NE(original_function_, intercept_function_);
 273     intercept_function_ = new_function;
 274   }
 275
 276   return error;
 277 }
 278
 279 DWORD IATPatchFunction::Unpatch() {
 280   DWORD error = RestoreImportedFunction(intercept_function_,
 281                                         original_function_,
 282                                         iat_thunk_);
 283   DCHECK_EQ(static_cast<DWORD>(NO_ERROR), error);
 284
 285   // Hands off the intercept if we fail to unpatch.
 286   // If IATPatchFunction::Unpatch fails during RestoreImportedFunction
 287   // it means that we cannot safely unpatch the import address table
 288   // patch. In this case its better to be hands off the intercept as
 289   // trying to unpatch again in the destructor of IATPatchFunction is
 290   // not going to be any safer
 291   if (module_handle_)
 292     FreeLibrary(module_handle_);
 293   module_handle_ = NULL;
 294   intercept_function_ = NULL;
 295   original_function_ = NULL;
 296   iat_thunk_ = NULL;
 297
 298   return error;
 299 }
 300
 301 void* IATPatchFunction::original_function() const {
 302   DCHECK(is_patched());
 303   return original_function_;
 304 }
 305
 306 }  // namespace win
 307 }  // namespace base