Change KeepAliveImpl to handle extension unload and shutdown correctly.

Previously, if a KeepAlive was active (such as when a mojo-backed API
call is in-progress) when an extension is reloaded, the KeepAlive could
be destroyed after the ProcessManager resets the KeepAlive count,
leading to an incorrect KeepAlive count. If a KeepAlive was active
during browser shutdown, it would attempt to access its ProcessManager
after KeyedServices have been destroyed, which crashes.

BUG=509959

Review URL: https://codereview.chromium.org/1236133002

Cr-Commit-Position: refs/heads/master@{#338806}

3 files changed, 92 insertions, 3 deletions

diff --git a/extensions/browser/mojo/keep_alive_impl.cc b/extensions/browser/mojo/keep_alive_impl.cc
index 9b80598..adfccaa 100644
--- a/extensions/browser/mojo/keep_alive_impl.cc
+++ b/extensions/browser/mojo/keep_alive_impl.cc
@@ -4,6 +4,8 @@
 
 #include "extensions/browser/mojo/keep_alive_impl.h"
 
+#include "base/bind.h"
+#include "extensions/browser/extension_registry.h"
 #include "extensions/browser/process_manager.h"
 
 namespace extensions {
@@ -18,11 +20,31 @@ void KeepAliveImpl::Create(content::BrowserContext* context,
 KeepAliveImpl::KeepAliveImpl(content::BrowserContext* context,
                              const Extension* extension,
                              mojo::InterfaceRequest<KeepAlive> request)
-    : context_(context), extension_(extension), binding_(this, request.Pass()) {
+    : context_(context),
+      extension_(extension),
+      extension_registry_observer_(this),
+      binding_(this, request.Pass()) {
   ProcessManager::Get(context_)->IncrementLazyKeepaliveCount(extension_);
+  binding_.set_connection_error_handler(
+      base::Bind(&KeepAliveImpl::OnDisconnected, base::Unretained(this)));
+  extension_registry_observer_.Add(ExtensionRegistry::Get(context_));
 }
 
-KeepAliveImpl::~KeepAliveImpl() {
+KeepAliveImpl::~KeepAliveImpl() = default;
+
+void KeepAliveImpl::OnExtensionUnloaded(
+    content::BrowserContext* browser_context,
+    const Extension* extension,
+    UnloadedExtensionInfo::Reason reason) {
+  if (browser_context == context_ && extension == extension_)
+    delete this;
+}
+
+void KeepAliveImpl::OnShutdown(ExtensionRegistry* registry) {
+  delete this;
+}
+
+void KeepAliveImpl::OnDisconnected() {
   ProcessManager::Get(context_)->DecrementLazyKeepaliveCount(extension_);
 }
 
diff --git a/extensions/browser/mojo/keep_alive_impl.h b/extensions/browser/mojo/keep_alive_impl.h
index 58108b8..df365f3 100644
--- a/extensions/browser/mojo/keep_alive_impl.h
+++ b/extensions/browser/mojo/keep_alive_impl.h
@@ -6,6 +6,8 @@
 #define EXTENSIONS_BROWSER_MOJO_KEEP_ALIVE_IMPL_H_
 
 #include "base/callback.h"
+#include "base/scoped_observer.h"
+#include "extensions/browser/extension_registry_observer.h"
 #include "extensions/common/mojo/keep_alive.mojom.h"
 #include "third_party/mojo/src/mojo/public/cpp/bindings/interface_request.h"
 #include "third_party/mojo/src/mojo/public/cpp/bindings/strong_binding.h"
@@ -19,7 +21,7 @@ class Extension;
 
 // An RAII mojo service implementation for extension keep alives. This adds a
 // keep alive on construction and removes it on destruction.
-class KeepAliveImpl : public KeepAlive {
+class KeepAliveImpl : public KeepAlive, public ExtensionRegistryObserver {
  public:
   // Create a keep alive for |extension| running in |context| and connect it to
   // |request|. When the requester closes its pipe, the keep alive ends.
@@ -33,8 +35,18 @@ class KeepAliveImpl : public KeepAlive {
                 mojo::InterfaceRequest<KeepAlive> request);
   ~KeepAliveImpl() override;
 
+  // ExtensionRegistryObserver overrides.
+  void OnExtensionUnloaded(content::BrowserContext* browser_context,
+                           const Extension* extension,
+                           UnloadedExtensionInfo::Reason reason) override;
+  void OnShutdown(ExtensionRegistry* registry) override;
+
+  // Invoked when the mojo connection is disconnected.
+  void OnDisconnected();
+
   content::BrowserContext* context_;
   const Extension* extension_;
+  ScopedObserver<ExtensionRegistry, KeepAliveImpl> extension_registry_observer_;
   mojo::StrongBinding<KeepAlive> binding_;
 
   DISALLOW_COPY_AND_ASSIGN(KeepAliveImpl);
diff --git a/extensions/browser/mojo/keep_alive_impl_unittest.cc b/extensions/browser/mojo/keep_alive_impl_unittest.cc
index 4bdb588b..0268358 100644
--- a/extensions/browser/mojo/keep_alive_impl_unittest.cc
+++ b/extensions/browser/mojo/keep_alive_impl_unittest.cc
@@ -7,6 +7,7 @@
 #include "base/message_loop/message_loop.h"
 #include "base/run_loop.h"
 #include "content/public/browser/notification_service.h"
+#include "extensions/browser/extension_registry.h"
 #include "extensions/browser/extensions_test.h"
 #include "extensions/browser/process_manager.h"
 #include "extensions/common/extension_builder.h"
@@ -98,4 +99,58 @@ TEST_F(KeepAliveTest, TwoKeepAlives) {
   EXPECT_EQ(0, GetKeepAliveCount());
 }
 
+TEST_F(KeepAliveTest, UnloadExtension) {
+  mojo::InterfacePtr<KeepAlive> keep_alive;
+  CreateKeepAlive(mojo::GetProxy(&keep_alive));
+  EXPECT_EQ(1, GetKeepAliveCount());
+
+  scoped_refptr<const Extension> other_extension =
+      ExtensionBuilder()
+          .SetManifest(
+               DictionaryBuilder()
+                   .Set("name", "app")
+                   .Set("version", "1")
+                   .Set("manifest_version", 2)
+                   .Set("app", DictionaryBuilder().Set(
+                                   "background",
+                                   DictionaryBuilder().Set(
+                                       "scripts",
+                                       ListBuilder().Append("background.js")))))
+          .SetID("bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb")
+          .Build();
+
+  ExtensionRegistry::Get(browser_context())
+      ->TriggerOnUnloaded(other_extension.get(),
+                          UnloadedExtensionInfo::REASON_DISABLE);
+  EXPECT_EQ(1, GetKeepAliveCount());
+
+  ExtensionRegistry::Get(browser_context())
+      ->TriggerOnUnloaded(extension(), UnloadedExtensionInfo::REASON_DISABLE);
+  // When its extension is unloaded, the KeepAliveImpl should not modify the
+  // keep-alive count for its extension. However, ProcessManager resets its
+  // keep-alive count for an unloaded extension.
+  EXPECT_EQ(0, GetKeepAliveCount());
+
+  // Wait for |keep_alive| to disconnect.
+  base::RunLoop run_loop;
+  keep_alive.set_connection_error_handler(run_loop.QuitClosure());
+  run_loop.Run();
+}
+
+TEST_F(KeepAliveTest, Shutdown) {
+  mojo::InterfacePtr<KeepAlive> keep_alive;
+  CreateKeepAlive(mojo::GetProxy(&keep_alive));
+  EXPECT_EQ(1, GetKeepAliveCount());
+
+  ExtensionRegistry::Get(browser_context())->Shutdown();
+  // After a shutdown event, the KeepAliveImpl should not access its
+  // ProcessManager and so the keep-alive count should remain unchanged.
+  EXPECT_EQ(1, GetKeepAliveCount());
+
+  // Wait for |keep_alive| to disconnect.
+  base::RunLoop run_loop;
+  keep_alive.set_connection_error_handler(run_loop.QuitClosure());
+  run_loop.Run();
+}
+
 }  // namespace extensions