Fix memory leak in unit tests for ProtocolHandlerRegistry introduced by observers.

Unit tests need to call Finalize at teardown. This has an additional change to
http://codereview.chromium.org/7080023/.

BUG=83556

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

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

6 files changed, 129 insertions, 15 deletions

diff --git a/chrome/browser/custom_handlers/protocol_handler_registry.cc b/chrome/browser/custom_handlers/protocol_handler_registry.cc
index 09585f9..f2b30f7 100644
--- a/chrome/browser/custom_handlers/protocol_handler_registry.cc
+++ b/chrome/browser/custom_handlers/protocol_handler_registry.cc
@@ -6,13 +6,14 @@
 
 #include <utility>
 
+#include "base/logging.h"
 #include "base/memory/scoped_ptr.h"
+#include "base/stl_util-inl.h"
 #include "chrome/browser/custom_handlers/protocol_handler.h"
 #include "chrome/browser/custom_handlers/register_protocol_handler_infobar_delegate.h"
 #include "chrome/browser/net/chrome_url_request_context.h"
 #include "chrome/browser/prefs/pref_service.h"
 #include "chrome/browser/profiles/profile_io_data.h"
-#include "chrome/browser/shell_integration.h"
 #include "chrome/common/pref_names.h"
 #include "content/browser/browser_thread.h"
 #include "content/browser/child_process_security_policy.h"
@@ -30,6 +31,12 @@ ProtocolHandlerRegistry::ProtocolHandlerRegistry(Profile* profile,
 }
 
 ProtocolHandlerRegistry::~ProtocolHandlerRegistry() {
+  DCHECK(default_client_observers_.empty());
+}
+
+void ProtocolHandlerRegistry::Finalize() {
+  DCHECK(BrowserThread::CurrentlyOn(BrowserThread::UI));
+  STLDeleteElements(&default_client_observers_);
 }
 
 const ProtocolHandlerRegistry::ProtocolHandlerList*
@@ -144,6 +151,7 @@ ProtocolHandlerRegistry::GetHandlersFromPref(const char* pref_name) const {
 }
 
 void ProtocolHandlerRegistry::Load() {
+  DCHECK(BrowserThread::CurrentlyOn(BrowserThread::UI));
   base::AutoLock auto_lock(lock_);
   is_loading_ = true;
   PrefService* prefs = profile_->GetPrefs();
@@ -170,6 +178,20 @@ void ProtocolHandlerRegistry::Load() {
     IgnoreProtocolHandler(ProtocolHandler::CreateProtocolHandler(*p));
   }
   is_loading_ = false;
+
+  // For each default protocol handler, check that we are still registered
+  // with the OS as the default application.
+  for (ProtocolHandlerMap::const_iterator p = default_handlers_.begin();
+       p != default_handlers_.end(); ++p) {
+    ProtocolHandler handler = p->second;
+    DefaultClientObserver* observer = new DefaultClientObserver(this);
+    scoped_refptr<ShellIntegration::DefaultProtocolClientWorker> worker;
+    worker = new ShellIntegration::DefaultProtocolClientWorker(
+        observer, handler.protocol());
+    observer->SetWorker(worker);
+    default_client_observers_.push_back(observer);
+    worker->StartCheckIsDefault();
+  }
 }
 
 void ProtocolHandlerRegistry::Save() {
@@ -277,22 +299,36 @@ bool ProtocolHandlerRegistry::IsHandledProtocolInternal(
 void ProtocolHandlerRegistry::RemoveHandler(const ProtocolHandler& handler) {
   {
     base::AutoLock auto_lock(lock_);
-    ProtocolHandlerList& handlers = protocol_handlers_[handler.protocol()];
-    ProtocolHandlerList::iterator p =
-        std::find(handlers.begin(), handlers.end(), handler);
-    if (p != handlers.end()) {
-      handlers.erase(p);
-    }
+    RemoveHandlerInternal(handler);
+  }
+  NotifyChanged();
+}
 
-    ProtocolHandlerMap::iterator q = default_handlers_.find(handler.protocol());
-    if (q != default_handlers_.end() && q->second == handler) {
-      default_handlers_.erase(q);
-    }
+void ProtocolHandlerRegistry::RemoveHandlerInternal(
+    const ProtocolHandler& handler) {
+  ProtocolHandlerList& handlers = protocol_handlers_[handler.protocol()];
+  ProtocolHandlerList::iterator p =
+      std::find(handlers.begin(), handlers.end(), handler);
+  if (p != handlers.end()) {
+    handlers.erase(p);
+  }
+  ProtocolHandlerMap::iterator q = default_handlers_.find(handler.protocol());
+  if (q != default_handlers_.end() && q->second == handler) {
+    default_handlers_.erase(q);
+  }
 
-    if (!IsHandledProtocolInternal(handler.protocol())) {
-      delegate_->DeregisterExternalHandler(handler.protocol());
-    }
-    Save();
+  if (!IsHandledProtocolInternal(handler.protocol())) {
+    delegate_->DeregisterExternalHandler(handler.protocol());
+  }
+  Save();
+}
+
+void ProtocolHandlerRegistry::RemoveDefaultHandler(const std::string& scheme) {
+  {
+    base::AutoLock auto_lock(lock_);
+    ProtocolHandler current_default = GetHandlerForInternal(scheme);
+    if (!current_default.IsEmpty())
+      RemoveHandlerInternal(current_default);
   }
   NotifyChanged();
 }
@@ -480,3 +516,32 @@ bool ProtocolHandlerRegistry::Delegate::IsExternalHandlerRegistered(
   // in ProfileIOData.
   return ProfileIOData::IsHandledProtocol(protocol);
 }
+
+// DefaultClientObserver ------------------------------------------------------
+
+ProtocolHandlerRegistry::DefaultClientObserver::DefaultClientObserver(
+    ProtocolHandlerRegistry* registry)
+    : registry_(registry), worker_(NULL) {
+  DCHECK(registry_);
+}
+
+ProtocolHandlerRegistry::DefaultClientObserver::~DefaultClientObserver() {
+  if (worker_)
+    worker_->ObserverDestroyed();
+}
+
+void
+ProtocolHandlerRegistry::DefaultClientObserver::SetDefaultWebClientUIState(
+    ShellIntegration::DefaultWebClientUIState state) {
+  if (worker_) {
+    if (state == ShellIntegration::STATE_NOT_DEFAULT)
+      registry_->RemoveDefaultHandler(worker_->protocol());
+  } else {
+    NOTREACHED();
+  }
+}
+
+void ProtocolHandlerRegistry::DefaultClientObserver::SetWorker(
+    ShellIntegration::DefaultProtocolClientWorker* worker) {
+  worker_ = worker;
+}
diff --git a/chrome/browser/custom_handlers/protocol_handler_registry.h b/chrome/browser/custom_handlers/protocol_handler_registry.h
index 7532bba..fa588c5 100644
--- a/chrome/browser/custom_handlers/protocol_handler_registry.h
+++ b/chrome/browser/custom_handlers/protocol_handler_registry.h
@@ -16,6 +16,7 @@
 #include "base/values.h"
 #include "chrome/browser/custom_handlers/protocol_handler.h"
 #include "chrome/browser/profiles/profile.h"
+#include "chrome/browser/shell_integration.h"
 #include "content/common/notification_service.h"
 #include "net/url_request/url_request.h"
 #include "net/url_request/url_request_job.h"
@@ -41,9 +42,35 @@ class ProtocolHandlerRegistry
     virtual void RegisterWithOSAsDefaultClient(const std::string& protocol);
   };
 
+  class DefaultClientObserver
+      : public ShellIntegration::DefaultWebClientObserver {
+   public:
+    explicit DefaultClientObserver(ProtocolHandlerRegistry* registry);
+    virtual ~DefaultClientObserver();
+
+    // Get response from the worker regarding whether Chrome is the default
+    // handler for the protocol.
+    virtual void SetDefaultWebClientUIState(
+        ShellIntegration::DefaultWebClientUIState state);
+
+    // Give the observer a handle to the worker, so we can find out the protocol
+    // when we're called and also tell the worker if we get deleted.
+    void SetWorker(ShellIntegration::DefaultProtocolClientWorker* worker);
+
+   private:
+    // This is a raw pointer, not reference counted, intentionally. In general
+    // subclasses of DefaultWebClientObserver are not able to be refcounted
+    // e.g. the browser options page
+    ProtocolHandlerRegistry* registry_;
+    scoped_refptr<ShellIntegration::DefaultProtocolClientWorker> worker_;
+
+    DISALLOW_COPY_AND_ASSIGN(DefaultClientObserver);
+  };
+
   typedef std::map<std::string, ProtocolHandler> ProtocolHandlerMap;
   typedef std::vector<ProtocolHandler> ProtocolHandlerList;
   typedef std::map<std::string, ProtocolHandlerList> ProtocolHandlerMultiMap;
+  typedef std::vector<DefaultClientObserver*> DefaultClientObserverList;
 
   ProtocolHandlerRegistry(Profile* profile, Delegate* delegate);
   ~ProtocolHandlerRegistry();
@@ -97,6 +124,9 @@ class ProtocolHandlerRegistry
   // Removes the given protocol handler from the registry.
   void RemoveHandler(const ProtocolHandler& handler);
 
+  // Remove the default handler for the given protocol.
+  void RemoveDefaultHandler(const std::string& scheme);
+
   // Returns the default handler for this protocol, or an empty handler if none
   // exists.
   const ProtocolHandler& GetHandlerFor(const std::string& scheme) const;
@@ -113,6 +143,10 @@ class ProtocolHandlerRegistry
   // will not handle requests.
   void Disable();
 
+  // This is called by the UI thread when the system is shutting down. This
+  // does finalization which must be done on the UI thread.
+  void Finalize();
+
   // Registers the preferences that we store registered protocol handlers in.
   static void RegisterPrefs(PrefService* prefService);
 
@@ -138,6 +172,9 @@ class ProtocolHandlerRegistry
   // exists.
   const ProtocolHandler& GetHandlerForInternal(const std::string& scheme) const;
 
+  // Removes the given protocol handler from the registry.
+  void RemoveHandlerInternal(const ProtocolHandler& handler);
+
   // Saves a user's registered protocol handlers.
   void Save();
 
@@ -204,6 +241,8 @@ class ProtocolHandlerRegistry
   // TODO(koz): Remove the necessity of this lock by using message passing.
   mutable base::Lock lock_;
 
+  DefaultClientObserverList default_client_observers_;
+
   friend class ProtocolHandlerRegistryTest;
 
   DISALLOW_COPY_AND_ASSIGN(ProtocolHandlerRegistry);
diff --git a/chrome/browser/custom_handlers/protocol_handler_registry_unittest.cc b/chrome/browser/custom_handlers/protocol_handler_registry_unittest.cc
index eb268b1..bc0829b 100644
--- a/chrome/browser/custom_handlers/protocol_handler_registry_unittest.cc
+++ b/chrome/browser/custom_handlers/protocol_handler_registry_unittest.cc
@@ -140,6 +140,10 @@ class ProtocolHandlerRegistryTest : public RenderViewHostTestHarness {
     ProtocolHandlerRegistry::RegisterPrefs(pref_service());
   }
 
+  virtual void TearDown() {
+    registry_->Finalize();
+  }
+
   BrowserThread ui_thread_;
   FakeDelegate* delegate_;
   scoped_ptr<TestingProfile> profile_;
diff --git a/chrome/browser/profiles/profile_impl.cc b/chrome/browser/profiles/profile_impl.cc
index e25d9f7..2e0396f 100644
--- a/chrome/browser/profiles/profile_impl.cc
+++ b/chrome/browser/profiles/profile_impl.cc
@@ -682,6 +682,9 @@ ProfileImpl::~ProfileImpl() {
   if (pref_proxy_config_tracker_)
     pref_proxy_config_tracker_->DetachFromPrefService();
 
+  if (protocol_handler_registry_)
+    protocol_handler_registry_->Finalize();
+
   // This causes the Preferences file to be written to disk.
   MarkAsCleanShutdown();
 }
diff --git a/chrome/browser/shell_integration.cc b/chrome/browser/shell_integration.cc
index fd0e035..c332965 100644
--- a/chrome/browser/shell_integration.cc
+++ b/chrome/browser/shell_integration.cc
@@ -92,6 +92,7 @@ void ShellIntegration::DefaultWebClientWorker::StartSetAsDefault() {
 void ShellIntegration::DefaultWebClientWorker::ObserverDestroyed() {
   // Our associated view has gone away, so we shouldn't call back to it if
   // our worker thread returns after the view is dead.
+  DCHECK(BrowserThread::CurrentlyOn(BrowserThread::UI));
   observer_ = NULL;
 }
 
diff --git a/chrome/browser/shell_integration.h b/chrome/browser/shell_integration.h
index dd295a8..ab570ce 100644
--- a/chrome/browser/shell_integration.h
+++ b/chrome/browser/shell_integration.h
@@ -246,6 +246,8 @@ class ShellIntegration {
     DefaultProtocolClientWorker(DefaultWebClientObserver* observer,
                                 const std::string& protocol);
 
+    const std::string& protocol() const { return protocol_; }
+
    private:
     virtual ~DefaultProtocolClientWorker() {}