Bug 1660307 - Get rid off nsSupportsWeakReference for WebSocketImpl, r=baku

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

dom/websocket/WebSocket.cpp

@@ -80,10 +80,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 {
@@ -230,6 +258,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 +272,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:
@@ -524,6 +575,11 @@ void WebSocketImpl::FailConnection(uint16_t aReasonCode,
   ConsoleError();
   mFailed = true;
   CloseConnection(aReasonCode, aReasonString);
+
+  if (NS_IsMainThread() && mImplProxy) {
+    mImplProxy->Disconnect();
+    mImplProxy = nullptr;
+  }
 }
 
 namespace {
@@ -613,10 +669,15 @@ void WebSocketImpl::DisconnectInternal() {
   if (!mWorkerRef) {
     nsCOMPtr<nsIObserverService> os = mozilla::services::GetObserverService();
     if (os) {
-      os->RemoveObserver(this, DOM_WINDOW_DESTROYED_TOPIC);
-      os->RemoveObserver(this, DOM_WINDOW_FROZEN_TOPIC);
+      os->RemoveObserver(mImplProxy, DOM_WINDOW_DESTROYED_TOPIC);
+      os->RemoveObserver(mImplProxy, DOM_WINDOW_FROZEN_TOPIC);
     }
   }
+
+  if (mImplProxy) {
+    mImplProxy->Disconnect();
+    mImplProxy = nullptr;
+  }
 }
 
 //-----------------------------------------------------------------------------
@@ -1487,16 +1548,18 @@ nsresult WebSocketImpl::Init(JSContext* aCx, nsIPrincipal* aLoadingPrincipal,
 
   // Shut down websocket if window is frozen or destroyed (only needed for
   // "ghost" websockets--see bug 696085)
+  RefPtr<WebSocketImplProxy> proxy;
   if (mIsMainThread) {
     nsCOMPtr<nsIObserverService> os = mozilla::services::GetObserverService();
     if (NS_WARN_IF(!os)) {
       return NS_ERROR_FAILURE;
     }
 
-    rv = os->AddObserver(this, DOM_WINDOW_DESTROYED_TOPIC, true);
+    proxy = new WebSocketImplProxy(this);
+    rv = os->AddObserver(proxy, DOM_WINDOW_DESTROYED_TOPIC, true);
     NS_ENSURE_SUCCESS(rv, rv);
 
-    rv = os->AddObserver(this, DOM_WINDOW_FROZEN_TOPIC, true);
+    rv = os->AddObserver(proxy, DOM_WINDOW_FROZEN_TOPIC, true);
     NS_ENSURE_SUCCESS(rv, rv);
   }
 
@@ -1653,6 +1716,9 @@ nsresult WebSocketImpl::Init(JSContext* aCx, nsIPrincipal* aLoadingPrincipal,
     AppendUTF16toUTF8(aProtocolArray[index], mRequestedProtocolList);
   }
 
+  if (mIsMainThread) {
+    mImplProxy = std::move(proxy);
+  }
   return NS_OK;
 }
 
@@ -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) {