Bug 1535913 - Part 2: Hold a keepalive for extension BCGs in WebExtensionPolicy, r=kmag

Differential Revision: https://phabricator.services.mozilla.com/D114889

docshell/base/BrowsingContextGroup.cpp

@@ -272,6 +272,11 @@ void BrowsingContextGroup::RemoveKeepAlive() {
   MaybeDestroy();
 }
 
+auto BrowsingContextGroup::MakeKeepAlivePtr() -> KeepAlivePtr {
+  AddKeepAlive();
+  return KeepAlivePtr{do_AddRef(this).take()};
+}
+
 void BrowsingContextGroup::MaybeDestroy() {
   // Once there are no synced contexts referencing a `BrowsingContextGroup`, we
   // can clear subscribers and destroy this group. We only do this in the parent

docshell/base/BrowsingContextGroup.h

@@ -77,6 +77,19 @@ class BrowsingContextGroup final : public nsWrapperCache {
   void AddKeepAlive();
   void RemoveKeepAlive();
 
+  // A `KeepAlivePtr` will hold both a strong reference to the
+  // `BrowsingContextGroup` and holds a `KeepAlive`. When the pointer is
+  // dropped, it will release both the strong reference and the keepalive.
+  struct KeepAliveDeleter {
+    void operator()(BrowsingContextGroup* aPtr) {
+      if (RefPtr<BrowsingContextGroup> ptr = already_AddRefed(aPtr)) {
+        ptr->RemoveKeepAlive();
+      }
+    }
+  };
+  using KeepAlivePtr = UniquePtr<BrowsingContextGroup, KeepAliveDeleter>;
+  KeepAlivePtr MakeKeepAlivePtr();
+
   // Call when we want to check if we should suspend or resume all top level
   // contexts.
   void UpdateToplevelsSuspendedIfNeeded();
@@ -232,4 +245,16 @@ class BrowsingContextGroup final : public nsWrapperCache {
 }  // namespace dom
 }  // namespace mozilla
 
+inline void ImplCycleCollectionUnlink(
+    mozilla::dom::BrowsingContextGroup::KeepAlivePtr& aField) {
+  aField = nullptr;
+}
+
+inline void ImplCycleCollectionTraverse(
+    nsCycleCollectionTraversalCallback& aCallback,
+    mozilla::dom::BrowsingContextGroup::KeepAlivePtr& aField, const char* aName,
+    uint32_t aFlags = 0) {
+  CycleCollectionNoteChild(aCallback, aField.get(), aName, aFlags);
+}
+
 #endif  // !defined(mozilla_dom_BrowsingContextGroup_h)

toolkit/components/extensions/WebExtensionPolicy.cpp

@@ -309,8 +309,9 @@ bool WebExtensionPolicy::Enable() {
   }
 
   if (XRE_IsParentProcess()) {
-    // Reserve a BrowsingContextGroup ID for use by this WebExtensionPolicy.
-    mBrowsingContextGroupId = nsContentUtils::GenerateBrowsingContextId();
+    // Reserve a BrowsingContextGroup for use by this WebExtensionPolicy.
+    RefPtr<BrowsingContextGroup> group = BrowsingContextGroup::Create();
+    mBrowsingContextGroup = group->MakeKeepAlivePtr();
   }
 
   Unused << Proto()->SetSubstitution(MozExtensionHostname(), mBaseURI);
@@ -327,6 +328,12 @@ bool WebExtensionPolicy::Disable() {
     return false;
   }
 
+  if (XRE_IsParentProcess()) {
+    // Clear our BrowsingContextGroup reference. A new instance will be created
+    // when the extension is next activated.
+    mBrowsingContextGroup = nullptr;
+  }
+
   Unused << Proto()->SetSubstitution(MozExtensionHostname(), nullptr);
 
   mActive = false;
@@ -570,7 +577,7 @@ void WebExtensionPolicy::GetReadyPromise(
 
 uint64_t WebExtensionPolicy::GetBrowsingContextGroupId() const {
   MOZ_ASSERT(XRE_IsParentProcess() && mActive);
-  return mBrowsingContextGroupId;
+  return mBrowsingContextGroup ? mBrowsingContextGroup->Id() : 0;
 }
 
 uint64_t WebExtensionPolicy::GetBrowsingContextGroupId(ErrorResult& aRv) {
@@ -583,11 +590,9 @@ uint64_t WebExtensionPolicy::GetBrowsingContextGroupId(ErrorResult& aRv) {
   return 0;
 }
 
-NS_IMPL_CYCLE_COLLECTION_WRAPPERCACHE_WEAK_PTR(WebExtensionPolicy, mParent,
-                                               mLocalizeCallback,
-                                               mHostPermissions,
-                                               mWebAccessibleResources,
-                                               mContentScripts)
+NS_IMPL_CYCLE_COLLECTION_WRAPPERCACHE_WEAK_PTR(
+    WebExtensionPolicy, mParent, mBrowsingContextGroup, mLocalizeCallback,
+    mHostPermissions, mWebAccessibleResources, mContentScripts)
 
 NS_INTERFACE_MAP_BEGIN_CYCLE_COLLECTION(WebExtensionPolicy)
   NS_WRAPPERCACHE_INTERFACE_MAP_ENTRY

toolkit/components/extensions/WebExtensionPolicy.h

@@ -7,6 +7,7 @@
 #define mozilla_extensions_WebExtensionPolicy_h
 
 #include "mozilla/dom/BindingDeclarations.h"
+#include "mozilla/dom/BrowsingContextGroup.h"
 #include "mozilla/dom/Nullable.h"
 #include "mozilla/dom/WebExtensionPolicyBinding.h"
 #include "mozilla/dom/WindowProxyHolder.h"
@@ -239,7 +240,7 @@ class WebExtensionPolicy final : public nsISupports,
   nsString mExtensionPageCSP;
   nsString mBaseCSP;
 
-  uint64_t mBrowsingContextGroupId = 0;
+  dom::BrowsingContextGroup::KeepAlivePtr mBrowsingContextGroup;
 
   bool mActive = false;
   bool mAllowPrivateBrowsingByDefault = true;