iat_patch_function.cc 8.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250
  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. #include "base/win/iat_patch_function.h"
  5. #include "base/check_op.h"
  6. #include "base/notreached.h"
  7. #include "base/win/patch_util.h"
  8. #include "base/win/pe_image.h"
  9. namespace base {
  10. namespace win {
  11. namespace {
  12. struct InterceptFunctionInformation {
  13. bool finished_operation;
  14. const char* imported_from_module;
  15. const char* function_name;
  16. void* new_function;
  17. void** old_function;
  18. IMAGE_THUNK_DATA** iat_thunk;
  19. DWORD return_code;
  20. };
  21. void* GetIATFunction(IMAGE_THUNK_DATA* iat_thunk) {
  22. if (!iat_thunk) {
  23. NOTREACHED();
  24. return nullptr;
  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. iat_function.thunk = *iat_thunk;
  35. return iat_function.pointer;
  36. }
  37. bool InterceptEnumCallback(const base::win::PEImage& image,
  38. const char* module,
  39. DWORD ordinal,
  40. const char* name,
  41. DWORD hint,
  42. IMAGE_THUNK_DATA* iat,
  43. void* cookie) {
  44. InterceptFunctionInformation* intercept_information =
  45. reinterpret_cast<InterceptFunctionInformation*>(cookie);
  46. if (!intercept_information) {
  47. NOTREACHED();
  48. return false;
  49. }
  50. DCHECK(module);
  51. if (name && (0 == lstrcmpiA(name, intercept_information->function_name))) {
  52. // Save the old pointer.
  53. if (intercept_information->old_function) {
  54. *(intercept_information->old_function) = GetIATFunction(iat);
  55. }
  56. if (intercept_information->iat_thunk) {
  57. *(intercept_information->iat_thunk) = iat;
  58. }
  59. // portability check
  60. static_assert(
  61. sizeof(iat->u1.Function) == sizeof(intercept_information->new_function),
  62. "unknown IAT thunk format");
  63. // Patch the function.
  64. intercept_information->return_code = internal::ModifyCode(
  65. &(iat->u1.Function), &(intercept_information->new_function),
  66. sizeof(intercept_information->new_function));
  67. // Terminate further enumeration.
  68. intercept_information->finished_operation = true;
  69. return false;
  70. }
  71. return true;
  72. }
  73. // Helper to intercept a function in an import table of a specific
  74. // module.
  75. //
  76. // Arguments:
  77. // module_handle Module to be intercepted
  78. // imported_from_module Module that exports the symbol
  79. // function_name Name of the API to be intercepted
  80. // new_function Interceptor function
  81. // old_function Receives the original function pointer
  82. // iat_thunk Receives pointer to IAT_THUNK_DATA
  83. // for the API from the import table.
  84. //
  85. // Returns: Returns NO_ERROR on success or Windows error code
  86. // as defined in winerror.h
  87. DWORD InterceptImportedFunction(HMODULE module_handle,
  88. const char* imported_from_module,
  89. const char* function_name,
  90. void* new_function,
  91. void** old_function,
  92. IMAGE_THUNK_DATA** iat_thunk) {
  93. if (!module_handle || !imported_from_module || !function_name ||
  94. !new_function) {
  95. NOTREACHED();
  96. return ERROR_INVALID_PARAMETER;
  97. }
  98. base::win::PEImage target_image(module_handle);
  99. if (!target_image.VerifyMagic()) {
  100. NOTREACHED();
  101. return ERROR_INVALID_PARAMETER;
  102. }
  103. InterceptFunctionInformation intercept_information = {false,
  104. imported_from_module,
  105. function_name,
  106. new_function,
  107. old_function,
  108. iat_thunk,
  109. ERROR_GEN_FAILURE};
  110. // First go through the IAT. If we don't find the import we are looking
  111. // for in IAT, search delay import table.
  112. target_image.EnumAllImports(InterceptEnumCallback, &intercept_information,
  113. imported_from_module);
  114. if (!intercept_information.finished_operation) {
  115. target_image.EnumAllDelayImports(
  116. InterceptEnumCallback, &intercept_information, imported_from_module);
  117. }
  118. return intercept_information.return_code;
  119. }
  120. // Restore intercepted IAT entry with the original function.
  121. //
  122. // Arguments:
  123. // intercept_function Interceptor function
  124. // original_function Receives the original function pointer
  125. //
  126. // Returns: Returns NO_ERROR on success or Windows error code
  127. // as defined in winerror.h
  128. DWORD RestoreImportedFunction(void* intercept_function,
  129. void* original_function,
  130. IMAGE_THUNK_DATA* iat_thunk) {
  131. if (!intercept_function || !original_function || !iat_thunk) {
  132. NOTREACHED();
  133. return ERROR_INVALID_PARAMETER;
  134. }
  135. if (GetIATFunction(iat_thunk) != intercept_function) {
  136. // Check if someone else has intercepted on top of us.
  137. // We cannot unpatch in this case, just raise a red flag.
  138. NOTREACHED();
  139. return ERROR_INVALID_FUNCTION;
  140. }
  141. return internal::ModifyCode(&(iat_thunk->u1.Function), &original_function,
  142. sizeof(original_function));
  143. }
  144. } // namespace
  145. IATPatchFunction::IATPatchFunction() = default;
  146. IATPatchFunction::~IATPatchFunction() {
  147. if (intercept_function_) {
  148. DWORD error = Unpatch();
  149. DCHECK_EQ(static_cast<DWORD>(NO_ERROR), error);
  150. }
  151. }
  152. DWORD IATPatchFunction::Patch(const wchar_t* module,
  153. const char* imported_from_module,
  154. const char* function_name,
  155. void* new_function) {
  156. HMODULE module_handle = LoadLibraryW(module);
  157. if (!module_handle) {
  158. NOTREACHED();
  159. return GetLastError();
  160. }
  161. DWORD error = PatchFromModule(module_handle, imported_from_module,
  162. function_name, new_function);
  163. if (NO_ERROR == error) {
  164. module_handle_ = module_handle;
  165. } else {
  166. FreeLibrary(module_handle);
  167. }
  168. return error;
  169. }
  170. DWORD IATPatchFunction::PatchFromModule(HMODULE module,
  171. const char* imported_from_module,
  172. const char* function_name,
  173. void* new_function) {
  174. DCHECK_EQ(nullptr, original_function_);
  175. DCHECK_EQ(nullptr, iat_thunk_);
  176. DCHECK_EQ(nullptr, intercept_function_);
  177. DCHECK(module);
  178. DWORD error =
  179. InterceptImportedFunction(module, imported_from_module, function_name,
  180. new_function, &original_function_, &iat_thunk_);
  181. if (NO_ERROR == error) {
  182. DCHECK_NE(original_function_, intercept_function_);
  183. intercept_function_ = new_function;
  184. }
  185. return error;
  186. }
  187. DWORD IATPatchFunction::Unpatch() {
  188. DWORD error = RestoreImportedFunction(intercept_function_, original_function_,
  189. iat_thunk_);
  190. DCHECK_EQ(static_cast<DWORD>(NO_ERROR), error);
  191. // Hands off the intercept if we fail to unpatch.
  192. // If IATPatchFunction::Unpatch fails during RestoreImportedFunction
  193. // it means that we cannot safely unpatch the import address table
  194. // patch. In this case its better to be hands off the intercept as
  195. // trying to unpatch again in the destructor of IATPatchFunction is
  196. // not going to be any safer
  197. if (module_handle_)
  198. FreeLibrary(module_handle_);
  199. module_handle_ = nullptr;
  200. intercept_function_ = nullptr;
  201. original_function_ = nullptr;
  202. iat_thunk_ = nullptr;
  203. return error;
  204. }
  205. void* IATPatchFunction::original_function() const {
  206. DCHECK(is_patched());
  207. return original_function_;
  208. }
  209. } // namespace win
  210. } // namespace base