base/iat_patch.cc

   1 // Copyright (c) 2010 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/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     : module_handle_(NULL),
 173       original_function_(NULL),
 174       iat_thunk_(NULL),
 175       intercept_function_(NULL) {
 176 }
 177
 178 IATPatchFunction::~IATPatchFunction() {
 179   if (NULL != intercept_function_) {
 180     DWORD error = Unpatch();
 181     DCHECK(error == NO_ERROR);
 182   }
 183 }
 184
 185 DWORD IATPatchFunction::Patch(const wchar_t* module,
 186                               const char* imported_from_module,
 187                               const char* function_name,
 188                               void* new_function) {
 189   DCHECK_EQ(static_cast<void*>(NULL), original_function_);
 190   DCHECK_EQ(static_cast<IMAGE_THUNK_DATA*>(NULL), iat_thunk_);
 191   DCHECK_EQ(static_cast<void*>(NULL), intercept_function_);
 192
 193   HMODULE module_handle = LoadLibraryW(module);
 194
 195   if (module_handle == NULL) {
 196     NOTREACHED();
 197     return GetLastError();
 198   }
 199
 200   DWORD error = InterceptImportedFunction(module_handle,
 201                                           imported_from_module,
 202                                           function_name,
 203                                           new_function,
 204                                           &original_function_,
 205                                           &iat_thunk_);
 206
 207   if (NO_ERROR == error) {
 208     DCHECK_NE(original_function_, intercept_function_);
 209     module_handle_ = module_handle;
 210     intercept_function_ = new_function;
 211   } else {
 212     FreeLibrary(module_handle);
 213   }
 214
 215   return error;
 216 }
 217
 218 DWORD IATPatchFunction::Unpatch() {
 219   DWORD error = RestoreImportedFunction(intercept_function_,
 220                                         original_function_,
 221                                         iat_thunk_);
 222   DCHECK(NO_ERROR == error);
 223
 224   // Hands off the intercept if we fail to unpatch.
 225   // If IATPatchFunction::Unpatch fails during RestoreImportedFunction
 226   // it means that we cannot safely unpatch the import address table
 227   // patch. In this case its better to be hands off the intercept as
 228   // trying to unpatch again in the destructor of IATPatchFunction is
 229   // not going to be any safer
 230   if (module_handle_)
 231     FreeLibrary(module_handle_);
 232   module_handle_ = NULL;
 233   intercept_function_ = NULL;
 234   original_function_ = NULL;
 235   iat_thunk_ = NULL;
 236
 237   return error;
 238 }
 239
 240 }  // namespace iat_patch