Home | History | Annotate | Download | only in base 
      1 // Copyright (c) 2006-2008 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_EQ(NO_ERROR, 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
    241