| Index: base/iat_patch.cc
|
| diff --git a/base/iat_patch.cc b/base/iat_patch.cc
|
| index 2417608a82f9e8ac4e5305988f16551b0e838f1a..fb4a79a7d92d4ccec933f711748404af187f30a3 100644
|
| --- a/base/iat_patch.cc
|
| +++ b/base/iat_patch.cc
|
| @@ -181,7 +181,7 @@ IATPatchFunction::~IATPatchFunction() {
|
| }
|
| }
|
|
|
| -DWORD IATPatchFunction::Patch(HMODULE module_handle,
|
| +DWORD IATPatchFunction::Patch(const wchar_t* module,
|
| const char* imported_from_module,
|
| const char* function_name,
|
| void* new_function) {
|
| @@ -189,6 +189,13 @@ DWORD IATPatchFunction::Patch(HMODULE module_handle,
|
| DCHECK_EQ(static_cast<IMAGE_THUNK_DATA*>(NULL), iat_thunk_);
|
| DCHECK_EQ(static_cast<void*>(NULL), intercept_function_);
|
|
|
| + HMODULE module_handle = LoadLibraryW(module);
|
| +
|
| + if (module_handle == NULL) {
|
| + NOTREACHED();
|
| + return GetLastError();
|
| + }
|
| +
|
| DWORD error = InterceptImportedFunction(module_handle,
|
| imported_from_module,
|
| function_name,
|
| @@ -198,32 +205,19 @@ DWORD IATPatchFunction::Patch(HMODULE module_handle,
|
|
|
| if (NO_ERROR == error) {
|
| DCHECK_NE(original_function_, intercept_function_);
|
| + module_handle_ = module_handle;
|
| intercept_function_ = new_function;
|
| + } else {
|
| + FreeLibrary(module_handle);
|
| }
|
|
|
| return error;
|
| }
|
|
|
| DWORD IATPatchFunction::Unpatch() {
|
| - DWORD error = 0;
|
| - MEMORY_BASIC_INFORMATION memory_info = {0};
|
| -
|
| - // If the module has already unloaded, no point trying to unpatch.
|
| - if (!VirtualQuery(original_function_, &memory_info,
|
| - sizeof(memory_info))) {
|
| - error = GetLastError();
|
| - NOTREACHED();
|
| - return error;
|
| - }
|
| -
|
| - if ((memory_info.State & MEM_COMMIT) != MEM_COMMIT) {
|
| - NOTREACHED();
|
| - return ERROR_ACCESS_DENIED;
|
| - }
|
| -
|
| - error = RestoreImportedFunction(intercept_function_,
|
| - original_function_,
|
| - iat_thunk_);
|
| + DWORD error = RestoreImportedFunction(intercept_function_,
|
| + original_function_,
|
| + iat_thunk_);
|
| DCHECK(NO_ERROR == error);
|
|
|
| // Hands off the intercept if we fail to unpatch.
|
| @@ -232,6 +226,9 @@ DWORD IATPatchFunction::Unpatch() {
|
| // patch. In this case its better to be hands off the intercept as
|
| // trying to unpatch again in the destructor of IATPatchFunction is
|
| // not going to be any safer
|
| + if (module_handle_)
|
| + FreeLibrary(module_handle_);
|
| + module_handle_ = NULL;
|
| intercept_function_ = NULL;
|
| original_function_ = NULL;
|
| iat_thunk_ = NULL;
|
| @@ -240,4 +237,3 @@ DWORD IATPatchFunction::Unpatch() {
|
| }
|
|
|
| } // namespace iat_patch
|
| -
|
|
|