5 files changed, 29 insertions, 23 deletions

diff --git a/chrome_frame/bho.cc b/chrome_frame/bho.cc
index f8d44e6..5e0c7f6 100644
--- a/chrome_frame/bho.cc
+++ b/chrome_frame/bho.cc
@@ -312,34 +312,12 @@ void Bho::ProcessOptInUrls(IWebBrowser2* browser, BSTR url) {
   }
 }
 
-namespace {
-// Utility function that prevents the current module from ever being unloaded.
-void PinModule() {
-  FilePath module_path;
-  if (PathService::Get(base::FILE_MODULE, &module_path)) {
-    HMODULE unused;
-    if (!GetModuleHandleEx(GET_MODULE_HANDLE_EX_FLAG_PIN,
-                           module_path.value().c_str(), &unused)) {
-      NOTREACHED() << "Failed to pin module " << module_path.value().c_str()
-                   << " , last error: " << GetLastError();
-    }
-  } else {
-    NOTREACHED() << "Could not get module path.";
-  }
-}
-}  // namespace
-
 bool PatchHelper::InitializeAndPatchProtocolsIfNeeded() {
   bool ret = false;
 
   _pAtlModule->m_csStaticDataInitAndTypeInfo.Lock();
 
   if (state_ == UNKNOWN) {
-    // If we're going to start patching things for reals, we'd better make sure
-    // that we stick around for ever more:
-    if (!IsUnpinnedMode())
-      PinModule();
-
     ProtocolPatchMethod patch_method = GetPatchMethod();
     if (patch_method == PATCH_METHOD_INET_PROTOCOL) {
       g_trans_hooks.InstallHooks();
diff --git a/chrome_frame/buggy_bho_handling.cc b/chrome_frame/buggy_bho_handling.cc
index fd2c96c..edcfb20 100644
--- a/chrome_frame/buggy_bho_handling.cc
+++ b/chrome_frame/buggy_bho_handling.cc
@@ -128,6 +128,7 @@ HRESULT BuggyBhoTls::PatchInvokeMethod(PROC* invoke) {
       hr = E_UNEXPECTED;
       FunctionStub::Destroy(stub);
     } else {
+      PinModule();  // No backing out now.
       ::FlushInstructionCache(::GetCurrentProcess(), invoke, sizeof(PROC));
     }
   }
diff --git a/chrome_frame/utils.cc b/chrome_frame/utils.cc
index c2c2d48..e8dfb4b 100644
--- a/chrome_frame/utils.cc
+++ b/chrome_frame/utils.cc
@@ -942,7 +942,9 @@ bool IsHeadlessMode() {
 }
 
 bool IsUnpinnedMode() {
-  bool unpinned = GetConfigBool(false, kChromeFrameUnpinnedMode);
+  // We only check this value once and then cache it since changing the registry
+  // once we've pinned the DLL won't have any effect.
+  static bool unpinned = GetConfigBool(false, kChromeFrameUnpinnedMode);
   return unpinned;
 }
 
@@ -1390,3 +1392,22 @@ bool CanNavigate(const GURL& url, IInternetSecurityManager* security_manager,
   return true;
 }
 
+void PinModule() {
+  static bool s_pinned = false;
+  if (!s_pinned && !IsUnpinnedMode()) {
+    FilePath module_path;
+    if (PathService::Get(base::FILE_MODULE, &module_path)) {
+      HMODULE unused;
+      if (!GetModuleHandleEx(GET_MODULE_HANDLE_EX_FLAG_PIN,
+                             module_path.value().c_str(), &unused)) {
+        NOTREACHED() << "Failed to pin module " << module_path.value().c_str()
+                     << " , last error: " << GetLastError();
+      } else {
+        s_pinned = true;
+      }
+    } else {
+      NOTREACHED() << "Could not get module path.";
+    }
+  }
+}
+
diff --git a/chrome_frame/utils.h b/chrome_frame/utils.h
index 2c8c700..477084d 100644
--- a/chrome_frame/utils.h
+++ b/chrome_frame/utils.h
@@ -538,4 +538,8 @@ class ChromeFrameUrl {
 bool CanNavigate(const GURL& url, IInternetSecurityManager* security_manager,
                  bool is_privileged);
 
+// Utility function that prevents the current module from ever being unloaded.
+// Call if you make irreversible patches.
+void PinModule();
+
 #endif  // CHROME_FRAME_UTILS_H_
diff --git a/chrome_frame/vtable_patch_manager.cc b/chrome_frame/vtable_patch_manager.cc
index 10852ee..9f60d3a 100644
--- a/chrome_frame/vtable_patch_manager.cc
+++ b/chrome_frame/vtable_patch_manager.cc
@@ -14,6 +14,7 @@
 #include "base/scoped_ptr.h"
 
 #include "chrome_frame/function_stub.h"
+#include "chrome_frame/utils.h"
 
 namespace vtable_patch {
 
@@ -144,6 +145,7 @@ HRESULT PatchInterfaceMethods(void* unknown, MethodPatchInfo* patches) {
     } else {
       // Success, save the stub we created.
       it->stub_ = stub;
+      PinModule();
     }
   }