1 files changed, 107 insertions, 113 deletions

diff --git a/ipc/glue/ProtocolUtils.cpp b/ipc/glue/ProtocolUtils.cpp
index b29221d476..42bab98320 100644
--- a/ipc/glue/ProtocolUtils.cpp
+++ b/ipc/glue/ProtocolUtils.cpp
@@ -251,12 +251,6 @@ ActorLifecycleProxy::ActorLifecycleProxy(IProtocol* aActor) : mActor(aActor) {
   MOZ_ASSERT(mActor->CanSend(),
              "Cannot create LifecycleProxy for non-connected actor!");
 
-  // Take a reference to our manager's lifecycle proxy to try to hold it &
-  // ensure it doesn't die before us.
-  if (mActor->mManager) {
-    mManager = mActor->mManager->mLifecycleProxy;
-  }
-
   // Record that we've taken our first reference to our actor.
   mActor->ActorAlloc();
 }
@@ -321,6 +315,7 @@ IProtocol::~IProtocol() {
   // Gecko, simply emit a warning and clear the weak backreference from our
   // LifecycleProxy back to us.
   if (mLifecycleProxy) {
+    MOZ_ASSERT(mLinkStatus != LinkStatus::Inactive);
     NS_WARNING(
         nsPrintfCString("Actor destructor for '%s%s' called before IPC "
                         "lifecycle complete!\n"
@@ -329,34 +324,13 @@ IProtocol::~IProtocol() {
             .get());
 
     mLifecycleProxy->mActor = nullptr;
-
-    // If we are somehow being destroyed while active, make sure that the
-    // existing IPC reference has been freed. If the status of the actor is
-    // `Destroyed`, the reference has already been freed, and we shouldn't free
-    // it a second time.
-    MOZ_ASSERT(mLinkStatus != LinkStatus::Inactive);
-    if (mLinkStatus != LinkStatus::Destroyed) {
-      NS_IF_RELEASE(mLifecycleProxy);
-    }
     mLifecycleProxy = nullptr;
   }
 }
 
 // The following methods either directly forward to the toplevel protocol, or
 // almost directly do.
-int32_t IProtocol::Register(IProtocol* aRouted) {
-  return mToplevel->Register(aRouted);
-}
-int32_t IProtocol::RegisterID(IProtocol* aRouted, int32_t aId) {
-  return mToplevel->RegisterID(aRouted, aId);
-}
 IProtocol* IProtocol::Lookup(int32_t aId) { return mToplevel->Lookup(aId); }
-void IProtocol::Unregister(int32_t aId) {
-  if (aId == mId) {
-    mId = kFreedActorId;
-  }
-  return mToplevel->Unregister(aId);
-}
 
 Shmem::SharedMemory* IProtocol::CreateSharedMemory(size_t aSize, bool aUnsafe,
                                                    int32_t* aId) {
@@ -383,11 +357,6 @@ nsISerialEventTarget* IProtocol::GetActorEventTarget() {
   return GetIPCChannel()->GetWorkerEventTarget();
 }
 
-void IProtocol::SetId(int32_t aId) {
-  MOZ_ASSERT(mId == aId || mLinkStatus == LinkStatus::Inactive);
-  mId = aId;
-}
-
 Maybe<IProtocol*> IProtocol::ReadActor(IPC::MessageReader* aReader,
                                        bool aNullable,
                                        const char* aActorDescription,
@@ -487,26 +456,60 @@ bool IProtocol::DeallocShmem(Shmem& aMem) {
   return ok;
 }
 
-void IProtocol::SetManager(IProtocol* aManager) {
+void IProtocol::SetManager(IRefCountedProtocol* aManager) {
   MOZ_RELEASE_ASSERT(!mManager || mManager == aManager);
   mManager = aManager;
   mToplevel = aManager->mToplevel;
 }
 
-void IProtocol::SetManagerAndRegister(IProtocol* aManager) {
-  // Set the manager prior to registering so registering properly inherits
-  // the manager's event target.
-  SetManager(aManager);
+bool IProtocol::SetManagerAndRegister(IRefCountedProtocol* aManager,
+                                      int32_t aId) {
+  MOZ_RELEASE_ASSERT(mLinkStatus == LinkStatus::Inactive,
+                     "Actor must be inactive to SetManagerAndRegister");
 
-  aManager->Register(this);
-}
+  // Set to `false` if the actor is to be torn down after registration.
+  bool success = true;
 
-void IProtocol::SetManagerAndRegister(IProtocol* aManager, int32_t aId) {
   // Set the manager prior to registering so registering properly inherits
   // the manager's event target.
   SetManager(aManager);
 
-  aManager->RegisterID(this, aId);
+  mId = aId == kNullActorId ? mToplevel->NextId() : aId;
+  while (mToplevel->mActorMap.Contains(mId)) {
+    // The ID already existing is an error case, but we want to proceed with
+    // registration so that we can tear down the actor cleanly - generate a new
+    // ID for that case.
+    NS_WARNING("Actor already exists with the selected ID!");
+    mId = mToplevel->NextId();
+    success = false;
+  }
+
+  RefPtr<ActorLifecycleProxy> proxy = ActorConnected();
+  mToplevel->mActorMap.InsertOrUpdate(mId, proxy);
+  MOZ_ASSERT(proxy->Get() == this);
+
+  // If our manager is already dying, mark ourselves as doomed as well.
+  if (aManager && aManager->mLinkStatus != LinkStatus::Connected) {
+    mLinkStatus = LinkStatus::Doomed;
+    if (aManager->mLinkStatus != LinkStatus::Doomed) {
+      // Our manager is already fully dead, make sure we call
+      // `ActorDisconnected`.
+      success = false;
+    }
+  }
+
+  // If setting the manager failed, call `ActorDisconnected` and return false.
+  if (!success) {
+    ActorDisconnected(FailedConstructor);
+    MOZ_ASSERT(mLinkStatus == LinkStatus::Destroyed);
+    return false;
+  }
+  return true;
+}
+
+void IProtocol::UnlinkManager() {
+  mToplevel = nullptr;
+  mManager = nullptr;
 }
 
 bool IProtocol::ChannelSend(UniquePtr<IPC::Message> aMsg) {
@@ -540,9 +543,9 @@ void IProtocol::WarnMessageDiscarded(IPC::Message* aMsg) {
 }
 #endif
 
-void IProtocol::ActorConnected() {
+already_AddRefed<ActorLifecycleProxy> IProtocol::ActorConnected() {
   if (mLinkStatus != LinkStatus::Inactive) {
-    return;
+    return nullptr;
   }
 
 #ifdef FUZZING_SNAPSHOT
@@ -552,73 +555,81 @@ void IProtocol::ActorConnected() {
   mLinkStatus = LinkStatus::Connected;
 
   MOZ_ASSERT(!mLifecycleProxy, "double-connecting live actor");
-  mLifecycleProxy = new ActorLifecycleProxy(this);
-  NS_ADDREF(mLifecycleProxy);  // Reference freed in DestroySubtree();
+  RefPtr<ActorLifecycleProxy> proxy = new ActorLifecycleProxy(this);
+  mLifecycleProxy = proxy;
+  return proxy.forget();
 }
 
-void IProtocol::DoomSubtree() {
-  MOZ_ASSERT(CanSend(), "dooming non-connected actor");
-  MOZ_ASSERT(mLifecycleProxy, "dooming zombie actor");
-
-  nsTArray<RefPtr<ActorLifecycleProxy>> managed;
-  AllManagedActors(managed);
-  for (ActorLifecycleProxy* proxy : managed) {
-    // Guard against actor being disconnected or destroyed during previous Doom
-    IProtocol* actor = proxy->Get();
-    if (actor && actor->CanSend()) {
-      actor->DoomSubtree();
-    }
+void IProtocol::ActorDisconnected(ActorDestroyReason aWhy) {
+  MOZ_ASSERT(mLifecycleProxy, "destroying zombie actor");
+  // If the actor has already been marked as `Destroyed`, there's nothing to do.
+  if (mLinkStatus != LinkStatus::Connected &&
+      mLinkStatus != LinkStatus::Doomed) {
+    return;
   }
 
-  // ActorDoom is called immediately before changing state, this allows messages
-  // to be sent during ActorDoom immediately before the channel is closed and
-  // sending messages is disabled.
-  ActorDoom();
-  mLinkStatus = LinkStatus::Doomed;
-}
+  // Mark the entire subtree as doomed so that no further messages can be
+  // sent/recieved, and newly created managed actors are immediately marked as
+  // doomed on creation.
+  DoomSubtree();
 
-void IProtocol::DestroySubtree(ActorDestroyReason aWhy) {
-  MOZ_ASSERT(CanRecv(), "destroying non-connected actor");
-  MOZ_ASSERT(mLifecycleProxy, "destroying zombie actor");
+  // Perform the steps to fully destroy an actor after it has been unregistered
+  // from its manager.
+  auto doActorDestroy = [toplevel = mToplevel, ipcChannel = GetIPCChannel()](
+                            IProtocol* actor, ActorDestroyReason why) {
+    MOZ_ASSERT(actor->mLinkStatus == LinkStatus::Doomed,
+               "Actor must be doomed when calling doActorDestroy");
+    MOZ_ASSERT(actor->AllManagedActorsCount() == 0,
+               "All managed actors must have been destroyed first");
+
+    // Mark the actor as Destroyed, ensuring we can't re-enter `ActorDestroy`,
+    // even if an callback spins a nested event loop.
+    actor->mLinkStatus = LinkStatus::Destroyed;
 
 #ifdef FUZZING_SNAPSHOT
-  fuzzing::IPCFuzzController::instance().OnActorDestroyed(this);
+    fuzzing::IPCFuzzController::instance().OnActorDestroyed(actor);
 #endif
 
-  int32_t id = Id();
+    int32_t id = actor->mId;
+    if (IProtocol* manager = actor->Manager()) {
+      actor->mId = kFreedActorId;
+      auto entry = toplevel->mActorMap.Lookup(id);
+      MOZ_DIAGNOSTIC_ASSERT(entry && *entry == actor->GetLifecycleProxy(),
+                            "ID must be present and reference this actor");
+      entry.Remove();
+      manager->RemoveManagee(actor->GetProtocolId(), actor);
+    }
+
+    ipcChannel->RejectPendingResponsesForActor(id);
+    actor->ActorDestroy(why);
+  };
 
-  // If we're a managed actor, unregister from our manager
-  if (Manager()) {
-    Unregister(id);
-  }
+  // Hold all ActorLifecycleProxy instances for managed actors until we return.
+  nsTArray<RefPtr<ActorLifecycleProxy>> proxyHolder;
+  proxyHolder.AppendElement(GetLifecycleProxy());
 
-  // Destroy subtree
+  // Invoke `ActorDestroy` for all managed actors in the subtree. These are
+  // handled one at a time, so that new actors which are potentially registered
+  // during `ActorDestroy` callbacks are not missed.
   ActorDestroyReason subtreeWhy = aWhy;
   if (aWhy == Deletion || aWhy == FailedConstructor) {
     subtreeWhy = AncestorDeletion;
   }
-
-  nsTArray<RefPtr<ActorLifecycleProxy>> managed;
-  AllManagedActors(managed);
-  for (ActorLifecycleProxy* proxy : managed) {
-    // Guard against actor being disconnected or destroyed during previous
-    // Destroy
-    IProtocol* actor = proxy->Get();
-    if (actor && actor->CanRecv()) {
-      actor->DestroySubtree(subtreeWhy);
+  while (IProtocol* actor = PeekManagedActor()) {
+    // If the selected actor manages other actors, destroy those first.
+    while (IProtocol* inner = actor->PeekManagedActor()) {
+      actor = inner;
     }
-  }
 
-  // Ensure that we don't send any messages while we're calling `ActorDestroy`
-  // by setting our state to `Doomed`.
-  mLinkStatus = LinkStatus::Doomed;
+    proxyHolder.AppendElement(actor->GetLifecycleProxy());
+    doActorDestroy(actor, subtreeWhy);
+  }
 
-  // The actor is being destroyed, reject any pending responses, invoke
-  // `ActorDestroy` to destroy it, and then clear our status to
-  // `LinkStatus::Destroyed`.
-  GetIPCChannel()->RejectPendingResponsesForActor(id);
-  ActorDestroy(aWhy);
-  mLinkStatus = LinkStatus::Destroyed;
+  // Destroy ourselves if we were not not otherwise destroyed while destroying
+  // managed actors.
+  if (mLinkStatus == LinkStatus::Doomed) {
+    doActorDestroy(this, aWhy);
+  }
 }
 
 IToplevelProtocol::IToplevelProtocol(const char* aName, ProtocolId aProtoId,
@@ -689,28 +700,11 @@ int32_t IToplevelProtocol::NextId() {
   return (++mLastLocalId << 2) | tag;
 }
 
-int32_t IToplevelProtocol::Register(IProtocol* aRouted) {
-  if (aRouted->Id() != kNullActorId && aRouted->Id() != kFreedActorId) {
-    // If there's already an ID, just return that.
-    return aRouted->Id();
+IProtocol* IToplevelProtocol::Lookup(int32_t aId) {
+  if (auto entry = mActorMap.Lookup(aId)) {
+    return entry.Data()->Get();
   }
-  return RegisterID(aRouted, NextId());
-}
-
-int32_t IToplevelProtocol::RegisterID(IProtocol* aRouted, int32_t aId) {
-  aRouted->SetId(aId);
-  aRouted->ActorConnected();
-  MOZ_ASSERT(!mActorMap.Contains(aId), "Don't insert with an existing ID");
-  mActorMap.InsertOrUpdate(aId, aRouted);
-  return aId;
-}
-
-IProtocol* IToplevelProtocol::Lookup(int32_t aId) { return mActorMap.Get(aId); }
-
-void IToplevelProtocol::Unregister(int32_t aId) {
-  MOZ_ASSERT(mActorMap.Contains(aId),
-             "Attempting to remove an ID not in the actor map");
-  mActorMap.Remove(aId);
+  return nullptr;
 }
 
 Shmem::SharedMemory* IToplevelProtocol::CreateSharedMemory(size_t aSize,