From 5dbaaf5ed36c31eac38276e42da334ae2a8734a3 Mon Sep 17 00:00:00 2001
From: Ben Kelly <ben@wanderview.com>
Date: Fri, 22 Dec 2017 21:09:18 -0500
Subject: [PATCH] Bug 1425975 P2 Add ServiceWorkerManager mControlledClients to
 track controlled ClientHandle references. r=asuth

---
 dom/workers/ServiceWorkerManager.cpp | 65 ++++++++++++++++++++++++----
 dom/workers/ServiceWorkerManager.h   | 23 ++++++++++
 2 files changed, 79 insertions(+), 9 deletions(-)

diff --git a/dom/workers/ServiceWorkerManager.cpp b/dom/workers/ServiceWorkerManager.cpp
index 565a92fb0311..aa05ce397f29 100644
--- a/dom/workers/ServiceWorkerManager.cpp
+++ b/dom/workers/ServiceWorkerManager.cpp
@@ -313,6 +313,41 @@ ServiceWorkerManager::Init(ServiceWorkerRegistrar* aRegistrar)
   mActor = static_cast<ServiceWorkerManagerChild*>(actor);
 }
 
+void
+ServiceWorkerManager::StartControllingClient(const ClientInfo& aClientInfo,
+                                             ServiceWorkerRegistrationInfo* aRegistrationInfo)
+{
+  RefPtr<ClientHandle> clientHandle =
+    ClientManager::CreateHandle(aClientInfo,
+                                SystemGroup::EventTargetFor(TaskCategory::Other));
+
+  auto entry = mControlledClients.LookupForAdd(aClientInfo.Id());
+  if (entry) {
+    entry.Data()->mRegistrationInfo = aRegistrationInfo;
+  } else {
+    entry.OrInsert([&] {
+      return new ControlledClientData(clientHandle, aRegistrationInfo);
+    });
+
+    RefPtr<ServiceWorkerManager> self(this);
+    clientHandle->OnDetach()->Then(
+      SystemGroup::EventTargetFor(TaskCategory::Other), __func__,
+      [self = Move(self), aClientInfo] {
+        self->StopControllingClient(aClientInfo);
+      });
+  }
+}
+
+void
+ServiceWorkerManager::StopControllingClient(const ClientInfo& aClientInfo)
+{
+  auto entry = mControlledClients.Lookup(aClientInfo.Id());
+  if (!entry) {
+    return;
+  }
+  entry.Remove();
+}
+
 void
 ServiceWorkerManager::MaybeStartShutdown()
 {
@@ -2362,23 +2397,35 @@ ServiceWorkerManager::StartControllingADocument(ServiceWorkerRegistrationInfo* a
   MOZ_DIAGNOSTIC_ASSERT(storageAllowed == nsContentUtils::StorageAccess::eAllow);
 #endif // MOZ_DIAGNOSTIC_ASSERT_ENABLED
 
-  RefPtr<GenericPromise> ref = GenericPromise::CreateAndResolve(true, __func__);
+  RefPtr<GenericPromise> ref;
+
+  ServiceWorkerInfo* activeWorker = aRegistration->GetActive();
+  if (NS_WARN_IF(!activeWorker)) {
+    ref = GenericPromise::CreateAndReject(NS_ERROR_DOM_INVALID_STATE_ERR,
+                                          __func__);
+    return ref.forget();
+  }
+
+  Maybe<ClientInfo> clientInfo = aDoc->GetClientInfo();
+  if (NS_WARN_IF(clientInfo.isNothing())) {
+    ref = GenericPromise::CreateAndReject(NS_ERROR_DOM_INVALID_STATE_ERR,
+                                          __func__);
+    return ref.forget();
+  }
 
   aRegistration->StartControllingADocument();
   mControlledDocuments.Put(aDoc, aRegistration);
 
+  StartControllingClient(clientInfo.ref(), aRegistration);
+
   // Mark the document's ClientSource as controlled using the ClientHandle
   // interface.  While we could get at the ClientSource directly from the
   // document here, our goal is to move ServiceWorkerManager to a separate
   // process.  Using the ClientHandle supports this remote operation.
-  ServiceWorkerInfo* activeWorker = aRegistration->GetActive();
-  Maybe<ClientInfo> clientInfo = aDoc->GetClientInfo();
-  if (activeWorker && clientInfo.isSome()) {
-    RefPtr<ClientHandle> clientHandle =
-      ClientManager::CreateHandle(clientInfo.ref(),
-                                  SystemGroup::EventTargetFor(TaskCategory::Other));
-    ref = Move(clientHandle->Control(activeWorker->Descriptor()));
-  }
+  RefPtr<ClientHandle> clientHandle =
+    ClientManager::CreateHandle(clientInfo.ref(),
+                                SystemGroup::EventTargetFor(TaskCategory::Other));
+  ref = Move(clientHandle->Control(activeWorker->Descriptor()));
 
   Telemetry::Accumulate(Telemetry::SERVICE_WORKER_CONTROLLED_DOCUMENTS, 1);
   return Move(ref);
diff --git a/dom/workers/ServiceWorkerManager.h b/dom/workers/ServiceWorkerManager.h
index a8d5fcbbd9e0..a0eb0b1b29a1 100644
--- a/dom/workers/ServiceWorkerManager.h
+++ b/dom/workers/ServiceWorkerManager.h
@@ -21,6 +21,7 @@
 #include "mozilla/UniquePtr.h"
 #include "mozilla/WeakPtr.h"
 #include "mozilla/dom/BindingUtils.h"
+#include "mozilla/dom/ClientHandle.h"
 #include "mozilla/dom/Promise.h"
 #include "mozilla/dom/ServiceWorkerCommon.h"
 #include "mozilla/dom/ServiceWorkerRegistrar.h"
@@ -105,6 +106,21 @@ public:
 
   nsRefPtrHashtable<nsISupportsHashKey, ServiceWorkerRegistrationInfo> mControlledDocuments;
 
+  struct ControlledClientData
+  {
+    RefPtr<ClientHandle> mClientHandle;
+    RefPtr<ServiceWorkerRegistrationInfo> mRegistrationInfo;
+
+    ControlledClientData(ClientHandle* aClientHandle,
+                         ServiceWorkerRegistrationInfo* aRegistrationInfo)
+      : mClientHandle(aClientHandle)
+      , mRegistrationInfo(aRegistrationInfo)
+    {
+    }
+  };
+
+  nsClassHashtable<nsIDHashKey, ControlledClientData> mControlledClients;
+
   // Track all documents that have attempted to register a service worker for a
   // given scope.
   typedef nsTArray<nsCOMPtr<nsIWeakReference>> WeakDocumentList;
@@ -328,6 +344,13 @@ private:
   void
   Init(ServiceWorkerRegistrar* aRegistrar);
 
+  void
+  StartControllingClient(const ClientInfo& aClientInfo,
+                         ServiceWorkerRegistrationInfo* aRegistrationInfo);
+
+  void
+  StopControllingClient(const ClientInfo& aClientInfo);
+
   void
   MaybeStartShutdown();