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