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

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
-

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_;

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

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);

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;

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,