1 files changed, 17 insertions, 21 deletions

diff --git a/base/iat_patch.cc b/base/iat_patch.cc
index 2417608..fb4a79a 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
-