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