1 // Copyright 2011 The Chromium Authors 2 // Use of this source code is governed by a BSD-style license that can be 3 // found in the LICENSE file. 4 5 #ifndef BASE_WIN_IAT_PATCH_FUNCTION_H_ 6 #define BASE_WIN_IAT_PATCH_FUNCTION_H_ 7 8 #include <windows.h> 9 10 #include "base/base_export.h" 11 #include "base/memory/raw_ptr.h" 12 13 namespace base { 14 namespace win { 15 16 // A class that encapsulates Import Address Table patching helpers and restores 17 // the original function in the destructor. 18 // 19 // It will intercept functions for a specific DLL imported from another DLL. 20 // This is the case when, for example, we want to intercept 21 // CertDuplicateCertificateContext function (exported from crypt32.dll) called 22 // by wininet.dll. 23 class BASE_EXPORT IATPatchFunction { 24 public: 25 IATPatchFunction(); 26 27 IATPatchFunction(const IATPatchFunction&) = delete; 28 IATPatchFunction& operator=(const IATPatchFunction&) = delete; 29 30 ~IATPatchFunction(); 31 32 // Intercept a function in an import table of a specific 33 // module. Save the original function and the import 34 // table address. These values will be used later 35 // during Unpatch 36 // 37 // Arguments: 38 // module Module to be intercepted 39 // imported_from_module Module that exports the 'function_name' 40 // function_name Name of the API to be intercepted 41 // 42 // Returns: Windows error code (winerror.h). NO_ERROR if successful 43 // 44 // Note: Patching a function will make the IAT patch take some "ownership" on 45 // |module|. It will LoadLibrary(module) to keep the DLL alive until a call 46 // to Unpatch(), which will call FreeLibrary() and allow the module to be 47 // unloaded. The idea is to help prevent the DLL from going away while a 48 // patch is still active. 49 DWORD Patch(const wchar_t* module, 50 const char* imported_from_module, 51 const char* function_name, 52 void* new_function); 53 54 // Same as Patch(), but uses a handle to a |module| instead of the DLL name. 55 DWORD PatchFromModule(HMODULE module, 56 const char* imported_from_module, 57 const char* function_name, 58 void* new_function); 59 60 // Unpatch the IAT entry using internally saved original 61 // function. 62 // 63 // Returns: Windows error code (winerror.h). NO_ERROR if successful 64 DWORD Unpatch(); 65 is_patched()66 bool is_patched() const { return (nullptr != intercept_function_); } 67 68 void* original_function() const; 69 70 private: 71 HMODULE module_handle_ = nullptr; 72 raw_ptr<void> intercept_function_ = nullptr; 73 raw_ptr<void> original_function_ = nullptr; 74 raw_ptr<IMAGE_THUNK_DATA> iat_thunk_ = nullptr; 75 }; 76 77 } // namespace win 78 } // namespace base 79 80 #endif // BASE_WIN_IAT_PATCH_FUNCTION_H_ 81