Try a new approach to fixing IAT unpatch crashes when the DLL is gone.

Have the IAT patcher take some "ownership" of the DLL, by taking a library name and then calling LoadLibrary() / FreeLibrary() to manage the reference count.  This means as long is there isn't some other reference count balancing bug happening in the process, the DLL will never be unloaded while we are patched.

This effectively reverts r9929, the VirtualQuery additional checks are removed.

BUG=7701

Review URL: http://codereview.chromium.org/21453

git-svn-id: svn://svn.chromium.org/chrome/trunk/src@10467 0039d316-1c4b-4281-b951-d872f2087c98

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
-