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

6 files changed, 99 insertions, 76 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
-
diff --git a/base/iat_patch.h b/base/iat_patch.h
index 819037c..9e470d6 100644
--- a/base/iat_patch.h
+++ b/base/iat_patch.h
@@ -80,13 +80,19 @@ class IATPatchFunction {
   // during Unpatch
   //
   // Arguments:
-  // module_handle          Module to be intercepted
+  // module                 Module to be intercepted
   // imported_from_module   Module that exports the 'function_name'
   // function_name          Name of the API to be intercepted
   //
   // Returns: Windows error code (winerror.h). NO_ERROR if successful
   //
-  DWORD Patch(HMODULE module_handle,
+  // Note: Patching a function will make the IAT patch take some "ownership" on
+  // |module|.  It will LoadLibrary(module) to keep the DLL alive until a call
+  // to Unpatch(), which will call FreeLibrary() and allow the module to be
+  // unloaded.  The idea is to help prevent the DLL from going away while a
+  // patch is still active.
+  //
+  DWORD Patch(const wchar_t* module,
               const char* imported_from_module,
               const char* function_name,
               void* new_function);
@@ -103,6 +109,7 @@ class IATPatchFunction {
   }
 
  private:
+  HMODULE module_handle_;
   void* intercept_function_;
   void* original_function_;
   IMAGE_THUNK_DATA* iat_thunk_;
diff --git a/chrome/browser/autocomplete/autocomplete_edit.cc b/chrome/browser/autocomplete/autocomplete_edit.cc
index b332832..862c2ed 100644
--- a/chrome/browser/autocomplete/autocomplete_edit.cc
+++ b/chrome/browser/autocomplete/autocomplete_edit.cc
@@ -7,8 +7,10 @@
 #include <locale>
 
 #include "base/base_drag_source.h"
+#include "base/basictypes.h"
 #include "base/clipboard.h"
 #include "base/iat_patch.h"
+#include "base/lazy_instance.h"
 #include "base/ref_counted.h"
 #include "base/scoped_clipboard_writer.h"
 #include "base/string_util.h"
@@ -633,26 +635,81 @@ AutocompleteEditView::ScopedSuspendUndo::~ScopedSuspendUndo() {
 // TODO (jcampan): these colors should be derived from the system colors to
 // ensure they show properly. Bug #948807.
 // Colors used to emphasize the scheme in the URL.
-static const COLORREF kSecureSchemeColor = RGB(0, 150, 20);
-static const COLORREF kInsecureSchemeColor = RGB(200, 0, 0);
+
+namespace {
+
+const COLORREF kSecureSchemeColor = RGB(0, 150, 20);
+const COLORREF kInsecureSchemeColor = RGB(200, 0, 0);
 
 // Colors used to strike-out the scheme when it is insecure.
-static const SkColor kSchemeStrikeoutColor = SkColorSetRGB(210, 0, 0);
-static const SkColor kSchemeSelectedStrikeoutColor =
+const SkColor kSchemeStrikeoutColor = SkColorSetRGB(210, 0, 0);
+const SkColor kSchemeSelectedStrikeoutColor =
     SkColorSetRGB(255, 255, 255);
 
 // These are used to hook the CRichEditCtrl's calls to BeginPaint() and
 // EndPaint() and provide a memory DC instead.  See OnPaint().
-static HWND edit_hwnd = NULL;
-static PAINTSTRUCT paint_struct;
+HWND edit_hwnd = NULL;
+PAINTSTRUCT paint_struct;
+
+HDC BeginPaintIntercept(HWND hWnd, LPPAINTSTRUCT lpPaint) {
+  if (!edit_hwnd || (hWnd != edit_hwnd))
+    return ::BeginPaint(hWnd, lpPaint);
+
+  *lpPaint = paint_struct;
+  return paint_struct.hdc;
+}
+
+BOOL EndPaintIntercept(HWND hWnd, const PAINTSTRUCT* lpPaint) {
+  return (edit_hwnd && (hWnd == edit_hwnd)) ?
+      true : ::EndPaint(hWnd, lpPaint);
+}
 
 // Returns a lazily initialized property bag accessor for saving our state in a
 // TabContents.
-static PropertyAccessor<AutocompleteEditState>* GetStateAccessor() {
+PropertyAccessor<AutocompleteEditState>* GetStateAccessor() {
   static PropertyAccessor<AutocompleteEditState> state;
   return &state;
 }
 
+class PaintPatcher {
+ public:
+  PaintPatcher() : refcount_(0) { }
+  ~PaintPatcher() { DCHECK(refcount_ == 0); }
+
+  void RefPatch() {
+    if (refcount_ == 0) {
+      DCHECK(!begin_paint_.is_patched());
+      DCHECK(!end_paint_.is_patched());
+      begin_paint_.Patch(L"riched20.dll", "user32.dll", "BeginPaint",
+                         &BeginPaintIntercept);
+      end_paint_.Patch(L"riched20.dll", "user32.dll", "EndPaint",
+                       &EndPaintIntercept);
+    }
+    ++refcount_;
+  }
+
+  void DerefPatch() {
+    DCHECK(begin_paint_.is_patched());
+    DCHECK(end_paint_.is_patched());
+    --refcount_;
+    if (refcount_ == 0) {
+      begin_paint_.Unpatch();
+      end_paint_.Unpatch();
+    }
+  }
+
+ private:
+  size_t refcount_;
+  iat_patch::IATPatchFunction begin_paint_;
+  iat_patch::IATPatchFunction end_paint_;
+
+  DISALLOW_COPY_AND_ASSIGN(PaintPatcher);
+};
+
+base::LazyInstance<PaintPatcher> g_paint_patcher(base::LINKER_INITIALIZED);
+
+}  // namespace
+
 AutocompleteEditView::AutocompleteEditView(
     const ChromeFont& font,
     AutocompleteEditController* controller,
@@ -685,22 +742,7 @@ AutocompleteEditView::AutocompleteEditView(
 
   saved_selection_for_focus_change_.cpMin = -1;
 
-  // Statics used for global patching of riched20.dll.
-  static HMODULE richedit_module = NULL;
-  static iat_patch::IATPatchFunction patch_begin_paint;
-  static iat_patch::IATPatchFunction patch_end_paint;
-
-  if (!richedit_module) {
-    richedit_module = LoadLibrary(L"riched20.dll");
-    if (richedit_module) {
-      DCHECK(!patch_begin_paint.is_patched());
-      patch_begin_paint.Patch(richedit_module, "user32.dll", "BeginPaint",
-                              &BeginPaintIntercept);
-      DCHECK(!patch_end_paint.is_patched());
-      patch_end_paint.Patch(richedit_module, "user32.dll", "EndPaint",
-                            &EndPaintIntercept);
-    }
-  }
+  g_paint_patcher.Pointer()->RefPatch();
 
   Create(hwnd, 0, 0, 0, l10n_util::GetExtendedStyles());
   SetReadOnly(popup_window_mode_);
@@ -787,6 +829,11 @@ AutocompleteEditView::~AutocompleteEditView() {
       NotificationType::AUTOCOMPLETE_EDIT_DESTROYED,
       Source<AutocompleteEditView>(this),
       NotificationService::NoDetails());
+
+  // We balance our reference count and unpatch when the last instance has
+  // been destroyed.  This prevents us from relying on the AtExit or static
+  // destructor sequence to do our unpatching, which is generally fragile.
+  g_paint_patcher.Pointer()->DerefPatch();
 }
 
 void AutocompleteEditView::SaveStateToTab(TabContents* tab) {
@@ -1374,23 +1421,6 @@ bool AutocompleteEditView::SchemeEnd(LPTSTR edit_text,
          (edit_text[current_pos + 2] == '/');
 }
 
-// static
-HDC AutocompleteEditView::BeginPaintIntercept(HWND hWnd,
-                                              LPPAINTSTRUCT lpPaint) {
-  if (!edit_hwnd || (hWnd != edit_hwnd))
-    return ::BeginPaint(hWnd, lpPaint);
-
-  *lpPaint = paint_struct;
-  return paint_struct.hdc;
-}
-
-// static
-BOOL AutocompleteEditView::EndPaintIntercept(HWND hWnd,
-                                             const PAINTSTRUCT* lpPaint) {
-  return (edit_hwnd && (hWnd == edit_hwnd)) ?
-      true : ::EndPaint(hWnd, lpPaint);
-}
-
 void AutocompleteEditView::OnChar(TCHAR ch, UINT repeat_count, UINT flags) {
   // Don't let alt-enter beep.  Not sure this is necessary, as the standard
   // alt-enter will hit DiscardWMSysChar() and get thrown away, and
diff --git a/chrome/browser/autocomplete/autocomplete_edit.h b/chrome/browser/autocomplete/autocomplete_edit.h
index ffb98c4..5fd07d9 100755
--- a/chrome/browser/autocomplete/autocomplete_edit.h
+++ b/chrome/browser/autocomplete/autocomplete_edit.h
@@ -669,10 +669,6 @@ class AutocompleteEditView
   // Returns true if |edit_text| starting at |current_pos| is "://".
   static bool SchemeEnd(LPTSTR edit_text, int current_pos, int length);
 
-  // Intercepts.  See OnPaint().
-  static HDC WINAPI BeginPaintIntercept(HWND hWnd, LPPAINTSTRUCT lpPaint);
-  static BOOL WINAPI EndPaintIntercept(HWND hWnd, CONST PAINTSTRUCT* lpPaint);
-
   // Message handlers
   void OnChar(TCHAR ch, UINT repeat_count, UINT flags);
   void OnContextMenu(HWND window, const CPoint& point);
diff --git a/webkit/glue/plugins/webplugin_delegate_impl.cc b/webkit/glue/plugins/webplugin_delegate_impl.cc
index 4163b50..564659d 100644
--- a/webkit/glue/plugins/webplugin_delegate_impl.cc
+++ b/webkit/glue/plugins/webplugin_delegate_impl.cc
@@ -155,8 +155,7 @@ WebPluginDelegateImpl::WebPluginDelegateImpl(
       handle_event_depth_(0),
       user_gesture_message_posted_(false),
 #pragma warning(suppress: 4355)  // can use this
-      user_gesture_msg_factory_(this),
-      plugin_module_handle_(NULL) {
+      user_gesture_msg_factory_(this) {
   memset(&window_, 0, sizeof(window_));
 
   const WebPluginInfo& plugin_info = instance_->plugin_lib()->plugin_info();
@@ -201,8 +200,6 @@ WebPluginDelegateImpl::WebPluginDelegateImpl(
     quirks_ |= PLUGIN_QUIRK_PATCH_TRACKPOPUP_MENU;
     quirks_ |= PLUGIN_QUIRK_PATCH_SETCURSOR;
   }
-
-  plugin_module_handle_ = ::GetModuleHandle(plugin_info.path.value().c_str());
 }
 
 WebPluginDelegateImpl::~WebPluginDelegateImpl() {
@@ -283,7 +280,7 @@ bool WebPluginDelegateImpl::Initialize(const GURL& url,
   if (windowless_ && !g_iat_patch_track_popup_menu.Pointer()->is_patched() &&
       (quirks_ & PLUGIN_QUIRK_PATCH_TRACKPOPUP_MENU)) {
     g_iat_patch_track_popup_menu.Pointer()->Patch(
-        plugin_module_handle_, "user32.dll", "TrackPopupMenu",
+        GetPluginPath().value().c_str(), "user32.dll", "TrackPopupMenu",
         WebPluginDelegateImpl::TrackPopupMenuPatch);
   }
 
@@ -296,7 +293,7 @@ bool WebPluginDelegateImpl::Initialize(const GURL& url,
   if (windowless_ && !g_iat_patch_set_cursor.Pointer()->is_patched() &&
       (quirks_ & PLUGIN_QUIRK_PATCH_SETCURSOR)) {
     g_iat_patch_set_cursor.Pointer()->Patch(
-        plugin_module_handle_, "user32.dll", "SetCursor",
+        GetPluginPath().value().c_str(), "user32.dll", "SetCursor",
         WebPluginDelegateImpl::SetCursorPatch);
   }
   return true;
diff --git a/webkit/glue/plugins/webplugin_delegate_impl.h b/webkit/glue/plugins/webplugin_delegate_impl.h
index 0474d26..27e17cf 100644
--- a/webkit/glue/plugins/webplugin_delegate_impl.h
+++ b/webkit/glue/plugins/webplugin_delegate_impl.h
@@ -247,9 +247,6 @@ class WebPluginDelegateImpl : public WebPluginDelegate {
 #endif
 
 #if defined(OS_WIN)
-  // The plugin module handle.
-  HMODULE plugin_module_handle_;
-
   // TrackPopupMenu interceptor. Parameters are the same as the Win32 function
   // TrackPopupMenu.
   static BOOL WINAPI TrackPopupMenuPatch(HMENU menu, unsigned int flags, int x,