From c34f63634cd831bd9857060fadf869e79f868850 Mon Sep 17 00:00:00 2001
From: Kershaw Chang <kershaw@mozilla.com>
Date: Mon, 12 Apr 2021 11:50:35 +0000
Subject: [PATCH] Bug 1660307 - Get rid off nsSupportsWeakReference for
 WebSocketImpl, r=baku

Differential Revision: https://phabricator.services.mozilla.com/D111059
---
 dom/websocket/WebSocket.cpp | 80 ++++++++++++++++++++++++++++++++++---
 1 file changed, 74 insertions(+), 6 deletions(-)

diff --git a/dom/websocket/WebSocket.cpp b/dom/websocket/WebSocket.cpp
index ff34e9c8819f..4e9838dcb2b4 100644
--- a/dom/websocket/WebSocket.cpp
+++ b/dom/websocket/WebSocket.cpp
@@ -29,6 +29,7 @@
 #include "mozilla/dom/WorkerRef.h"
 #include "mozilla/dom/WorkerRunnable.h"
 #include "mozilla/dom/WorkerScope.h"
+#include "mozilla/ScopeExit.h"
 #include "mozilla/StaticPrefs_dom.h"
 #include "mozilla/LoadInfo.h"
 #include "nsGlobalWindow.h"
@@ -80,10 +81,38 @@ using namespace mozilla::net;
 namespace mozilla {
 namespace dom {
 
+class WebSocketImpl;
+
+// This class is responsible for proxying nsIObserver and nsIWebSocketImpl
+// interfaces to WebSocketImpl. WebSocketImplProxy should be only accessed on
+// main thread, so we can let it support weak reference.
+class WebSocketImplProxy final : public nsIObserver,
+                                 public nsSupportsWeakReference,
+                                 public nsIWebSocketImpl {
+ public:
+  NS_DECL_ISUPPORTS
+  NS_DECL_NSIOBSERVER
+  NS_DECL_NSIWEBSOCKETIMPL
+
+  explicit WebSocketImplProxy(WebSocketImpl* aOwner) : mOwner(aOwner) {
+    MOZ_ASSERT(NS_IsMainThread());
+  }
+
+  void Disconnect() {
+    MOZ_ASSERT(NS_IsMainThread());
+
+    mOwner = nullptr;
+  }
+
+ private:
+  ~WebSocketImplProxy() = default;
+
+  RefPtr<WebSocketImpl> mOwner;
+};
+
 class WebSocketImpl final : public nsIInterfaceRequestor,
                             public nsIWebSocketListener,
                             public nsIObserver,
-                            public nsSupportsWeakReference,
                             public nsIRequest,
                             public nsIEventTarget,
                             public nsIWebSocketImpl {
@@ -116,6 +145,8 @@ class WebSocketImpl final : public nsIInterfaceRequestor,
         mWorkerShuttingDown(false) {
     if (!NS_IsMainThread()) {
       mIsMainThread = false;
+    } else {
+      mImplProxy = new WebSocketImplProxy(this);
     }
   }
 
@@ -230,6 +261,8 @@ class WebSocketImpl final : public nsIInterfaceRequestor,
   // For dispatching runnables to main thread.
   nsCOMPtr<nsIEventTarget> mMainThreadEventTarget;
 
+  RefPtr<WebSocketImplProxy> mImplProxy;
+
  private:
   ~WebSocketImpl() {
     MOZ_DIAGNOSTIC_ASSERT(NS_IsMainThread() == mIsMainThread ||
@@ -242,9 +275,30 @@ class WebSocketImpl final : public nsIInterfaceRequestor,
   }
 };
 
+NS_IMPL_ISUPPORTS(WebSocketImplProxy, nsIObserver, nsISupportsWeakReference,
+                  nsIWebSocketImpl)
+
+NS_IMETHODIMP
+WebSocketImplProxy::Observe(nsISupports* aSubject, const char* aTopic,
+                            const char16_t* aData) {
+  if (!mOwner) {
+    return NS_OK;
+  }
+
+  return mOwner->Observe(aSubject, aTopic, aData);
+}
+
+NS_IMETHODIMP
+WebSocketImplProxy::SendMessage(const nsAString& aMessage) {
+  if (!mOwner) {
+    return NS_OK;
+  }
+
+  return mOwner->SendMessage(aMessage);
+}
+
 NS_IMPL_ISUPPORTS(WebSocketImpl, nsIInterfaceRequestor, nsIWebSocketListener,
-                  nsIObserver, nsISupportsWeakReference, nsIRequest,
-                  nsIEventTarget, nsIWebSocketImpl)
+                  nsIObserver, nsIRequest, nsIEventTarget, nsIWebSocketImpl)
 
 class CallDispatchConnectionCloseEvents final : public DiscardableRunnable {
  public:
@@ -617,6 +671,11 @@ void WebSocketImpl::DisconnectInternal() {
       os->RemoveObserver(this, DOM_WINDOW_FROZEN_TOPIC);
     }
   }
+
+  if (mImplProxy) {
+    mImplProxy->Disconnect();
+    mImplProxy = nullptr;
+  }
 }
 
 //-----------------------------------------------------------------------------
@@ -1248,6 +1307,12 @@ already_AddRefed<WebSocket> WebSocket::ConstructorCommon(
   RefPtr<WebSocket> webSocket = new WebSocket(global);
   RefPtr<WebSocketImpl> webSocketImpl = webSocket->mImpl;
 
+  auto clearGuard = MakeScopeExit([&] {
+    if (NS_IsMainThread()) {
+      webSocketImpl->mImplProxy->Disconnect();
+    }
+  });
+
   bool connectionFailed = true;
 
   if (NS_IsMainThread()) {
@@ -1492,11 +1557,12 @@ nsresult WebSocketImpl::Init(JSContext* aCx, nsIPrincipal* aLoadingPrincipal,
     if (NS_WARN_IF(!os)) {
       return NS_ERROR_FAILURE;
     }
+    MOZ_ASSERT(mImplProxy);
 
-    rv = os->AddObserver(this, DOM_WINDOW_DESTROYED_TOPIC, true);
+    rv = os->AddObserver(mImplProxy, DOM_WINDOW_DESTROYED_TOPIC, true);
     NS_ENSURE_SUCCESS(rv, rv);
 
-    rv = os->AddObserver(this, DOM_WINDOW_FROZEN_TOPIC, true);
+    rv = os->AddObserver(mImplProxy, DOM_WINDOW_FROZEN_TOPIC, true);
     NS_ENSURE_SUCCESS(rv, rv);
   }
 
@@ -1781,7 +1847,9 @@ nsresult WebSocketImpl::InitializeConnection(
   mChannel = wsChannel;
 
   if (mIsMainThread) {
-    mService->AssociateWebSocketImplWithSerialID(this, mChannel->Serial());
+    MOZ_ASSERT(mImplProxy);
+    mService->AssociateWebSocketImplWithSerialID(mImplProxy,
+                                                 mChannel->Serial());
   }
 
   if (mIsMainThread && doc) {