Bug 1111971 - A better life-time management of aListener and aContext in WebSocketChannel. r=smaug

CLOSED TREE

dom/base/WebSocket.cpp

@@ -646,13 +646,14 @@ WebSocketImpl::DoOnMessageAvailable(const nsACString& aMsg, bool isBinary)
     if (NS_FAILED(rv)) {
       NS_WARNING("Failed to dispatch the message event");
     }
-  } else {
+
-    // CLOSING should be the only other state where it's possible to get msgs
+    return NS_OK;
-    // from channel: Spec says to drop them.
-    MOZ_ASSERT(readyState == WebSocket::CLOSING,
-               "Received message while CONNECTING or CLOSED");
   }
 
+  // CLOSING should be the only other state where it's possible to get msgs
+  // from channel: Spec says to drop them.
+  MOZ_ASSERT(readyState == WebSocket::CLOSING,
+             "Received message while CONNECTING or CLOSED");
   return NS_OK;
 }
 
@@ -718,14 +719,17 @@ WebSocketImpl::OnStart(nsISupports* aContext)
 
   mWebSocket->SetReadyState(WebSocket::OPEN);
 
+  // Let's keep the object alive because the webSocket can be CCed in the
+  // onopen callback.
+  nsRefPtr<WebSocket> webSocket = mWebSocket;
+
   // Call 'onopen'
-  rv = mWebSocket->CreateAndDispatchSimpleEvent(NS_LITERAL_STRING("open"));
+  rv = webSocket->CreateAndDispatchSimpleEvent(NS_LITERAL_STRING("open"));
   if (NS_FAILED(rv)) {
     NS_WARNING("Failed to dispatch the open event");
   }
 
-  mWebSocket->UpdateMustKeepAlive();
+  webSocket->UpdateMustKeepAlive();
-
   return NS_OK;
 }
 
@@ -1596,23 +1600,27 @@ WebSocketImpl::DispatchConnectionCloseEvents()
 
   mWebSocket->SetReadyState(WebSocket::CLOSED);
 
+  // Let's keep the object alive because the webSocket can be CCed in the
+  // onerror or in the onclose callback.
+  nsRefPtr<WebSocket> webSocket = mWebSocket;
+
   // Call 'onerror' if needed
   if (mFailed) {
     nsresult rv =
-      mWebSocket->CreateAndDispatchSimpleEvent(NS_LITERAL_STRING("error"));
+      webSocket->CreateAndDispatchSimpleEvent(NS_LITERAL_STRING("error"));
     if (NS_FAILED(rv)) {
       NS_WARNING("Failed to dispatch the error event");
     }
   }
 
-  nsresult rv = mWebSocket->CreateAndDispatchCloseEvent(mCloseEventWasClean,
+  nsresult rv = webSocket->CreateAndDispatchCloseEvent(mCloseEventWasClean,
-                                                        mCloseEventCode,
+                                                       mCloseEventCode,
-                                                        mCloseEventReason);
+                                                       mCloseEventReason);
   if (NS_FAILED(rv)) {
     NS_WARNING("Failed to dispatch the close event");
   }
 
-  mWebSocket->UpdateMustKeepAlive();
+  webSocket->UpdateMustKeepAlive();
   Disconnect();
 }
 

netwerk/protocol/websocket/BaseWebSocketChannel.cpp

@@ -10,6 +10,7 @@
 #include "nsILoadGroup.h"
 #include "nsIInterfaceRequestor.h"
 #include "nsAutoPtr.h"
+#include "nsProxyRelease.h"
 #include "nsStandardURL.h"
 
 #if defined(PR_LOGGING)
@@ -284,5 +285,26 @@ BaseWebSocketChannel::RetargetDeliveryTo(nsIEventTarget* aTargetThread)
   return NS_OK;
 }
 
+BaseWebSocketChannel::ListenerAndContextContainer::ListenerAndContextContainer(
+                                               nsIWebSocketListener* aListener,
+                                               nsISupports* aContext)
+  : mListener(aListener)
+  , mContext(aContext)
+{
+  MOZ_ASSERT(NS_IsMainThread());
+  MOZ_ASSERT(mListener);
+}
+
+BaseWebSocketChannel::ListenerAndContextContainer::~ListenerAndContextContainer()
+{
+  MOZ_ASSERT(mListener);
+
+  nsCOMPtr<nsIThread> mainThread;
+  NS_GetMainThread(getter_AddRefs(mainThread));
+
+  NS_ProxyRelease(mainThread, mListener, false);
+  NS_ProxyRelease(mainThread, mContext, false);
+}
+
 } // namespace net
 } // namespace mozilla

netwerk/protocol/websocket/BaseWebSocketChannel.h

@@ -57,11 +57,24 @@ class BaseWebSocketChannel : public nsIWebSocketChannel,
   virtual void GetEffectiveURL(nsAString& aEffectiveURL) const = 0;
   virtual bool IsEncrypted() const = 0;
 
+  class ListenerAndContextContainer MOZ_FINAL
+  {
+  public:
+    NS_INLINE_DECL_THREADSAFE_REFCOUNTING(ListenerAndContextContainer)
+
+    ListenerAndContextContainer(nsIWebSocketListener* aListener,
+                                nsISupports* aContext);
+
+    ~ListenerAndContextContainer();
+
+    nsCOMPtr<nsIWebSocketListener> mListener;
+    nsCOMPtr<nsISupports>          mContext;
+  };
+
  protected:
   nsCOMPtr<nsIURI>                mOriginalURI;
   nsCOMPtr<nsIURI>                mURI;
-  nsCOMPtr<nsIWebSocketListener>  mListener;
+  nsRefPtr<ListenerAndContextContainer> mListenerMT;
-  nsCOMPtr<nsISupports>           mContext;
   nsCOMPtr<nsIInterfaceRequestor> mCallbacks;
   nsCOMPtr<nsILoadGroup>          mLoadGroup;
   nsCOMPtr<nsILoadInfo>           mLoadInfo;

netwerk/protocol/websocket/WebSocketChannel.cpp

@@ -555,10 +555,11 @@ class CallOnMessageAvailable MOZ_FINAL : public nsIRunnable
 public:
   NS_DECL_THREADSAFE_ISUPPORTS
 
-  CallOnMessageAvailable(WebSocketChannel *aChannel,
+  CallOnMessageAvailable(WebSocketChannel* aChannel,
-                         nsCString        &aData,
+                         nsACString& aData,
-                         int32_t           aLen)
+                         int32_t aLen)
     : mChannel(aChannel),
+      mListenerMT(aChannel->mListenerMT),
       mData(aData),
       mLen(aLen) {}
 
@@ -566,19 +567,26 @@ public:
   {
     MOZ_ASSERT(mChannel->IsOnTargetThread());
 
-    if (mLen < 0)
+    if (mListenerMT) {
-      mChannel->mListener->OnMessageAvailable(mChannel->mContext, mData);
+      if (mLen < 0) {
-    else
+        mListenerMT->mListener->OnMessageAvailable(mListenerMT->mContext,
-      mChannel->mListener->OnBinaryMessageAvailable(mChannel->mContext, mData);
+                                                   mData);
+      } else {
+        mListenerMT->mListener->OnBinaryMessageAvailable(mListenerMT->mContext,
+                                                         mData);
+      }
+    }
+
     return NS_OK;
   }
 
 private:
   ~CallOnMessageAvailable() {}
 
-  nsRefPtr<WebSocketChannel>        mChannel;
+  nsRefPtr<WebSocketChannel> mChannel;
-  nsCString                         mData;
+  nsRefPtr<BaseWebSocketChannel::ListenerAndContextContainer> mListenerMT;
-  int32_t                           mLen;
+  nsCString mData;
+  int32_t mLen;
 };
 NS_IMPL_ISUPPORTS(CallOnMessageAvailable, nsIRunnable)
 
@@ -591,10 +599,12 @@ class CallOnStop MOZ_FINAL : public nsIRunnable
 public:
   NS_DECL_THREADSAFE_ISUPPORTS
 
-  CallOnStop(WebSocketChannel *aChannel,
+  CallOnStop(WebSocketChannel* aChannel,
-             nsresult          aReason)
+             nsresult aReason)
     : mChannel(aChannel),
-      mReason(aReason) {}
+      mListenerMT(mChannel->mListenerMT),
+      mReason(aReason)
+  {}
 
   NS_IMETHOD Run() MOZ_OVERRIDE
   {
@@ -602,19 +612,20 @@ public:
 
     nsWSAdmissionManager::OnStopSession(mChannel, mReason);
 
-    if (mChannel->mListener) {
+    if (mListenerMT) {
-      mChannel->mListener->OnStop(mChannel->mContext, mReason);
+      mListenerMT->mListener->OnStop(mListenerMT->mContext, mReason);
-      mChannel->mListener = nullptr;
+      mChannel->mListenerMT = nullptr;
-      mChannel->mContext = nullptr;
     }
+
     return NS_OK;
   }
 
 private:
   ~CallOnStop() {}
 
-  nsRefPtr<WebSocketChannel>        mChannel;
+  nsRefPtr<WebSocketChannel> mChannel;
-  nsresult                          mReason;
+  nsRefPtr<BaseWebSocketChannel::ListenerAndContextContainer> mListenerMT;
+  nsresult mReason;
 };
 NS_IMPL_ISUPPORTS(CallOnStop, nsIRunnable)
 
@@ -627,10 +638,11 @@ class CallOnServerClose MOZ_FINAL : public nsIRunnable
 public:
   NS_DECL_THREADSAFE_ISUPPORTS
 
-  CallOnServerClose(WebSocketChannel *aChannel,
+  CallOnServerClose(WebSocketChannel* aChannel,
-                    uint16_t          aCode,
+                    uint16_t aCode,
-                    nsCString        &aReason)
+                    nsACString& aReason)
     : mChannel(aChannel),
+      mListenerMT(mChannel->mListenerMT),
       mCode(aCode),
       mReason(aReason) {}
 
@@ -638,16 +650,20 @@ public:
   {
     MOZ_ASSERT(mChannel->IsOnTargetThread());
 
-    mChannel->mListener->OnServerClose(mChannel->mContext, mCode, mReason);
+    if (mListenerMT) {
+      mListenerMT->mListener->OnServerClose(mListenerMT->mContext, mCode,
+                                            mReason);
+    }
     return NS_OK;
   }
 
 private:
   ~CallOnServerClose() {}
 
-  nsRefPtr<WebSocketChannel>        mChannel;
+  nsRefPtr<WebSocketChannel> mChannel;
-  uint16_t                          mCode;
+  nsRefPtr<BaseWebSocketChannel::ListenerAndContextContainer> mListenerMT;
-  nsCString                         mReason;
+  uint16_t mCode;
+  nsCString mReason;
 };
 NS_IMPL_ISUPPORTS(CallOnServerClose, nsIRunnable)
 
@@ -658,9 +674,10 @@ NS_IMPL_ISUPPORTS(CallOnServerClose, nsIRunnable)
 class CallAcknowledge MOZ_FINAL : public nsCancelableRunnable
 {
 public:
-  CallAcknowledge(WebSocketChannel *aChannel,
+  CallAcknowledge(WebSocketChannel* aChannel,
-                  uint32_t          aSize)
+                  uint32_t aSize)
     : mChannel(aChannel),
+      mListenerMT(mChannel->mListenerMT),
       mSize(aSize) {}
 
   NS_IMETHOD Run()
@@ -668,15 +685,18 @@ public:
     MOZ_ASSERT(mChannel->IsOnTargetThread());
 
     LOG(("WebSocketChannel::CallAcknowledge: Size %u\n", mSize));
-    mChannel->mListener->OnAcknowledge(mChannel->mContext, mSize);
+    if (mListenerMT) {
+      mListenerMT->mListener->OnAcknowledge(mListenerMT->mContext, mSize);
+    }
     return NS_OK;
   }
 
 private:
   ~CallAcknowledge() {}
 
-  nsRefPtr<WebSocketChannel>        mChannel;
+  nsRefPtr<WebSocketChannel> mChannel;
-  uint32_t                          mSize;
+  nsRefPtr<BaseWebSocketChannel::ListenerAndContextContainer> mListenerMT;
+  uint32_t mSize;
 };
 
 //-----------------------------------------------------------------------------
@@ -1174,17 +1194,7 @@ WebSocketChannel::~WebSocketChannel()
     NS_ProxyRelease(mainThread, forgettable, false);
   }
 
-  if (mListener) {
+  mListenerMT = nullptr;
-    nsIWebSocketListener *forgettableListener;
-    mListener.forget(&forgettableListener);
-    NS_ProxyRelease(mainThread, forgettableListener, false);
-  }
-
-  if (mContext) {
-    nsISupports *forgettableContext;
-    mContext.forget(&forgettableContext);
-    NS_ProxyRelease(mainThread, forgettableContext, false);
-  }
 
   if (mLoadGroup) {
     nsILoadGroup *forgettableGroup;
@@ -1586,7 +1596,7 @@ WebSocketChannel::ProcessInput(uint8_t *buffer, uint32_t count)
       LOG(("WebSocketChannel:: %stext frame received\n",
            isDeflated ? "deflated " : ""));
 
-      if (mListener) {
+      if (mListenerMT) {
         nsCString utf8Data;
 
         if (isDeflated) {
@@ -1657,7 +1667,7 @@ WebSocketChannel::ProcessInput(uint8_t *buffer, uint32_t count)
           mCloseTimer->Cancel();
           mCloseTimer = nullptr;
         }
-        if (mListener) {
+        if (mListenerMT) {
           mTargetThread->Dispatch(new CallOnServerClose(this, mServerCloseCode,
                                                         mServerCloseReason),
                                   NS_DISPATCH_NORMAL);
@@ -1696,7 +1706,7 @@ WebSocketChannel::ProcessInput(uint8_t *buffer, uint32_t count)
       LOG(("WebSocketChannel:: %sbinary frame received\n",
            isDeflated ? "deflated " : ""));
 
-      if (mListener) {
+      if (mListenerMT) {
         nsCString binaryData;
 
         if (isDeflated) {
@@ -2265,8 +2275,11 @@ WebSocketChannel::StopSession(nsresult reason)
 
   if (!mCalledOnStop) {
     mCalledOnStop = 1;
-    mTargetThread->Dispatch(new CallOnStop(this, reason),
+
-                            NS_DISPATCH_NORMAL);
+    nsWSAdmissionManager::OnStopSession(this, reason);
+
+    nsRefPtr<CallOnStop> runnable = new CallOnStop(this, reason);
+    mTargetThread->Dispatch(runnable, NS_DISPATCH_NORMAL);
   }
 }
 
@@ -2577,10 +2590,10 @@ WebSocketChannel::StartWebsocketData()
   nsWSAdmissionManager::OnConnected(this);
 
   LOG(("WebSocketChannel::StartWebsocketData Notifying Listener %p\n",
-       mListener.get()));
+       mListenerMT ? mListenerMT->mListener.get() : nullptr));
 
-  if (mListener) {
+  if (mListenerMT) {
-    mListener->OnStart(mContext);
+    mListenerMT->mListener->OnStart(mListenerMT->mContext);
   }
 
   // Start keepalive ping timer, if we're using keepalive.
@@ -2943,7 +2956,7 @@ WebSocketChannel::AsyncOpen(nsIURI *aURI,
     return NS_ERROR_UNEXPECTED;
   }
 
-  if (mListener || mWasOpened)
+  if (mListenerMT || mWasOpened)
     return NS_ERROR_ALREADY_OPENED;
 
   nsresult rv;
@@ -3109,8 +3122,7 @@ WebSocketChannel::AsyncOpen(nsIURI *aURI,
   // Only set these if the open was successful:
   //
   mWasOpened = 1;
-  mListener = aListener;
+  mListenerMT = new ListenerAndContextContainer(aListener, aContext);
-  mContext = aContext;
   IncrementSessionCount();
 
   return rv;

netwerk/protocol/websocket/WebSocketChannelChild.cpp

@@ -204,9 +204,9 @@ WebSocketChannelChild::OnStart(const nsCString& aProtocol,
   mEffectiveURL = aEffectiveURL;
   mEncrypted = aEncrypted;
 
-  if (mListener) {
+  if (mListenerMT) {
     AutoEventEnqueuer ensureSerialDispatch(mEventQ);
-    mListener->OnStart(mContext);
+    mListenerMT->mListener->OnStart(mListenerMT->mContext);
   }
 }
 
@@ -246,9 +246,9 @@ void
 WebSocketChannelChild::OnStop(const nsresult& aStatusCode)
 {
   LOG(("WebSocketChannelChild::RecvOnStop() %p\n", this));
-  if (mListener) {
+  if (mListenerMT) {
     AutoEventEnqueuer ensureSerialDispatch(mEventQ);
-    mListener->OnStop(mContext, aStatusCode);
+    mListenerMT->mListener->OnStop(mListenerMT->mContext, aStatusCode);
   }
 }
 
@@ -295,9 +295,9 @@ void
 WebSocketChannelChild::OnMessageAvailable(const nsCString& aMsg)
 {
   LOG(("WebSocketChannelChild::RecvOnMessageAvailable() %p\n", this));
-  if (mListener) {
+  if (mListenerMT) {
     AutoEventEnqueuer ensureSerialDispatch(mEventQ);
-    mListener->OnMessageAvailable(mContext, aMsg);
+    mListenerMT->mListener->OnMessageAvailable(mListenerMT->mContext, aMsg);
   }
 }
 
@@ -319,9 +319,10 @@ void
 WebSocketChannelChild::OnBinaryMessageAvailable(const nsCString& aMsg)
 {
   LOG(("WebSocketChannelChild::RecvOnBinaryMessageAvailable() %p\n", this));
-  if (mListener) {
+  if (mListenerMT) {
     AutoEventEnqueuer ensureSerialDispatch(mEventQ);
-    mListener->OnBinaryMessageAvailable(mContext, aMsg);
+    mListenerMT->mListener->OnBinaryMessageAvailable(mListenerMT->mContext,
+                                                     aMsg);
   }
 }
 
@@ -361,9 +362,9 @@ void
 WebSocketChannelChild::OnAcknowledge(const uint32_t& aSize)
 {
   LOG(("WebSocketChannelChild::RecvOnAcknowledge() %p\n", this));
-  if (mListener) {
+  if (mListenerMT) {
     AutoEventEnqueuer ensureSerialDispatch(mEventQ);
-    mListener->OnAcknowledge(mContext, aSize);
+    mListenerMT->mListener->OnAcknowledge(mListenerMT->mContext, aSize);
   }
 }
 
@@ -409,9 +410,10 @@ WebSocketChannelChild::OnServerClose(const uint16_t& aCode,
                                      const nsCString& aReason)
 {
   LOG(("WebSocketChannelChild::RecvOnServerClose() %p\n", this));
-  if (mListener) {
+  if (mListenerMT) {
     AutoEventEnqueuer ensureSerialDispatch(mEventQ);
-    mListener->OnServerClose(mContext, aCode, aReason);
+    mListenerMT->mListener->OnServerClose(mListenerMT->mContext, aCode,
+                                          aReason);
   }
 }
 
@@ -424,7 +426,7 @@ WebSocketChannelChild::AsyncOpen(nsIURI *aURI,
   LOG(("WebSocketChannelChild::AsyncOpen() %p\n", this));
 
   NS_ABORT_IF_FALSE(NS_IsMainThread(), "not main thread");
-  NS_ABORT_IF_FALSE(aURI && aListener && !mListener, 
+  NS_ABORT_IF_FALSE(aURI && aListener && !mListenerMT,
                     "Invalid state for WebSocketChannelChild::AsyncOpen");
 
   mozilla::dom::TabChild* tabChild = nullptr;
@@ -454,8 +456,7 @@ WebSocketChannelChild::AsyncOpen(nsIURI *aURI,
 
   mOriginalURI = aURI;
   mURI = mOriginalURI;
-  mListener = aListener;
+  mListenerMT = new ListenerAndContextContainer(aListener, aContext);
-  mContext = aContext;
   mOrigin = aOrigin;
   mWasOpened = 1;