Begin to enable extension APIs in Extension Service Worker.

extensions/browser/extension_function_dispatcher.cc

Index: extensions/browser/extension_function_dispatcher.cc
diff --git a/extensions/browser/extension_function_dispatcher.cc b/extensions/browser/extension_function_dispatcher.cc
index 5978d8e7ed264d45030a1b08f21fd78cbbe3c59b..8d87b7ea6c344a8ddc35090930439f69d4c69587 100644
--- a/extensions/browser/extension_function_dispatcher.cc
+++ b/extensions/browser/extension_function_dispatcher.cc
@@ -21,6 +21,7 @@
 #include "content/public/browser/browser_thread.h"
 #include "content/public/browser/render_frame_host.h"
 #include "content/public/browser/render_process_host.h"
+#include "content/public/browser/render_process_host_observer.h"
 #include "content/public/browser/render_view_host.h"
 #include "content/public/browser/user_metrics.h"
 #include "content/public/browser/web_contents.h"
@@ -38,6 +39,7 @@
 #include "extensions/common/extension_api.h"
 #include "extensions/common/extension_messages.h"
 #include "extensions/common/extension_set.h"
+#include "extensions/common/extensions_client.h"
 #include "ipc/ipc_message.h"
 #include "ipc/ipc_message_macros.h"
 
@@ -208,6 +210,85 @@ class ExtensionFunctionDispatcher::UIThreadResponseCallbackWrapper
   DISALLOW_COPY_AND_ASSIGN(UIThreadResponseCallbackWrapper);
 };
 
+class ExtensionFunctionDispatcher::UIThreadWorkerResponseCallbackWrapper
+    : public content::RenderProcessHostObserver {
+ public:
+  UIThreadWorkerResponseCallbackWrapper(
+      const base::WeakPtr<ExtensionFunctionDispatcher>& dispatcher,
+      int render_process_id,
+      int worker_thread_id)
+      : dispatcher_(dispatcher),
+        render_process_id_(render_process_id),
+        worker_thread_id_(worker_thread_id),
+        weak_ptr_factory_(this) {
+    content::RenderProcessHost::FromID(render_process_id_)->AddObserver(this);
+    DCHECK(ExtensionsClient::Get()
+               ->ExtensionAPIEnabledInExtensionServiceWorkers());
+  }
+
+  ~UIThreadWorkerResponseCallbackWrapper() override {
+    content::RenderProcessHost* rph =
+        content::RenderProcessHost::FromID(render_process_id_);
+    // We can be deleted from
+    // ExtensionFunctionDispatcher::RemoveWorkerCallbacksForProcess(),
+    // drop our reference from RenderProcessHost in that case.
+    if (rph)
+      rph->RemoveObserver(this);

[Devlin, 2016/04/15 21:06:19]

Hmm... is this right?  If this is the object that first received the
RenderProcessExited (or other) notification, it's already removed itself as an
observer, but, I think, RPH::FromID will still return the RPH (since it's
kinda-sorta still valid).  So we'll try to remove a non-existent observer.

Since CleanUp() results in the synchronous deletion of this object, can we just
move the RemoveObserver call to be solely in the dtor?  And, in that case, maybe
just migrate to a ScopedObserver?

[lazyboy, 2016/04/18 18:42:14]

Done.

+  }
+
+  // content::RenderProcessHostObserver override.
+  void RenderProcessExited(content::RenderProcessHost* rph,
+                           base::TerminationStatus status,
+                           int exit_code) override {
+    rph->RemoveObserver(this);
+    CleanUp();
+  }
+
+  // content::RenderProcessHostObserver override.
+  void RenderProcessHostDestroyed(content::RenderProcessHost* rph) override {
+    rph->RemoveObserver(this);
+    CleanUp();
+  }
+
+  ExtensionFunction::ResponseCallback CreateCallback(int request_id) {
+    return base::Bind(
+        &UIThreadWorkerResponseCallbackWrapper::OnExtensionFunctionCompleted,
+        weak_ptr_factory_.GetWeakPtr(), request_id);
+  }
+
+ private:
+  void CleanUp() {
+    DCHECK_CURRENTLY_ON(BrowserThread::UI);
+    if (dispatcher_)
+      dispatcher_->RemoveWorkerCallbacksForProcess(render_process_id_);
+    // Note: we are deleted here!
+  }
+
+  void OnExtensionFunctionCompleted(int request_id,
+                                    ExtensionFunction::ResponseType type,
+                                    const base::ListValue& results,
+                                    const std::string& error,
+                                    functions::HistogramValue histogram_value) {
+    if (type == ExtensionFunction::BAD_MESSAGE) {
+      // TODO(lazyboy): Kill the offending process.
+      return;
+    }
+    content::RenderProcessHost* sender =
+        content::RenderProcessHost::FromID(render_process_id_);
+    DCHECK(sender);
+    sender->Send(new ExtensionMsg_ResponseWorker(
+        worker_thread_id_, request_id, type == ExtensionFunction::SUCCEEDED,
+        results, error));
+  }
+
+  base::WeakPtr<ExtensionFunctionDispatcher> dispatcher_;
+  const int render_process_id_;
+  const int worker_thread_id_;
+  base::WeakPtrFactory<UIThreadWorkerResponseCallbackWrapper> weak_ptr_factory_;
+
+  DISALLOW_COPY_AND_ASSIGN(UIThreadWorkerResponseCallbackWrapper);
+};
+
 WindowController*
 ExtensionFunctionDispatcher::Delegate::GetExtensionWindowController() const {
   return nullptr;
@@ -307,28 +388,53 @@ ExtensionFunctionDispatcher::~ExtensionFunctionDispatcher() {
 
 void ExtensionFunctionDispatcher::Dispatch(
     const ExtensionHostMsg_Request_Params& params,
-    content::RenderFrameHost* render_frame_host) {
-  UIThreadResponseCallbackWrapperMap::const_iterator
-      iter = ui_thread_response_callback_wrappers_.find(render_frame_host);
-  UIThreadResponseCallbackWrapper* callback_wrapper = nullptr;
-  if (iter == ui_thread_response_callback_wrappers_.end()) {
-    callback_wrapper = new UIThreadResponseCallbackWrapper(AsWeakPtr(),
-                                                           render_frame_host);
-    ui_thread_response_callback_wrappers_[render_frame_host] = callback_wrapper;
+    content::RenderFrameHost* render_frame_host,
+    int render_process_id) {
+  if (render_frame_host) {
+    // Extension API from a non Service Worker context, e.g. extension page,
+    // background page, content script.
+    UIThreadResponseCallbackWrapperMap::const_iterator iter =
+        ui_thread_response_callback_wrappers_.find(render_frame_host);
+    UIThreadResponseCallbackWrapper* callback_wrapper = nullptr;
+    if (iter == ui_thread_response_callback_wrappers_.end()) {
+      callback_wrapper =
+          new UIThreadResponseCallbackWrapper(AsWeakPtr(), render_frame_host);
+      ui_thread_response_callback_wrappers_[render_frame_host] =
+          callback_wrapper;
+    } else {
+      callback_wrapper = iter->second;
+    }
+    DispatchWithCallbackInternal(
+        params, render_frame_host, render_process_id,
+        callback_wrapper->CreateCallback(params.request_id));
   } else {
-    callback_wrapper = iter->second;
+    // Extension API from Service Worker.
+    DCHECK_GE(params.embedded_worker_id, 0);
+    UIThreadWorkerResponseCallbackWrapperMapKey key =
+        std::make_pair(render_process_id, params.embedded_worker_id);
+    UIThreadWorkerResponseCallbackWrapperMap::const_iterator iter =
+        ui_thread_response_callback_wrappers_for_worker_.find(key);
+    UIThreadWorkerResponseCallbackWrapper* callback_wrapper = nullptr;
+    if (iter == ui_thread_response_callback_wrappers_for_worker_.end()) {
+      callback_wrapper = new UIThreadWorkerResponseCallbackWrapper(
+          AsWeakPtr(), render_process_id, params.worker_thread_id);
+      ui_thread_response_callback_wrappers_for_worker_[key] =
+          std::unique_ptr<UIThreadWorkerResponseCallbackWrapper>(
+              callback_wrapper);
+    } else {
+      callback_wrapper = iter->second.get();
+    }
+    DispatchWithCallbackInternal(
+        params, nullptr, render_process_id,
+        callback_wrapper->CreateCallback(params.request_id));
   }
-
-  DispatchWithCallbackInternal(
-      params, render_frame_host,
-      callback_wrapper->CreateCallback(params.request_id));
 }
 
 void ExtensionFunctionDispatcher::DispatchWithCallbackInternal(
     const ExtensionHostMsg_Request_Params& params,
     content::RenderFrameHost* render_frame_host,
+    int render_process_id,
     const ExtensionFunction::ResponseCallback& callback) {
-  DCHECK(render_frame_host);
   // TODO(yzshen): There is some shared logic between this method and
   // DispatchOnIOThread(). It is nice to deduplicate.
   ProcessMap* process_map = ProcessMap::Get(browser_context_);
@@ -343,15 +449,12 @@ void ExtensionFunctionDispatcher::DispatchWithCallbackInternal(
         registry->enabled_extensions().GetHostedAppByURL(params.source_url);
   }
 
-  int process_id = render_frame_host->GetProcess()->GetID();
-  scoped_refptr<ExtensionFunction> function(
-      CreateExtensionFunction(params,
-                              extension,
-                              process_id,
-                              *process_map,
-                              ExtensionAPI::GetSharedInstance(),
-                              browser_context_,
-                              callback));
+  if (render_frame_host)
+    DCHECK_EQ(render_process_id, render_frame_host->GetProcess()->GetID());
+
+  scoped_refptr<ExtensionFunction> function(CreateExtensionFunction(
+      params, extension, render_process_id, *process_map,
+      ExtensionAPI::GetSharedInstance(), browser_context_, callback));
   if (!function.get())
     return;
 
@@ -418,11 +521,29 @@ void ExtensionFunctionDispatcher::DispatchWithCallbackInternal(
   // now, largely for simplicity's sake. This is OK because currently, only
   // the webRequest API uses IOThreadExtensionFunction, and that API is not
   // compatible with lazy background pages.
+  // TODO(lazyboy): API functions from extension Service Worker will incorrectly
+  // change keepalive count below.
   process_manager->IncrementLazyKeepaliveCount(extension);
 }
 
+void ExtensionFunctionDispatcher::RemoveWorkerCallbacksForProcess(
+    int render_process_id) {
+  UIThreadWorkerResponseCallbackWrapperMap& map =
+      ui_thread_response_callback_wrappers_for_worker_;
+  for (UIThreadWorkerResponseCallbackWrapperMap::iterator it = map.begin();
+       it != map.end();) {
+    if (it->first.first == render_process_id) {
+      it = map.erase(it);
+      continue;
+    }
+    ++it;
+  }
+}
+
 void ExtensionFunctionDispatcher::OnExtensionFunctionCompleted(
     const Extension* extension) {
+  // TODO(lazyboy): API functions from extension Service Worker will incorrectly
+  // change keepalive count below.
   if (extension) {
     ProcessManager::Get(browser_context_)
         ->DecrementLazyKeepaliveCount(extension);