dom/base/WebSocket.cpp

   1 /* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
   2 /* vim: set sw=2 ts=8 et tw=80 : */
   3 /* This Source Code Form is subject to the terms of the Mozilla Public
   4  * License, v. 2.0. If a copy of the MPL was not distributed with this
   5  * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
   6
   7 #include "WebSocket.h"
   8 #include "mozilla/dom/WebSocketBinding.h"
   9 #include "mozilla/net/WebSocketChannel.h"
  10
  11 #include "jsapi.h"
  12 #include "jsfriendapi.h"
  13 #include "mozilla/DOMEventTargetHelper.h"
  14 #include "mozilla/net/WebSocketChannel.h"
  15 #include "mozilla/dom/File.h"
  16 #include "mozilla/dom/ScriptSettings.h"
  17 #include "mozilla/dom/WorkerPrivate.h"
  18 #include "mozilla/dom/WorkerRunnable.h"
  19 #include "mozilla/dom/WorkerScope.h"
  20 #include "nsIScriptGlobalObject.h"
  21 #include "nsIDOMWindow.h"
  22 #include "nsIDocument.h"
  23 #include "nsXPCOM.h"
  24 #include "nsIXPConnect.h"
  25 #include "nsContentUtils.h"
  26 #include "nsError.h"
  27 #include "nsIScriptObjectPrincipal.h"
  28 #include "nsIURL.h"
  29 #include "nsIUnicodeEncoder.h"
  30 #include "nsThreadUtils.h"
  31 #include "nsIDOMMessageEvent.h"
  32 #include "nsIPromptFactory.h"
  33 #include "nsIWindowWatcher.h"
  34 #include "nsIPrompt.h"
  35 #include "nsIStringBundle.h"
  36 #include "nsIConsoleService.h"
  37 #include "mozilla/dom/CloseEvent.h"
  38 #include "nsICryptoHash.h"
  39 #include "nsJSUtils.h"
  40 #include "nsIScriptError.h"
  41 #include "nsNetUtil.h"
  42 #include "nsILoadGroup.h"
  43 #include "mozilla/Preferences.h"
  44 #include "xpcpublic.h"
  45 #include "nsContentPolicyUtils.h"
  46 #include "nsWrapperCacheInlines.h"
  47 #include "nsIObserverService.h"
  48 #include "nsIEventTarget.h"
  49 #include "nsIInterfaceRequestor.h"
  50 #include "nsIObserver.h"
  51 #include "nsIRequest.h"
  52 #include "nsIThreadRetargetableRequest.h"
  53 #include "nsIWebSocketChannel.h"
  54 #include "nsIWebSocketListener.h"
  55 #include "nsProxyRelease.h"
  56 #include "nsWeakReference.h"
  57
  58 using namespace mozilla::net;
  59 using namespace mozilla::dom::workers;
  60
  61 namespace mozilla {
  62 namespace dom {
  63
  64 class WebSocketImpl MOZ_FINAL : public nsIInterfaceRequestor
  65                               , public nsIWebSocketListener
  66                               , public nsIObserver
  67                               , public nsSupportsWeakReference
  68                               , public nsIRequest
  69                               , public nsIEventTarget
  70 {
  71 public:
  72   NS_DECL_NSIINTERFACEREQUESTOR
  73   NS_DECL_NSIWEBSOCKETLISTENER
  74   NS_DECL_NSIOBSERVER
  75   NS_DECL_NSIREQUEST
  76   NS_DECL_THREADSAFE_ISUPPORTS
  77   NS_DECL_NSIEVENTTARGET
  78
  79   explicit WebSocketImpl(WebSocket* aWebSocket)
  80   : mWebSocket(aWebSocket)
  81   , mOnCloseScheduled(false)
  82   , mFailed(false)
  83   , mDisconnectingOrDisconnected(false)
  84   , mCloseEventWasClean(false)
  85   , mCloseEventCode(nsIWebSocketChannel::CLOSE_ABNORMAL)
  86   , mScriptLine(0)
  87   , mInnerWindowID(0)
  88   , mWorkerPrivate(nullptr)
  89 #ifdef DEBUG
  90   , mHasFeatureRegistered(false)
  91 #endif
  92   , mIsMainThread(true)
  93   , mWorkerShuttingDown(false)
  94   {
  95     if (!NS_IsMainThread()) {
  96       mWorkerPrivate = GetCurrentThreadWorkerPrivate();
  97       MOZ_ASSERT(mWorkerPrivate);
  98       mIsMainThread = false;
  99     }
 100   }
 101
 102   void AssertIsOnTargetThread() const
 103   {
 104     MOZ_ASSERT(IsTargetThread());
 105   }
 106
 107   bool IsTargetThread() const;
 108
 109   void Init(JSContext* aCx,
 110             nsIPrincipal* aPrincipal,
 111             const nsAString& aURL,
 112             nsTArray<nsString>& aProtocolArray,
 113             const nsACString& aScriptFile,
 114             uint32_t aScriptLine,
 115             ErrorResult& aRv,
 116             bool* aConnectionFailed);
 117
 118   void AsyncOpen(ErrorResult& aRv);
 119
 120   nsresult ParseURL(const nsAString& aURL);
 121   nsresult InitializeConnection();
 122
 123   // These methods when called can release the WebSocket object
 124   void FailConnection(uint16_t reasonCode,
 125                       const nsACString& aReasonString = EmptyCString());
 126   nsresult CloseConnection(uint16_t reasonCode,
 127                            const nsACString& aReasonString = EmptyCString());
 128   nsresult Disconnect();
 129   void DisconnectInternal();
 130
 131   nsresult ConsoleError();
 132   nsresult PrintErrorOnConsole(const char* aBundleURI,
 133                                const char16_t* aError,
 134                                const char16_t** aFormatStrings,
 135                                uint32_t aFormatStringsLen);
 136
 137   nsresult DoOnMessageAvailable(const nsACString& aMsg,
 138                                 bool isBinary);
 139
 140   // ConnectionCloseEvents: 'error' event if needed, then 'close' event.
 141   // - These must not be dispatched while we are still within an incoming call
 142   //   from JS (ex: close()).  Set 'sync' to false in that case to dispatch in a
 143   //   separate new event.
 144   nsresult ScheduleConnectionCloseEvents(nsISupports* aContext,
 145                                          nsresult aStatusCode,
 146                                          bool sync);
 147   // 2nd half of ScheduleConnectionCloseEvents, sometimes run in its own event.
 148   void DispatchConnectionCloseEvents();
 149
 150   // Dispatch a runnable to the right thread.
 151   nsresult DispatchRunnable(nsIRunnable* aRunnable);
 152
 153   nsresult UpdateURI();
 154
 155   void AddRefObject();
 156   void ReleaseObject();
 157
 158   bool RegisterFeature();
 159   void UnregisterFeature();
 160
 161   nsresult CancelInternal();
 162
 163   nsRefPtr<WebSocket> mWebSocket;
 164
 165   nsCOMPtr<nsIWebSocketChannel> mChannel;
 166
 167   bool mSecure; // if true it is using SSL and the wss scheme,
 168                 // otherwise it is using the ws scheme with no SSL
 169
 170   bool mOnCloseScheduled;
 171   bool mFailed;
 172   bool mDisconnectingOrDisconnected;
 173
 174   // Set attributes of DOM 'onclose' message
 175   bool      mCloseEventWasClean;
 176   nsString  mCloseEventReason;
 177   uint16_t  mCloseEventCode;
 178
 179   nsCString mAsciiHost;  // hostname
 180   uint32_t  mPort;
 181   nsCString mResource; // [filepath[?query]]
 182   nsString  mUTF16Origin;
 183
 184   nsCString mURI;
 185   nsCString mRequestedProtocolList;
 186
 187   nsCOMPtr<nsIPrincipal> mPrincipal;
 188   nsWeakPtr              mOriginDocument;
 189
 190   // Web Socket owner information:
 191   // - the script file name, UTF8 encoded.
 192   // - source code line number where the Web Socket object was constructed.
 193   // - the ID of the inner window where the script lives. Note that this may not
 194   //   be the same as the Web Socket owner window.
 195   // These attributes are used for error reporting.
 196   nsCString mScriptFile;
 197   uint32_t mScriptLine;
 198   uint64_t mInnerWindowID;
 199
 200   WorkerPrivate* mWorkerPrivate;
 201   nsAutoPtr<WorkerFeature> mWorkerFeature;
 202
 203 #ifdef DEBUG
 204   // This is protected by mutex.
 205   bool mHasFeatureRegistered;
 206
 207   bool HasFeatureRegistered()
 208   {
 209     MOZ_ASSERT(mWebSocket);
 210     MutexAutoLock lock(mWebSocket->mMutex);
 211     return mHasFeatureRegistered;
 212   }
 213
 214   void SetHasFeatureRegistered(bool aValue)
 215   {
 216     MOZ_ASSERT(mWebSocket);
 217     MutexAutoLock lock(mWebSocket->mMutex);
 218     mHasFeatureRegistered = aValue;
 219   }
 220 #endif
 221
 222   nsWeakPtr mWeakLoadGroup;
 223
 224   bool mIsMainThread;
 225   bool mWorkerShuttingDown;
 226
 227 private:
 228   ~WebSocketImpl()
 229   {
 230     // If we threw during Init we never called disconnect
 231     if (!mDisconnectingOrDisconnected) {
 232       Disconnect();
 233     }
 234   }
 235 };
 236
 237 NS_IMPL_ISUPPORTS(WebSocketImpl,
 238                   nsIInterfaceRequestor,
 239                   nsIWebSocketListener,
 240                   nsIObserver,
 241                   nsISupportsWeakReference,
 242                   nsIRequest,
 243                   nsIEventTarget)
 244
 245 class CallDispatchConnectionCloseEvents MOZ_FINAL : public nsCancelableRunnable
 246 {
 247 public:
 248   explicit CallDispatchConnectionCloseEvents(WebSocketImpl* aWebSocketImpl)
 249     : mWebSocketImpl(aWebSocketImpl)
 250   {
 251     aWebSocketImpl->AssertIsOnTargetThread();
 252   }
 253
 254   NS_IMETHOD Run()
 255   {
 256     mWebSocketImpl->AssertIsOnTargetThread();
 257     mWebSocketImpl->DispatchConnectionCloseEvents();
 258     return NS_OK;
 259   }
 260
 261 private:
 262   nsRefPtr<WebSocketImpl> mWebSocketImpl;
 263 };
 264
 265 //-----------------------------------------------------------------------------
 266 // WebSocketImpl
 267 //-----------------------------------------------------------------------------
 268
 269 namespace {
 270
 271 class PrintErrorOnConsoleRunnable MOZ_FINAL : public WorkerMainThreadRunnable
 272 {
 273 public:
 274   PrintErrorOnConsoleRunnable(WebSocketImpl* aImpl,
 275                               const char* aBundleURI,
 276                               const char16_t* aError,
 277                               const char16_t** aFormatStrings,
 278                               uint32_t aFormatStringsLen)
 279     : WorkerMainThreadRunnable(aImpl->mWorkerPrivate)
 280     , mImpl(aImpl)
 281     , mBundleURI(aBundleURI)
 282     , mError(aError)
 283     , mFormatStrings(aFormatStrings)
 284     , mFormatStringsLen(aFormatStringsLen)
 285     , mRv(NS_ERROR_FAILURE)
 286   { }
 287
 288   bool MainThreadRun() MOZ_OVERRIDE
 289   {
 290     mRv = mImpl->PrintErrorOnConsole(mBundleURI, mError, mFormatStrings,
 291                                      mFormatStringsLen);
 292     return true;
 293   }
 294
 295   nsresult ErrorCode() const
 296   {
 297     return mRv;
 298   }
 299
 300 private:
 301   // Raw pointer because this runnable is sync.
 302   WebSocketImpl* mImpl;
 303
 304   const char* mBundleURI;
 305   const char16_t* mError;
 306   const char16_t** mFormatStrings;
 307   uint32_t mFormatStringsLen;
 308   nsresult mRv;
 309 };
 310
 311 } // anonymous namespace
 312
 313 nsresult
 314 WebSocketImpl::PrintErrorOnConsole(const char *aBundleURI,
 315                                    const char16_t *aError,
 316                                    const char16_t **aFormatStrings,
 317                                    uint32_t aFormatStringsLen)
 318 {
 319   // This method must run on the main thread.
 320
 321   if (!NS_IsMainThread()) {
 322     MOZ_ASSERT(mWorkerPrivate);
 323
 324     nsRefPtr<PrintErrorOnConsoleRunnable> runnable =
 325       new PrintErrorOnConsoleRunnable(this, aBundleURI, aError, aFormatStrings,
 326                                       aFormatStringsLen);
 327     runnable->Dispatch(mWorkerPrivate->GetJSContext());
 328     return runnable->ErrorCode();
 329   }
 330
 331   nsresult rv;
 332   nsCOMPtr<nsIStringBundleService> bundleService =
 333     do_GetService(NS_STRINGBUNDLE_CONTRACTID, &rv);
 334   NS_ENSURE_SUCCESS(rv, rv);
 335
 336   nsCOMPtr<nsIStringBundle> strBundle;
 337   rv = bundleService->CreateBundle(aBundleURI, getter_AddRefs(strBundle));
 338   NS_ENSURE_SUCCESS(rv, rv);
 339
 340   nsCOMPtr<nsIConsoleService> console(
 341     do_GetService(NS_CONSOLESERVICE_CONTRACTID, &rv));
 342   NS_ENSURE_SUCCESS(rv, rv);
 343
 344   nsCOMPtr<nsIScriptError> errorObject(
 345     do_CreateInstance(NS_SCRIPTERROR_CONTRACTID, &rv));
 346   NS_ENSURE_SUCCESS(rv, rv);
 347
 348   // Localize the error message
 349   nsXPIDLString message;
 350   if (aFormatStrings) {
 351     rv = strBundle->FormatStringFromName(aError, aFormatStrings,
 352                                          aFormatStringsLen,
 353                                          getter_Copies(message));
 354   } else {
 355     rv = strBundle->GetStringFromName(aError, getter_Copies(message));
 356   }
 357   NS_ENSURE_SUCCESS(rv, rv);
 358
 359   if (mInnerWindowID) {
 360     rv = errorObject->InitWithWindowID(message,
 361                                        NS_ConvertUTF8toUTF16(mScriptFile),
 362                                        EmptyString(), mScriptLine, 0,
 363                                        nsIScriptError::errorFlag, "Web Socket",
 364                                        mInnerWindowID);
 365   } else {
 366     rv = errorObject->Init(message,
 367                            NS_ConvertUTF8toUTF16(mScriptFile),
 368                            EmptyString(), mScriptLine, 0,
 369                            nsIScriptError::errorFlag, "Web Socket");
 370   }
 371
 372   NS_ENSURE_SUCCESS(rv, rv);
 373
 374   // print the error message directly to the JS console
 375   rv = console->LogMessage(errorObject);
 376   NS_ENSURE_SUCCESS(rv, rv);
 377
 378   return NS_OK;
 379 }
 380
 381 namespace {
 382
 383 class CloseRunnable MOZ_FINAL : public WorkerMainThreadRunnable
 384 {
 385 public:
 386   CloseRunnable(WebSocketImpl* aImpl, uint16_t aReasonCode,
 387                 const nsACString& aReasonString)
 388     : WorkerMainThreadRunnable(aImpl->mWorkerPrivate)
 389     , mImpl(aImpl)
 390     , mReasonCode(aReasonCode)
 391     , mReasonString(aReasonString)
 392     , mRv(NS_ERROR_FAILURE)
 393   { }
 394
 395   bool MainThreadRun() MOZ_OVERRIDE
 396   {
 397     mRv = mImpl->mChannel->Close(mReasonCode, mReasonString);
 398     return true;
 399   }
 400
 401   nsresult ErrorCode() const
 402   {
 403     return mRv;
 404   }
 405
 406 private:
 407   // A raw pointer because this runnable is sync.
 408   WebSocketImpl* mImpl;
 409
 410   uint16_t mReasonCode;
 411   const nsACString& mReasonString;
 412   nsresult mRv;
 413 };
 414
 415 class MOZ_STACK_CLASS MaybeDisconnect
 416 {
 417 public:
 418   explicit MaybeDisconnect(WebSocketImpl* aImpl)
 419     : mImpl(aImpl)
 420   {
 421   }
 422
 423   ~MaybeDisconnect()
 424   {
 425     if (mImpl->mWorkerShuttingDown) {
 426       mImpl->Disconnect();
 427     }
 428   }
 429
 430 private:
 431   WebSocketImpl* mImpl;
 432 };
 433
 434 } // anonymous namespace
 435
 436 nsresult
 437 WebSocketImpl::CloseConnection(uint16_t aReasonCode,
 438                                const nsACString& aReasonString)
 439 {
 440   AssertIsOnTargetThread();
 441
 442   if (mDisconnectingOrDisconnected) {
 443     return NS_OK;
 444   }
 445
 446   // If this method is called because the worker is going away, we will not
 447   // receive the OnStop() method and we have to disconnect the WebSocket and
 448   // release the WorkerFeature.
 449   MaybeDisconnect md(this);
 450
 451   uint16_t readyState = mWebSocket->ReadyState();
 452   if (readyState == WebSocket::CLOSING ||
 453       readyState == WebSocket::CLOSED) {
 454     return NS_OK;
 455   }
 456
 457   // The common case...
 458   if (mChannel) {
 459     mWebSocket->SetReadyState(WebSocket::CLOSING);
 460
 461     // The channel has to be closed on the main-thread.
 462
 463     if (NS_IsMainThread()) {
 464       return mChannel->Close(aReasonCode, aReasonString);
 465     }
 466
 467     nsRefPtr<CloseRunnable> runnable =
 468       new CloseRunnable(this, aReasonCode, aReasonString);
 469     runnable->Dispatch(mWorkerPrivate->GetJSContext());
 470     return runnable->ErrorCode();
 471   }
 472
 473   // No channel, but not disconnected: canceled or failed early
 474   //
 475   MOZ_ASSERT(readyState == WebSocket::CONNECTING,
 476              "Should only get here for early websocket cancel/error");
 477
 478   // Server won't be sending us a close code, so use what's passed in here.
 479   mCloseEventCode = aReasonCode;
 480   CopyUTF8toUTF16(aReasonString, mCloseEventReason);
 481
 482   mWebSocket->SetReadyState(WebSocket::CLOSING);
 483
 484   // Can be called from Cancel() or Init() codepaths, so need to dispatch
 485   // onerror/onclose asynchronously
 486   ScheduleConnectionCloseEvents(
 487                     nullptr,
 488                     (aReasonCode == nsIWebSocketChannel::CLOSE_NORMAL ||
 489                      aReasonCode == nsIWebSocketChannel::CLOSE_GOING_AWAY) ?
 490                      NS_OK : NS_ERROR_FAILURE,
 491                     false);
 492
 493   return NS_OK;
 494 }
 495
 496 nsresult
 497 WebSocketImpl::ConsoleError()
 498 {
 499   AssertIsOnTargetThread();
 500
 501   NS_ConvertUTF8toUTF16 specUTF16(mURI);
 502   const char16_t* formatStrings[] = { specUTF16.get() };
 503
 504   if (mWebSocket->ReadyState() < WebSocket::OPEN) {
 505     PrintErrorOnConsole("chrome://global/locale/appstrings.properties",
 506                         MOZ_UTF16("connectionFailure"),
 507                         formatStrings, ArrayLength(formatStrings));
 508   } else {
 509     PrintErrorOnConsole("chrome://global/locale/appstrings.properties",
 510                         MOZ_UTF16("netInterrupt"),
 511                         formatStrings, ArrayLength(formatStrings));
 512   }
 513   /// todo some specific errors - like for message too large
 514   return NS_OK;
 515 }
 516
 517 void
 518 WebSocketImpl::FailConnection(uint16_t aReasonCode,
 519                               const nsACString& aReasonString)
 520 {
 521   AssertIsOnTargetThread();
 522
 523   if (mDisconnectingOrDisconnected) {
 524     return;
 525   }
 526
 527   ConsoleError();
 528   mFailed = true;
 529   CloseConnection(aReasonCode, aReasonString);
 530 }
 531
 532 namespace {
 533
 534 class DisconnectInternalRunnable MOZ_FINAL : public WorkerMainThreadRunnable
 535 {
 536 public:
 537   explicit DisconnectInternalRunnable(WebSocketImpl* aImpl)
 538     : WorkerMainThreadRunnable(aImpl->mWorkerPrivate)
 539     , mImpl(aImpl)
 540   { }
 541
 542   bool MainThreadRun() MOZ_OVERRIDE
 543   {
 544     mImpl->DisconnectInternal();
 545     return true;
 546   }
 547
 548 private:
 549   // A raw pointer because this runnable is sync.
 550   WebSocketImpl* mImpl;
 551 };
 552
 553 } // anonymous namespace
 554
 555 nsresult
 556 WebSocketImpl::Disconnect()
 557 {
 558   if (mDisconnectingOrDisconnected) {
 559     return NS_OK;
 560   }
 561
 562   AssertIsOnTargetThread();
 563
 564   // Disconnect can be called from some control event (such as Notify() of
 565   // WorkerFeature). This will be schedulated before any other sync/async
 566   // runnable. In order to prevent some double Disconnect() calls, we use this
 567   // boolean.
 568   mDisconnectingOrDisconnected = true;
 569
 570   // DisconnectInternal touches observers and nsILoadGroup and it must run on
 571   // the main thread.
 572
 573   if (NS_IsMainThread()) {
 574     DisconnectInternal();
 575   } else {
 576     nsRefPtr<DisconnectInternalRunnable> runnable =
 577       new DisconnectInternalRunnable(this);
 578     runnable->Dispatch(mWorkerPrivate->GetJSContext());
 579   }
 580
 581   // DontKeepAliveAnyMore() can release the object. So hold a reference to this
 582   // until the end of the method.
 583   nsRefPtr<WebSocketImpl> kungfuDeathGrip = this;
 584
 585   nsCOMPtr<nsIThread> mainThread;
 586   if (NS_FAILED(NS_GetMainThread(getter_AddRefs(mainThread))) ||
 587       NS_FAILED(NS_ProxyRelease(mainThread, mChannel))) {
 588     NS_WARNING("Failed to proxy release of channel, leaking instead!");
 589   }
 590
 591   mWebSocket->DontKeepAliveAnyMore();
 592   mWebSocket->mImpl = nullptr;
 593
 594   if (mWorkerPrivate && mWorkerFeature) {
 595     UnregisterFeature();
 596   }
 597
 598   // We want to release the WebSocket in the correct thread.
 599   mWebSocket = nullptr;
 600
 601   return NS_OK;
 602 }
 603
 604 void
 605 WebSocketImpl::DisconnectInternal()
 606 {
 607   AssertIsOnMainThread();
 608
 609   nsCOMPtr<nsILoadGroup> loadGroup = do_QueryReferent(mWeakLoadGroup);
 610   if (loadGroup) {
 611     loadGroup->RemoveRequest(this, nullptr, NS_OK);
 612     // mWeakLoadGroup has to be release on main-thread because WeakReferences
 613     // are not thread-safe.
 614     mWeakLoadGroup = nullptr;
 615   }
 616
 617   nsCOMPtr<nsIObserverService> os = mozilla::services::GetObserverService();
 618   if (os) {
 619     os->RemoveObserver(this, DOM_WINDOW_DESTROYED_TOPIC);
 620     os->RemoveObserver(this, DOM_WINDOW_FROZEN_TOPIC);
 621   }
 622 }
 623
 624 //-----------------------------------------------------------------------------
 625 // WebSocketImpl::nsIWebSocketListener methods:
 626 //-----------------------------------------------------------------------------
 627
 628 nsresult
 629 WebSocketImpl::DoOnMessageAvailable(const nsACString& aMsg, bool isBinary)
 630 {
 631   AssertIsOnTargetThread();
 632
 633   if (mDisconnectingOrDisconnected) {
 634     return NS_OK;
 635   }
 636
 637   int16_t readyState = mWebSocket->ReadyState();
 638   if (readyState == WebSocket::CLOSED) {
 639     NS_ERROR("Received message after CLOSED");
 640     return NS_ERROR_UNEXPECTED;
 641   }
 642
 643   if (readyState == WebSocket::OPEN) {
 644     // Dispatch New Message
 645     nsresult rv = mWebSocket->CreateAndDispatchMessageEvent(aMsg, isBinary);
 646     if (NS_FAILED(rv)) {
 647       NS_WARNING("Failed to dispatch the message event");
 648     }
 649
 650     return NS_OK;
 651   }
 652
 653   // CLOSING should be the only other state where it's possible to get msgs
 654   // from channel: Spec says to drop them.
 655   MOZ_ASSERT(readyState == WebSocket::CLOSING,
 656              "Received message while CONNECTING or CLOSED");
 657   return NS_OK;
 658 }
 659
 660 NS_IMETHODIMP
 661 WebSocketImpl::OnMessageAvailable(nsISupports* aContext,
 662                                   const nsACString& aMsg)
 663 {
 664   AssertIsOnTargetThread();
 665
 666   if (mDisconnectingOrDisconnected) {
 667     return NS_OK;
 668   }
 669
 670   return DoOnMessageAvailable(aMsg, false);
 671 }
 672
 673 NS_IMETHODIMP
 674 WebSocketImpl::OnBinaryMessageAvailable(nsISupports* aContext,
 675                                         const nsACString& aMsg)
 676 {
 677   AssertIsOnTargetThread();
 678
 679   if (mDisconnectingOrDisconnected) {
 680     return NS_OK;
 681   }
 682
 683   return DoOnMessageAvailable(aMsg, true);
 684 }
 685
 686 NS_IMETHODIMP
 687 WebSocketImpl::OnStart(nsISupports* aContext)
 688 {
 689   AssertIsOnTargetThread();
 690
 691   if (mDisconnectingOrDisconnected) {
 692     return NS_OK;
 693   }
 694
 695   int16_t readyState = mWebSocket->ReadyState();
 696
 697   // This is the only function that sets OPEN, and should be called only once
 698   MOZ_ASSERT(readyState != WebSocket::OPEN,
 699              "readyState already OPEN! OnStart called twice?");
 700
 701   // Nothing to do if we've already closed/closing
 702   if (readyState != WebSocket::CONNECTING) {
 703     return NS_OK;
 704   }
 705
 706   // Attempt to kill "ghost" websocket: but usually too early for check to fail
 707   nsresult rv = mWebSocket->CheckInnerWindowCorrectness();
 708   if (NS_FAILED(rv)) {
 709     CloseConnection(nsIWebSocketChannel::CLOSE_GOING_AWAY);
 710     return rv;
 711   }
 712
 713   if (!mRequestedProtocolList.IsEmpty()) {
 714     mChannel->GetProtocol(mWebSocket->mEstablishedProtocol);
 715   }
 716
 717   mChannel->GetExtensions(mWebSocket->mEstablishedExtensions);
 718   UpdateURI();
 719
 720   mWebSocket->SetReadyState(WebSocket::OPEN);
 721
 722   // Let's keep the object alive because the webSocket can be CCed in the
 723   // onopen callback.
 724   nsRefPtr<WebSocket> webSocket = mWebSocket;
 725
 726   // Call 'onopen'
 727   rv = webSocket->CreateAndDispatchSimpleEvent(NS_LITERAL_STRING("open"));
 728   if (NS_FAILED(rv)) {
 729     NS_WARNING("Failed to dispatch the open event");
 730   }
 731
 732   webSocket->UpdateMustKeepAlive();
 733   return NS_OK;
 734 }
 735
 736 NS_IMETHODIMP
 737 WebSocketImpl::OnStop(nsISupports* aContext, nsresult aStatusCode)
 738 {
 739   AssertIsOnTargetThread();
 740
 741   if (mDisconnectingOrDisconnected) {
 742     return NS_OK;
 743   }
 744
 745   // We can be CONNECTING here if connection failed.
 746   // We can be OPEN if we have encountered a fatal protocol error
 747   // We can be CLOSING if close() was called and/or server initiated close.
 748   MOZ_ASSERT(mWebSocket->ReadyState() != WebSocket::CLOSED,
 749              "Shouldn't already be CLOSED when OnStop called");
 750
 751   // called by network stack, not JS, so can dispatch JS events synchronously
 752   return ScheduleConnectionCloseEvents(aContext, aStatusCode, true);
 753 }
 754
 755 nsresult
 756 WebSocketImpl::ScheduleConnectionCloseEvents(nsISupports* aContext,
 757                                              nsresult aStatusCode,
 758                                              bool sync)
 759 {
 760   AssertIsOnTargetThread();
 761
 762   // no-op if some other code has already initiated close event
 763   if (!mOnCloseScheduled) {
 764     mCloseEventWasClean = NS_SUCCEEDED(aStatusCode);
 765
 766     if (aStatusCode == NS_BASE_STREAM_CLOSED) {
 767       // don't generate an error event just because of an unclean close
 768       aStatusCode = NS_OK;
 769     }
 770
 771     if (NS_FAILED(aStatusCode)) {
 772       ConsoleError();
 773       mFailed = true;
 774     }
 775
 776     mOnCloseScheduled = true;
 777
 778     if (sync) {
 779       DispatchConnectionCloseEvents();
 780     } else {
 781       NS_DispatchToCurrentThread(new CallDispatchConnectionCloseEvents(this));
 782     }
 783   }
 784
 785   return NS_OK;
 786 }
 787
 788 NS_IMETHODIMP
 789 WebSocketImpl::OnAcknowledge(nsISupports *aContext, uint32_t aSize)
 790 {
 791   AssertIsOnTargetThread();
 792
 793   if (mDisconnectingOrDisconnected) {
 794     return NS_OK;
 795   }
 796
 797   if (aSize > mWebSocket->mOutgoingBufferedAmount) {
 798     return NS_ERROR_UNEXPECTED;
 799   }
 800
 801   mWebSocket->mOutgoingBufferedAmount -= aSize;
 802   return NS_OK;
 803 }
 804
 805 NS_IMETHODIMP
 806 WebSocketImpl::OnServerClose(nsISupports *aContext, uint16_t aCode,
 807                              const nsACString &aReason)
 808 {
 809   AssertIsOnTargetThread();
 810
 811   if (mDisconnectingOrDisconnected) {
 812     return NS_OK;
 813   }
 814
 815   int16_t readyState = mWebSocket->ReadyState();
 816
 817   MOZ_ASSERT(readyState != WebSocket::CONNECTING,
 818              "Received server close before connected?");
 819   MOZ_ASSERT(readyState != WebSocket::CLOSED,
 820              "Received server close after already closed!");
 821
 822   // store code/string for onclose DOM event
 823   mCloseEventCode = aCode;
 824   CopyUTF8toUTF16(aReason, mCloseEventReason);
 825
 826   if (readyState == WebSocket::OPEN) {
 827     // Server initiating close.
 828     // RFC 6455, 5.5.1: "When sending a Close frame in response, the endpoint
 829     // typically echos the status code it received".
 830     // But never send certain codes, per section 7.4.1
 831     if (aCode == 1005 || aCode == 1006 || aCode == 1015) {
 832       CloseConnection(0, EmptyCString());
 833     } else {
 834       CloseConnection(aCode, aReason);
 835     }
 836   } else {
 837     // We initiated close, and server has replied: OnStop does rest of the work.
 838     MOZ_ASSERT(readyState == WebSocket::CLOSING, "unknown state");
 839   }
 840
 841   return NS_OK;
 842 }
 843
 844 //-----------------------------------------------------------------------------
 845 // WebSocketImpl::nsIInterfaceRequestor
 846 //-----------------------------------------------------------------------------
 847
 848 NS_IMETHODIMP
 849 WebSocketImpl::GetInterface(const nsIID& aIID, void** aResult)
 850 {
 851   AssertIsOnMainThread();
 852
 853   if (!mWebSocket || mWebSocket->ReadyState() == WebSocket::CLOSED) {
 854     return NS_ERROR_FAILURE;
 855   }
 856
 857   if (aIID.Equals(NS_GET_IID(nsIAuthPrompt)) ||
 858       aIID.Equals(NS_GET_IID(nsIAuthPrompt2))) {
 859     nsresult rv;
 860     nsIScriptContext* sc = mWebSocket->GetContextForEventHandlers(&rv);
 861     nsCOMPtr<nsIDocument> doc =
 862       nsContentUtils::GetDocumentFromScriptContext(sc);
 863     if (!doc) {
 864       return NS_ERROR_NOT_AVAILABLE;
 865     }
 866
 867     nsCOMPtr<nsIPromptFactory> wwatch =
 868       do_GetService(NS_WINDOWWATCHER_CONTRACTID, &rv);
 869     NS_ENSURE_SUCCESS(rv, rv);
 870
 871     nsCOMPtr<nsPIDOMWindow> outerWindow = doc->GetWindow();
 872     return wwatch->GetPrompt(outerWindow, aIID, aResult);
 873   }
 874
 875   return QueryInterface(aIID, aResult);
 876 }
 877
 878 ////////////////////////////////////////////////////////////////////////////////
 879 // WebSocket
 880 ////////////////////////////////////////////////////////////////////////////////
 881
 882 WebSocket::WebSocket(nsPIDOMWindow* aOwnerWindow)
 883   : DOMEventTargetHelper(aOwnerWindow)
 884   , mIsMainThread(true)
 885   , mKeepingAlive(false)
 886   , mCheckMustKeepAlive(true)
 887   , mOutgoingBufferedAmount(0)
 888   , mBinaryType(dom::BinaryType::Blob)
 889   , mMutex("WebSocketImpl::mMutex")
 890   , mReadyState(CONNECTING)
 891 {
 892   mImpl = new WebSocketImpl(this);
 893   mIsMainThread = mImpl->mIsMainThread;
 894 }
 895
 896 WebSocket::~WebSocket()
 897 {
 898 }
 899
 900 JSObject*
 901 WebSocket::WrapObject(JSContext* cx)
 902 {
 903   return WebSocketBinding::Wrap(cx, this);
 904 }
 905
 906 //---------------------------------------------------------------------------
 907 // WebIDL
 908 //---------------------------------------------------------------------------
 909
 910 // Constructor:
 911 already_AddRefed<WebSocket>
 912 WebSocket::Constructor(const GlobalObject& aGlobal,
 913                        const nsAString& aUrl,
 914                        ErrorResult& aRv)
 915 {
 916   Sequence<nsString> protocols;
 917   return WebSocket::Constructor(aGlobal, aUrl, protocols, aRv);
 918 }
 919
 920 already_AddRefed<WebSocket>
 921 WebSocket::Constructor(const GlobalObject& aGlobal,
 922                        const nsAString& aUrl,
 923                        const nsAString& aProtocol,
 924                        ErrorResult& aRv)
 925 {
 926   Sequence<nsString> protocols;
 927   protocols.AppendElement(aProtocol);
 928   return WebSocket::Constructor(aGlobal, aUrl, protocols, aRv);
 929 }
 930
 931 namespace {
 932
 933 // This class is used to clear any exception.
 934 class MOZ_STACK_CLASS ClearException
 935 {
 936 public:
 937   explicit ClearException(JSContext* aCx)
 938     : mCx(aCx)
 939   {
 940   }
 941
 942   ~ClearException()
 943   {
 944     JS_ClearPendingException(mCx);
 945   }
 946
 947 private:
 948   JSContext* mCx;
 949 };
 950
 951 class InitRunnable MOZ_FINAL : public WorkerMainThreadRunnable
 952 {
 953 public:
 954   InitRunnable(WebSocketImpl* aImpl, const nsAString& aURL,
 955                nsTArray<nsString>& aProtocolArray,
 956                const nsACString& aScriptFile, uint32_t aScriptLine,
 957                ErrorResult& aRv, bool* aConnectionFailed)
 958     : WorkerMainThreadRunnable(aImpl->mWorkerPrivate)
 959     , mImpl(aImpl)
 960     , mURL(aURL)
 961     , mProtocolArray(aProtocolArray)
 962     , mScriptFile(aScriptFile)
 963     , mScriptLine(aScriptLine)
 964     , mRv(aRv)
 965     , mConnectionFailed(aConnectionFailed)
 966   {
 967     MOZ_ASSERT(mWorkerPrivate);
 968     mWorkerPrivate->AssertIsOnWorkerThread();
 969   }
 970
 971   bool MainThreadRun() MOZ_OVERRIDE
 972   {
 973     AssertIsOnMainThread();
 974
 975     // Walk up to our containing page
 976     WorkerPrivate* wp = mWorkerPrivate;
 977     while (wp->GetParent()) {
 978       wp = wp->GetParent();
 979     }
 980
 981     nsPIDOMWindow* window = wp->GetWindow();
 982     if (window) {
 983       return InitWithWindow(window);
 984     }
 985
 986     return InitWindowless();
 987   }
 988
 989 private:
 990   bool InitWithWindow(nsPIDOMWindow* aWindow)
 991   {
 992     AutoJSAPI jsapi;
 993     if (NS_WARN_IF(!jsapi.Init(aWindow))) {
 994       mRv.Throw(NS_ERROR_FAILURE);
 995       return true;
 996     }
 997
 998     ClearException ce(jsapi.cx());
 999
1000     nsIDocument* doc = aWindow->GetExtantDoc();
1001     if (!doc) {
1002       mRv.Throw(NS_ERROR_FAILURE);
1003       return true;
1004     }
1005
1006     nsCOMPtr<nsIPrincipal> principal = doc->NodePrincipal();
1007     if (!principal) {
1008       mRv.Throw(NS_ERROR_FAILURE);
1009       return true;
1010     }
1011
1012     mImpl->Init(jsapi.cx(), principal, mURL, mProtocolArray, mScriptFile,
1013                 mScriptLine, mRv, mConnectionFailed);
1014     return true;
1015   }
1016
1017   bool InitWindowless()
1018   {
1019     MOZ_ASSERT(NS_IsMainThread());
1020
1021     WorkerPrivate* wp = mWorkerPrivate;
1022     while (wp->GetParent()) {
1023       wp = wp->GetParent();
1024     }
1025
1026     MOZ_ASSERT(!wp->GetWindow());
1027
1028     mImpl->Init(nullptr, wp->GetPrincipal(), mURL, mProtocolArray, mScriptFile,
1029                 mScriptLine, mRv, mConnectionFailed);
1030     return true;
1031   }
1032
1033   // Raw pointer. This worker runs synchronously.
1034   WebSocketImpl* mImpl;
1035
1036   const nsAString& mURL;
1037   nsTArray<nsString>& mProtocolArray;
1038   nsCString mScriptFile;
1039   uint32_t mScriptLine;
1040   ErrorResult& mRv;
1041   bool* mConnectionFailed;
1042 };
1043
1044 class AsyncOpenRunnable MOZ_FINAL : public WorkerMainThreadRunnable
1045 {
1046 public:
1047   AsyncOpenRunnable(WebSocketImpl* aImpl, ErrorResult& aRv)
1048     : WorkerMainThreadRunnable(aImpl->mWorkerPrivate)
1049     , mImpl(aImpl)
1050     , mRv(aRv)
1051   {
1052     MOZ_ASSERT(mWorkerPrivate);
1053     mWorkerPrivate->AssertIsOnWorkerThread();
1054   }
1055
1056   bool MainThreadRun() MOZ_OVERRIDE
1057   {
1058     AssertIsOnMainThread();
1059     mImpl->AsyncOpen(mRv);
1060     return true;
1061   }
1062
1063 private:
1064   // Raw pointer. This worker runs synchronously.
1065   WebSocketImpl* mImpl;
1066
1067   ErrorResult& mRv;
1068 };
1069
1070 } // anonymous namespace
1071
1072 already_AddRefed<WebSocket>
1073 WebSocket::Constructor(const GlobalObject& aGlobal,
1074                        const nsAString& aUrl,
1075                        const Sequence<nsString>& aProtocols,
1076                        ErrorResult& aRv)
1077 {
1078   nsCOMPtr<nsIPrincipal> principal;
1079   nsCOMPtr<nsPIDOMWindow> ownerWindow;
1080
1081   if (NS_IsMainThread()) {
1082     nsCOMPtr<nsIScriptObjectPrincipal> scriptPrincipal =
1083       do_QueryInterface(aGlobal.GetAsSupports());
1084     if (!scriptPrincipal) {
1085       aRv.Throw(NS_ERROR_FAILURE);
1086       return nullptr;
1087     }
1088
1089     principal = scriptPrincipal->GetPrincipal();
1090     if (!principal) {
1091       aRv.Throw(NS_ERROR_FAILURE);
1092       return nullptr;
1093     }
1094
1095     nsCOMPtr<nsIScriptGlobalObject> sgo =
1096       do_QueryInterface(aGlobal.GetAsSupports());
1097     if (!sgo) {
1098       aRv.Throw(NS_ERROR_FAILURE);
1099       return nullptr;
1100     }
1101
1102     ownerWindow = do_QueryInterface(aGlobal.GetAsSupports());
1103     if (!ownerWindow) {
1104       aRv.Throw(NS_ERROR_FAILURE);
1105       return nullptr;
1106     }
1107   }
1108
1109   nsTArray<nsString> protocolArray;
1110
1111   for (uint32_t index = 0, len = aProtocols.Length(); index < len; ++index) {
1112
1113     const nsString& protocolElement = aProtocols[index];
1114
1115     if (protocolElement.IsEmpty()) {
1116       aRv.Throw(NS_ERROR_DOM_SYNTAX_ERR);
1117       return nullptr;
1118     }
1119     if (protocolArray.Contains(protocolElement)) {
1120       aRv.Throw(NS_ERROR_DOM_SYNTAX_ERR);
1121       return nullptr;
1122     }
1123     if (protocolElement.FindChar(',') != -1)  /* interferes w/list */ {
1124       aRv.Throw(NS_ERROR_DOM_SYNTAX_ERR);
1125       return nullptr;
1126     }
1127
1128     protocolArray.AppendElement(protocolElement);
1129   }
1130
1131   nsRefPtr<WebSocket> webSocket = new WebSocket(ownerWindow);
1132   nsRefPtr<WebSocketImpl> kungfuDeathGrip = webSocket->mImpl;
1133
1134   bool connectionFailed = true;
1135
1136   if (NS_IsMainThread()) {
1137     webSocket->mImpl->Init(aGlobal.Context(), principal, aUrl, protocolArray,
1138                            EmptyCString(), 0, aRv, &connectionFailed);
1139   } else {
1140     // In workers we have to keep the worker alive using a feature in order to
1141     // dispatch messages correctly.
1142     if (!webSocket->mImpl->RegisterFeature()) {
1143       aRv.Throw(NS_ERROR_FAILURE);
1144       return nullptr;
1145     }
1146
1147     unsigned lineno;
1148     JS::AutoFilename file;
1149     if (!JS::DescribeScriptedCaller(aGlobal.Context(), &file, &lineno)) {
1150       NS_WARNING("Failed to get line number and filename in workers.");
1151     }
1152
1153     nsRefPtr<InitRunnable> runnable =
1154       new InitRunnable(webSocket->mImpl, aUrl, protocolArray,
1155                        nsAutoCString(file.get()), lineno, aRv,
1156                        &connectionFailed);
1157     runnable->Dispatch(aGlobal.Context());
1158   }
1159
1160   if (NS_WARN_IF(aRv.Failed())) {
1161     return nullptr;
1162   }
1163
1164   // It can be that we have been already disconnected because the WebSocket is
1165   // gone away while we where initializing the webSocket.
1166   if (!webSocket->mImpl) {
1167     aRv.Throw(NS_ERROR_FAILURE);
1168     return nullptr;
1169   }
1170
1171   // We don't return an error if the connection just failed. Instead we dispatch
1172   // an event.
1173   if (connectionFailed) {
1174     webSocket->mImpl->FailConnection(nsIWebSocketChannel::CLOSE_ABNORMAL);
1175   }
1176
1177   // If we don't have a channel, the connection is failed and onerror() will be
1178   // called asynchrounsly.
1179   if (!webSocket->mImpl->mChannel) {
1180     return webSocket.forget();
1181   }
1182
1183   class MOZ_STACK_CLASS ClearWebSocket
1184   {
1185   public:
1186     explicit ClearWebSocket(WebSocketImpl* aWebSocketImpl)
1187       : mWebSocketImpl(aWebSocketImpl)
1188       , mDone(false)
1189     {
1190     }
1191
1192     void Done()
1193     {
1194        mDone = true;
1195     }
1196
1197     ~ClearWebSocket()
1198     {
1199       if (!mDone) {
1200         mWebSocketImpl->mChannel = nullptr;
1201         mWebSocketImpl->FailConnection(nsIWebSocketChannel::CLOSE_ABNORMAL);
1202       }
1203     }
1204
1205     WebSocketImpl* mWebSocketImpl;
1206     bool mDone;
1207   };
1208
1209   ClearWebSocket cws(webSocket->mImpl);
1210
1211   // This operation must be done on the correct thread. The rest must run on the
1212   // main-thread.
1213   aRv = webSocket->mImpl->mChannel->SetNotificationCallbacks(webSocket->mImpl);
1214   if (NS_WARN_IF(aRv.Failed())) {
1215     return nullptr;
1216   }
1217
1218   if (NS_IsMainThread()) {
1219     webSocket->mImpl->AsyncOpen(aRv);
1220   } else {
1221     nsRefPtr<AsyncOpenRunnable> runnable =
1222       new AsyncOpenRunnable(webSocket->mImpl, aRv);
1223     runnable->Dispatch(aGlobal.Context());
1224   }
1225
1226   if (NS_WARN_IF(aRv.Failed())) {
1227     return nullptr;
1228   }
1229
1230   // It can be that we have been already disconnected because the WebSocket is
1231   // gone away while we where initializing the webSocket.
1232   if (!webSocket->mImpl) {
1233     aRv.Throw(NS_ERROR_FAILURE);
1234     return nullptr;
1235   }
1236
1237   cws.Done();
1238   return webSocket.forget();
1239 }
1240
1241 NS_IMPL_CYCLE_COLLECTION_CLASS(WebSocket)
1242
1243 NS_IMPL_CYCLE_COLLECTION_CAN_SKIP_BEGIN(WebSocket)
1244   bool isBlack = tmp->IsBlack();
1245   if (isBlack || tmp->mKeepingAlive) {
1246     if (tmp->mListenerManager) {
1247       tmp->mListenerManager->MarkForCC();
1248     }
1249     if (!isBlack && tmp->PreservingWrapper()) {
1250       // This marks the wrapper black.
1251       tmp->GetWrapper();
1252     }
1253     return true;
1254   }
1255 NS_IMPL_CYCLE_COLLECTION_CAN_SKIP_END
1256
1257 NS_IMPL_CYCLE_COLLECTION_CAN_SKIP_IN_CC_BEGIN(WebSocket)
1258   return tmp->IsBlack();
1259 NS_IMPL_CYCLE_COLLECTION_CAN_SKIP_IN_CC_END
1260
1261 NS_IMPL_CYCLE_COLLECTION_CAN_SKIP_THIS_BEGIN(WebSocket)
1262   return tmp->IsBlack();
1263 NS_IMPL_CYCLE_COLLECTION_CAN_SKIP_THIS_END
1264
1265 NS_IMPL_CYCLE_COLLECTION_TRACE_BEGIN_INHERITED(WebSocket,
1266                                                DOMEventTargetHelper)
1267 NS_IMPL_CYCLE_COLLECTION_TRACE_END
1268
1269 NS_IMPL_CYCLE_COLLECTION_TRAVERSE_BEGIN_INHERITED(WebSocket,
1270                                                   DOMEventTargetHelper)
1271   if (tmp->mImpl) {
1272     NS_IMPL_CYCLE_COLLECTION_TRAVERSE(mImpl->mPrincipal)
1273     NS_IMPL_CYCLE_COLLECTION_TRAVERSE(mImpl->mChannel)
1274   }
1275 NS_IMPL_CYCLE_COLLECTION_TRAVERSE_END
1276
1277 NS_IMPL_CYCLE_COLLECTION_UNLINK_BEGIN_INHERITED(WebSocket,
1278                                                 DOMEventTargetHelper)
1279   if (tmp->mImpl) {
1280     NS_IMPL_CYCLE_COLLECTION_UNLINK(mImpl->mPrincipal)
1281     NS_IMPL_CYCLE_COLLECTION_UNLINK(mImpl->mChannel)
1282     tmp->mImpl->Disconnect();
1283     MOZ_ASSERT(!tmp->mImpl);
1284   }
1285 NS_IMPL_CYCLE_COLLECTION_UNLINK_END
1286
1287 NS_INTERFACE_MAP_BEGIN_CYCLE_COLLECTION_INHERITED(WebSocket)
1288 NS_INTERFACE_MAP_END_INHERITING(DOMEventTargetHelper)
1289
1290 NS_IMPL_ADDREF_INHERITED(WebSocket, DOMEventTargetHelper)
1291 NS_IMPL_RELEASE_INHERITED(WebSocket, DOMEventTargetHelper)
1292
1293 void
1294 WebSocket::DisconnectFromOwner()
1295 {
1296   AssertIsOnMainThread();
1297   DOMEventTargetHelper::DisconnectFromOwner();
1298
1299   if (mImpl) {
1300     mImpl->CloseConnection(nsIWebSocketChannel::CLOSE_GOING_AWAY);
1301   }
1302
1303   DontKeepAliveAnyMore();
1304 }
1305
1306 //-----------------------------------------------------------------------------
1307 // WebSocketImpl:: initialization
1308 //-----------------------------------------------------------------------------
1309
1310 void
1311 WebSocketImpl::Init(JSContext* aCx,
1312                     nsIPrincipal* aPrincipal,
1313                     const nsAString& aURL,
1314                     nsTArray<nsString>& aProtocolArray,
1315                     const nsACString& aScriptFile,
1316                     uint32_t aScriptLine,
1317                     ErrorResult& aRv,
1318                     bool* aConnectionFailed)
1319 {
1320   AssertIsOnMainThread();
1321   MOZ_ASSERT(aPrincipal);
1322
1323   // We need to keep the implementation alive in case the init disconnects it
1324   // because of some error.
1325   nsRefPtr<WebSocketImpl> kungfuDeathGrip = this;
1326
1327   mPrincipal = aPrincipal;
1328
1329   // Attempt to kill "ghost" websocket: but usually too early for check to fail
1330   aRv = mWebSocket->CheckInnerWindowCorrectness();
1331   if (NS_WARN_IF(aRv.Failed())) {
1332     return;
1333   }
1334
1335   // Shut down websocket if window is frozen or destroyed (only needed for
1336   // "ghost" websockets--see bug 696085)
1337   if (!mWorkerPrivate) {
1338     nsCOMPtr<nsIObserverService> os = mozilla::services::GetObserverService();
1339     if (NS_WARN_IF(!os)) {
1340       aRv.Throw(NS_ERROR_FAILURE);
1341       return;
1342     }
1343
1344     aRv = os->AddObserver(this, DOM_WINDOW_DESTROYED_TOPIC, true);
1345     if (NS_WARN_IF(aRv.Failed())) {
1346       return;
1347     }
1348
1349     aRv = os->AddObserver(this, DOM_WINDOW_FROZEN_TOPIC, true);
1350     if (NS_WARN_IF(aRv.Failed())) {
1351       return;
1352     }
1353   }
1354
1355   if (mWorkerPrivate) {
1356     mScriptFile = aScriptFile;
1357     mScriptLine = aScriptLine;
1358   } else {
1359     MOZ_ASSERT(aCx);
1360
1361     unsigned lineno;
1362     JS::AutoFilename file;
1363     if (JS::DescribeScriptedCaller(aCx, &file, &lineno)) {
1364       mScriptFile = file.get();
1365       mScriptLine = lineno;
1366     }
1367   }
1368
1369   // If we don't have aCx, we are window-less, so we don't have a
1370   // inner-windowID. This can happen in sharedWorkers and ServiceWorkers or in
1371   // DedicateWorkers created by JSM.
1372   if (aCx) {
1373     mInnerWindowID = nsJSUtils::GetCurrentlyRunningCodeInnerWindowID(aCx);
1374   }
1375
1376   // parses the url
1377   aRv = ParseURL(PromiseFlatString(aURL));
1378   if (NS_WARN_IF(aRv.Failed())) {
1379     return;
1380   }
1381
1382   nsIScriptContext* sc = nullptr;
1383   {
1384     nsresult rv;
1385     sc = mWebSocket->GetContextForEventHandlers(&rv);
1386     if (NS_WARN_IF(NS_FAILED(rv))) {
1387       aRv.Throw(rv);
1388       return;
1389     }
1390   }
1391
1392   // Don't allow https:// to open ws://
1393   if (!mSecure &&
1394       !Preferences::GetBool("network.websocket.allowInsecureFromHTTPS",
1395                             false)) {
1396     // Confirmed we are opening plain ws:// and want to prevent this from a
1397     // secure context (e.g. https). Check the principal's uri to determine if
1398     // we were loaded from https.
1399     nsCOMPtr<nsIGlobalObject> globalObject(GetEntryGlobal());
1400     if (globalObject) {
1401       nsCOMPtr<nsIPrincipal> principal(globalObject->PrincipalOrNull());
1402       if (principal) {
1403         nsCOMPtr<nsIURI> uri;
1404         principal->GetURI(getter_AddRefs(uri));
1405         if (uri) {
1406           bool originIsHttps = false;
1407           aRv = uri->SchemeIs("https", &originIsHttps);
1408           if (NS_WARN_IF(aRv.Failed())) {
1409             return;
1410           }
1411
1412           if (originIsHttps) {
1413             aRv.Throw(NS_ERROR_DOM_SECURITY_ERR);
1414             return;
1415           }
1416         }
1417       }
1418     }
1419   }
1420
1421   // Assign the sub protocol list and scan it for illegal values
1422   for (uint32_t index = 0; index < aProtocolArray.Length(); ++index) {
1423     for (uint32_t i = 0; i < aProtocolArray[index].Length(); ++i) {
1424       if (aProtocolArray[index][i] < static_cast<char16_t>(0x0021) ||
1425           aProtocolArray[index][i] > static_cast<char16_t>(0x007E)) {
1426         aRv.Throw(NS_ERROR_DOM_SYNTAX_ERR);
1427         return;
1428       }
1429     }
1430
1431     if (!mRequestedProtocolList.IsEmpty()) {
1432       mRequestedProtocolList.AppendLiteral(", ");
1433     }
1434
1435     AppendUTF16toUTF8(aProtocolArray[index], mRequestedProtocolList);
1436   }
1437
1438   nsCOMPtr<nsIURI> uri;
1439   {
1440     nsresult rv = NS_NewURI(getter_AddRefs(uri), mURI);
1441
1442     // We crash here because we are sure that mURI is a valid URI, so either we
1443     // are OOM'ing or something else bad is happening.
1444     if (NS_FAILED(rv)) {
1445       MOZ_CRASH();
1446     }
1447   }
1448
1449   // Check content policy.
1450   int16_t shouldLoad = nsIContentPolicy::ACCEPT;
1451   nsCOMPtr<nsIDocument> originDoc = nsContentUtils::GetDocumentFromScriptContext(sc);
1452   mOriginDocument = do_GetWeakReference(originDoc);
1453   aRv = NS_CheckContentLoadPolicy(nsIContentPolicy::TYPE_WEBSOCKET,
1454                                   uri,
1455                                   mPrincipal,
1456                                   originDoc,
1457                                   EmptyCString(),
1458                                   nullptr,
1459                                   &shouldLoad,
1460                                   nsContentUtils::GetContentPolicy(),
1461                                   nsContentUtils::GetSecurityManager());
1462   if (NS_WARN_IF(aRv.Failed())) {
1463     return;
1464   }
1465
1466   if (NS_CP_REJECTED(shouldLoad)) {
1467     // Disallowed by content policy.
1468     aRv.Throw(NS_ERROR_CONTENT_BLOCKED);
1469     return;
1470   }
1471
1472   // the constructor should throw a SYNTAX_ERROR only if it fails to parse the
1473   // url parameter, so don't throw if InitializeConnection fails, and call
1474   // onerror/onclose asynchronously
1475   if (NS_FAILED(InitializeConnection())) {
1476     *aConnectionFailed = true;
1477   } else {
1478     *aConnectionFailed = false;
1479   }
1480 }
1481
1482 void
1483 WebSocketImpl::AsyncOpen(ErrorResult& aRv)
1484 {
1485   NS_ABORT_IF_FALSE(NS_IsMainThread(), "Not running on main thread");
1486
1487   nsCString asciiOrigin;
1488   aRv = nsContentUtils::GetASCIIOrigin(mPrincipal, asciiOrigin);
1489   if (NS_WARN_IF(aRv.Failed())) {
1490     return;
1491   }
1492
1493   ToLowerCase(asciiOrigin);
1494
1495   nsCOMPtr<nsIURI> uri;
1496   aRv = NS_NewURI(getter_AddRefs(uri), mURI);
1497   MOZ_ASSERT(!aRv.Failed());
1498
1499   aRv = mChannel->AsyncOpen(uri, asciiOrigin, this, nullptr);
1500   if (NS_WARN_IF(aRv.Failed())) {
1501     return;
1502   }
1503 }
1504
1505 //-----------------------------------------------------------------------------
1506 // WebSocketImpl methods:
1507 //-----------------------------------------------------------------------------
1508
1509 class nsAutoCloseWS MOZ_FINAL
1510 {
1511 public:
1512   explicit nsAutoCloseWS(WebSocketImpl* aWebSocketImpl)
1513     : mWebSocketImpl(aWebSocketImpl)
1514   {}
1515
1516   ~nsAutoCloseWS()
1517   {
1518     if (!mWebSocketImpl->mChannel) {
1519       mWebSocketImpl->CloseConnection(nsIWebSocketChannel::CLOSE_INTERNAL_ERROR);
1520     }
1521   }
1522 private:
1523   nsRefPtr<WebSocketImpl> mWebSocketImpl;
1524 };
1525
1526 nsresult
1527 WebSocketImpl::InitializeConnection()
1528 {
1529   AssertIsOnMainThread();
1530   NS_ABORT_IF_FALSE(!mChannel, "mChannel should be null");
1531
1532   nsCOMPtr<nsIWebSocketChannel> wsChannel;
1533   nsAutoCloseWS autoClose(this);
1534   nsresult rv;
1535
1536   if (mSecure) {
1537     wsChannel =
1538       do_CreateInstance("@mozilla.org/network/protocol;1?name=wss", &rv);
1539   } else {
1540     wsChannel =
1541       do_CreateInstance("@mozilla.org/network/protocol;1?name=ws", &rv);
1542   }
1543   NS_ENSURE_SUCCESS(rv, rv);
1544
1545   // add ourselves to the document's load group and
1546   // provide the http stack the loadgroup info too
1547   nsCOMPtr<nsILoadGroup> loadGroup;
1548   rv = GetLoadGroup(getter_AddRefs(loadGroup));
1549   if (loadGroup) {
1550     rv = wsChannel->SetLoadGroup(loadGroup);
1551     NS_ENSURE_SUCCESS(rv, rv);
1552     rv = loadGroup->AddRequest(this, nullptr);
1553     NS_ENSURE_SUCCESS(rv, rv);
1554
1555     mWeakLoadGroup = do_GetWeakReference(loadGroup);
1556   }
1557
1558   // manually adding loadinfo to the channel since it
1559   // was not set during channel creation.
1560   nsCOMPtr<nsIDocument> doc = do_QueryReferent(mOriginDocument);
1561
1562   // mOriginDocument has to be release on main-thread because WeakReferences
1563   // are not thread-safe.
1564   mOriginDocument = nullptr;
1565
1566   nsCOMPtr<nsILoadInfo> loadInfo =
1567     new LoadInfo(doc ?
1568                    doc->NodePrincipal() : mPrincipal.get(),
1569                  mPrincipal,
1570                  doc,
1571                  nsILoadInfo::SEC_NORMAL,
1572                  nsIContentPolicy::TYPE_WEBSOCKET);
1573   rv = wsChannel->SetLoadInfo(loadInfo);
1574   NS_ENSURE_SUCCESS(rv, rv);
1575
1576   if (!mRequestedProtocolList.IsEmpty()) {
1577     rv = wsChannel->SetProtocol(mRequestedProtocolList);
1578     NS_ENSURE_SUCCESS(rv, rv);
1579   }
1580
1581   nsCOMPtr<nsIThreadRetargetableRequest> rr = do_QueryInterface(wsChannel);
1582   NS_ENSURE_TRUE(rr, NS_ERROR_FAILURE);
1583
1584   rv = rr->RetargetDeliveryTo(this);
1585   NS_ENSURE_SUCCESS(rv, rv);
1586
1587   mChannel = wsChannel;
1588
1589   return NS_OK;
1590 }
1591
1592 void
1593 WebSocketImpl::DispatchConnectionCloseEvents()
1594 {
1595   AssertIsOnTargetThread();
1596
1597   if (mDisconnectingOrDisconnected) {
1598     return;
1599   }
1600
1601   mWebSocket->SetReadyState(WebSocket::CLOSED);
1602
1603   // Let's keep the object alive because the webSocket can be CCed in the
1604   // onerror or in the onclose callback.
1605   nsRefPtr<WebSocket> webSocket = mWebSocket;
1606
1607   // Call 'onerror' if needed
1608   if (mFailed) {
1609     nsresult rv =
1610       webSocket->CreateAndDispatchSimpleEvent(NS_LITERAL_STRING("error"));
1611     if (NS_FAILED(rv)) {
1612       NS_WARNING("Failed to dispatch the error event");
1613     }
1614   }
1615
1616   nsresult rv = webSocket->CreateAndDispatchCloseEvent(mCloseEventWasClean,
1617                                                        mCloseEventCode,
1618                                                        mCloseEventReason);
1619   if (NS_FAILED(rv)) {
1620     NS_WARNING("Failed to dispatch the close event");
1621   }
1622
1623   webSocket->UpdateMustKeepAlive();
1624   Disconnect();
1625 }
1626
1627 nsresult
1628 WebSocket::CreateAndDispatchSimpleEvent(const nsAString& aName)
1629 {
1630   MOZ_ASSERT(mImpl);
1631   AssertIsOnTargetThread();
1632
1633   nsresult rv = CheckInnerWindowCorrectness();
1634   if (NS_FAILED(rv)) {
1635     return NS_OK;
1636   }
1637
1638   nsCOMPtr<nsIDOMEvent> event;
1639   rv = NS_NewDOMEvent(getter_AddRefs(event), this, nullptr, nullptr);
1640   NS_ENSURE_SUCCESS(rv, rv);
1641
1642   // it doesn't bubble, and it isn't cancelable
1643   rv = event->InitEvent(aName, false, false);
1644   NS_ENSURE_SUCCESS(rv, rv);
1645
1646   event->SetTrusted(true);
1647
1648   return DispatchDOMEvent(nullptr, event, nullptr, nullptr);
1649 }
1650
1651 nsresult
1652 WebSocket::CreateAndDispatchMessageEvent(const nsACString& aData,
1653                                          bool aIsBinary)
1654 {
1655   MOZ_ASSERT(mImpl);
1656   AssertIsOnTargetThread();
1657
1658   AutoJSAPI jsapi;
1659
1660   if (NS_IsMainThread()) {
1661     if (NS_WARN_IF(!jsapi.Init(GetOwner()))) {
1662       return NS_ERROR_FAILURE;
1663     }
1664   } else {
1665     MOZ_ASSERT(!mIsMainThread);
1666     MOZ_ASSERT(mImpl->mWorkerPrivate);
1667     if (NS_WARN_IF(!jsapi.Init(mImpl->mWorkerPrivate->GlobalScope()))) {
1668       return NS_ERROR_FAILURE;
1669     }
1670   }
1671
1672   return CreateAndDispatchMessageEvent(jsapi.cx(), aData, aIsBinary);
1673 }
1674
1675 nsresult
1676 WebSocket::CreateAndDispatchMessageEvent(JSContext* aCx,
1677                                          const nsACString& aData,
1678                                          bool aIsBinary)
1679 {
1680   MOZ_ASSERT(mImpl);
1681   AssertIsOnTargetThread();
1682
1683   nsresult rv = CheckInnerWindowCorrectness();
1684   if (NS_FAILED(rv)) {
1685     return NS_OK;
1686   }
1687
1688   // Create appropriate JS object for message
1689   JS::Rooted<JS::Value> jsData(aCx);
1690   if (aIsBinary) {
1691     if (mBinaryType == dom::BinaryType::Blob) {
1692       nsresult rv = nsContentUtils::CreateBlobBuffer(aCx, GetOwner(), aData,
1693                                                      &jsData);
1694       NS_ENSURE_SUCCESS(rv, rv);
1695     } else if (mBinaryType == dom::BinaryType::Arraybuffer) {
1696       JS::Rooted<JSObject*> arrayBuf(aCx);
1697       nsresult rv = nsContentUtils::CreateArrayBuffer(aCx, aData,
1698                                                       arrayBuf.address());
1699       NS_ENSURE_SUCCESS(rv, rv);
1700       jsData = OBJECT_TO_JSVAL(arrayBuf);
1701     } else {
1702       NS_RUNTIMEABORT("Unknown binary type!");
1703       return NS_ERROR_UNEXPECTED;
1704     }
1705   } else {
1706     // JS string
1707     NS_ConvertUTF8toUTF16 utf16Data(aData);
1708     JSString* jsString;
1709     jsString = JS_NewUCStringCopyN(aCx, utf16Data.get(), utf16Data.Length());
1710     NS_ENSURE_TRUE(jsString, NS_ERROR_FAILURE);
1711
1712     jsData = STRING_TO_JSVAL(jsString);
1713   }
1714
1715   // create an event that uses the MessageEvent interface,
1716   // which does not bubble, is not cancelable, and has no default action
1717
1718   nsCOMPtr<nsIDOMEvent> event;
1719   rv = NS_NewDOMMessageEvent(getter_AddRefs(event), this, nullptr, nullptr);
1720   NS_ENSURE_SUCCESS(rv, rv);
1721
1722   nsCOMPtr<nsIDOMMessageEvent> messageEvent = do_QueryInterface(event);
1723   rv = messageEvent->InitMessageEvent(NS_LITERAL_STRING("message"),
1724                                       false, false,
1725                                       jsData,
1726                                       mImpl->mUTF16Origin,
1727                                       EmptyString(), nullptr);
1728   NS_ENSURE_SUCCESS(rv, rv);
1729
1730   event->SetTrusted(true);
1731
1732   return DispatchDOMEvent(nullptr, event, nullptr, nullptr);
1733 }
1734
1735 nsresult
1736 WebSocket::CreateAndDispatchCloseEvent(bool aWasClean,
1737                                        uint16_t aCode,
1738                                        const nsAString &aReason)
1739 {
1740   MOZ_ASSERT(mImpl);
1741   AssertIsOnTargetThread();
1742
1743   nsresult rv = CheckInnerWindowCorrectness();
1744   if (NS_FAILED(rv)) {
1745     return NS_OK;
1746   }
1747
1748   CloseEventInit init;
1749   init.mBubbles = false;
1750   init.mCancelable = false;
1751   init.mWasClean = aWasClean;
1752   init.mCode = aCode;
1753   init.mReason = aReason;
1754
1755   nsRefPtr<CloseEvent> event =
1756     CloseEvent::Constructor(this, NS_LITERAL_STRING("close"), init);
1757   event->SetTrusted(true);
1758
1759   return DispatchDOMEvent(nullptr, event, nullptr, nullptr);
1760 }
1761
1762 namespace {
1763
1764 class PrefEnabledRunnable MOZ_FINAL : public WorkerMainThreadRunnable
1765 {
1766 public:
1767   explicit PrefEnabledRunnable(WorkerPrivate* aWorkerPrivate)
1768     : WorkerMainThreadRunnable(aWorkerPrivate)
1769     , mEnabled(false)
1770   { }
1771
1772   bool MainThreadRun() MOZ_OVERRIDE
1773   {
1774     AssertIsOnMainThread();
1775     mEnabled = Preferences::GetBool("dom.workers.websocket.enabled", false);
1776     return true;
1777   }
1778
1779   bool IsEnabled() const
1780   {
1781     return mEnabled;
1782   }
1783
1784 private:
1785   bool mEnabled;
1786 };
1787
1788 } // anonymous namespace
1789
1790 bool
1791 WebSocket::PrefEnabled(JSContext* /* aCx */, JSObject* /* aGlobal */)
1792 {
1793   // WebSockets are always enabled on main-thread.
1794   if (NS_IsMainThread()) {
1795     return true;
1796   }
1797
1798   WorkerPrivate* workerPrivate = GetCurrentThreadWorkerPrivate();
1799   MOZ_ASSERT(workerPrivate);
1800   workerPrivate->AssertIsOnWorkerThread();
1801
1802   nsRefPtr<PrefEnabledRunnable> runnable =
1803     new PrefEnabledRunnable(workerPrivate);
1804   runnable->Dispatch(workerPrivate->GetJSContext());
1805
1806   return runnable->IsEnabled();
1807 }
1808
1809 nsresult
1810 WebSocketImpl::ParseURL(const nsAString& aURL)
1811 {
1812   AssertIsOnMainThread();
1813   NS_ENSURE_TRUE(!aURL.IsEmpty(), NS_ERROR_DOM_SYNTAX_ERR);
1814
1815   nsCOMPtr<nsIURI> uri;
1816   nsresult rv = NS_NewURI(getter_AddRefs(uri), aURL);
1817   NS_ENSURE_SUCCESS(rv, NS_ERROR_DOM_SYNTAX_ERR);
1818
1819   nsCOMPtr<nsIURL> parsedURL = do_QueryInterface(uri, &rv);
1820   NS_ENSURE_SUCCESS(rv, NS_ERROR_DOM_SYNTAX_ERR);
1821
1822   nsAutoCString fragment;
1823   rv = parsedURL->GetRef(fragment);
1824   NS_ENSURE_TRUE(NS_SUCCEEDED(rv) && fragment.IsEmpty(),
1825                  NS_ERROR_DOM_SYNTAX_ERR);
1826
1827   nsAutoCString scheme;
1828   rv = parsedURL->GetScheme(scheme);
1829   NS_ENSURE_TRUE(NS_SUCCEEDED(rv) && !scheme.IsEmpty(),
1830                  NS_ERROR_DOM_SYNTAX_ERR);
1831
1832   nsAutoCString host;
1833   rv = parsedURL->GetAsciiHost(host);
1834   NS_ENSURE_TRUE(NS_SUCCEEDED(rv) && !host.IsEmpty(), NS_ERROR_DOM_SYNTAX_ERR);
1835
1836   int32_t port;
1837   rv = parsedURL->GetPort(&port);
1838   NS_ENSURE_SUCCESS(rv, NS_ERROR_DOM_SYNTAX_ERR);
1839
1840   rv = NS_CheckPortSafety(port, scheme.get());
1841   NS_ENSURE_SUCCESS(rv, NS_ERROR_DOM_SYNTAX_ERR);
1842
1843   nsAutoCString filePath;
1844   rv = parsedURL->GetFilePath(filePath);
1845   if (filePath.IsEmpty()) {
1846     filePath.Assign('/');
1847   }
1848   NS_ENSURE_SUCCESS(rv, NS_ERROR_DOM_SYNTAX_ERR);
1849
1850   nsAutoCString query;
1851   rv = parsedURL->GetQuery(query);
1852   NS_ENSURE_SUCCESS(rv, NS_ERROR_DOM_SYNTAX_ERR);
1853
1854   if (scheme.LowerCaseEqualsLiteral("ws")) {
1855      mSecure = false;
1856      mPort = (port == -1) ? DEFAULT_WS_SCHEME_PORT : port;
1857   } else if (scheme.LowerCaseEqualsLiteral("wss")) {
1858     mSecure = true;
1859     mPort = (port == -1) ? DEFAULT_WSS_SCHEME_PORT : port;
1860   } else {
1861     return NS_ERROR_DOM_SYNTAX_ERR;
1862   }
1863
1864   rv = nsContentUtils::GetUTFOrigin(parsedURL, mUTF16Origin);
1865   NS_ENSURE_SUCCESS(rv, NS_ERROR_DOM_SYNTAX_ERR);
1866
1867   mAsciiHost = host;
1868   ToLowerCase(mAsciiHost);
1869
1870   mResource = filePath;
1871   if (!query.IsEmpty()) {
1872     mResource.Append('?');
1873     mResource.Append(query);
1874   }
1875   uint32_t length = mResource.Length();
1876   uint32_t i;
1877   for (i = 0; i < length; ++i) {
1878     if (mResource[i] < static_cast<char16_t>(0x0021) ||
1879         mResource[i] > static_cast<char16_t>(0x007E)) {
1880       return NS_ERROR_DOM_SYNTAX_ERR;
1881     }
1882   }
1883
1884   mWebSocket->mOriginalURL = aURL;
1885
1886   rv = parsedURL->GetSpec(mURI);
1887   MOZ_ASSERT(NS_SUCCEEDED(rv));
1888
1889   return NS_OK;
1890 }
1891
1892 //-----------------------------------------------------------------------------
1893 // Methods that keep alive the WebSocket object when:
1894 //   1. the object has registered event listeners that can be triggered
1895 //      ("strong event listeners");
1896 //   2. there are outgoing not sent messages.
1897 //-----------------------------------------------------------------------------
1898
1899 void
1900 WebSocket::UpdateMustKeepAlive()
1901 {
1902   // Here we could not have mImpl.
1903   MOZ_ASSERT(NS_IsMainThread() == mIsMainThread);
1904
1905   if (!mCheckMustKeepAlive || !mImpl) {
1906     return;
1907   }
1908
1909   bool shouldKeepAlive = false;
1910   uint16_t readyState = ReadyState();
1911
1912   if (mListenerManager) {
1913     switch (readyState)
1914     {
1915       case CONNECTING:
1916       {
1917         if (mListenerManager->HasListenersFor(nsGkAtoms::onopen) ||
1918             mListenerManager->HasListenersFor(nsGkAtoms::onmessage) ||
1919             mListenerManager->HasListenersFor(nsGkAtoms::onerror) ||
1920             mListenerManager->HasListenersFor(nsGkAtoms::onclose)) {
1921           shouldKeepAlive = true;
1922         }
1923       }
1924       break;
1925
1926       case OPEN:
1927       case CLOSING:
1928       {
1929         if (mListenerManager->HasListenersFor(nsGkAtoms::onmessage) ||
1930             mListenerManager->HasListenersFor(nsGkAtoms::onerror) ||
1931             mListenerManager->HasListenersFor(nsGkAtoms::onclose) ||
1932             mOutgoingBufferedAmount != 0) {
1933           shouldKeepAlive = true;
1934         }
1935       }
1936       break;
1937
1938       case CLOSED:
1939       {
1940         shouldKeepAlive = false;
1941       }
1942     }
1943   }
1944
1945   if (mKeepingAlive && !shouldKeepAlive) {
1946     mKeepingAlive = false;
1947     mImpl->ReleaseObject();
1948   } else if (!mKeepingAlive && shouldKeepAlive) {
1949     mKeepingAlive = true;
1950     mImpl->AddRefObject();
1951   }
1952 }
1953
1954 void
1955 WebSocket::DontKeepAliveAnyMore()
1956 {
1957   // Here we could not have mImpl.
1958   MOZ_ASSERT(NS_IsMainThread() == mIsMainThread);
1959
1960   if (mKeepingAlive) {
1961     MOZ_ASSERT(mImpl);
1962
1963     mKeepingAlive = false;
1964     mImpl->ReleaseObject();
1965   }
1966
1967   mCheckMustKeepAlive = false;
1968 }
1969
1970 namespace {
1971
1972 class WebSocketWorkerFeature MOZ_FINAL : public WorkerFeature
1973 {
1974 public:
1975   explicit WebSocketWorkerFeature(WebSocketImpl* aWebSocketImpl)
1976     : mWebSocketImpl(aWebSocketImpl)
1977   {
1978   }
1979
1980   bool Notify(JSContext* aCx, Status aStatus) MOZ_OVERRIDE
1981   {
1982     MOZ_ASSERT(aStatus > workers::Running);
1983
1984     if (aStatus >= Canceling) {
1985       mWebSocketImpl->mWorkerShuttingDown = true;
1986       mWebSocketImpl->CloseConnection(nsIWebSocketChannel::CLOSE_GOING_AWAY);
1987     }
1988
1989     return true;
1990   }
1991
1992   bool Suspend(JSContext* aCx) MOZ_OVERRIDE
1993   {
1994     mWebSocketImpl->mWorkerShuttingDown = true;
1995     mWebSocketImpl->CloseConnection(nsIWebSocketChannel::CLOSE_GOING_AWAY);
1996     return true;
1997   }
1998
1999 private:
2000   WebSocketImpl* mWebSocketImpl;
2001 };
2002
2003 } // anonymous namespace
2004
2005 void
2006 WebSocketImpl::AddRefObject()
2007 {
2008   AssertIsOnTargetThread();
2009   AddRef();
2010 }
2011
2012 void
2013 WebSocketImpl::ReleaseObject()
2014 {
2015   AssertIsOnTargetThread();
2016   Release();
2017 }
2018
2019 bool
2020 WebSocketImpl::RegisterFeature()
2021 {
2022   mWorkerPrivate->AssertIsOnWorkerThread();
2023   MOZ_ASSERT(!mWorkerFeature);
2024   mWorkerFeature = new WebSocketWorkerFeature(this);
2025
2026   JSContext* cx = GetCurrentThreadJSContext();
2027   if (!mWorkerPrivate->AddFeature(cx, mWorkerFeature)) {
2028     NS_WARNING("Failed to register a feature.");
2029     mWorkerFeature = nullptr;
2030     return false;
2031   }
2032
2033 #ifdef DEBUG
2034   SetHasFeatureRegistered(true);
2035 #endif
2036
2037   return true;
2038 }
2039
2040 void
2041 WebSocketImpl::UnregisterFeature()
2042 {
2043   MOZ_ASSERT(mDisconnectingOrDisconnected);
2044   MOZ_ASSERT(mWorkerPrivate);
2045   mWorkerPrivate->AssertIsOnWorkerThread();
2046   MOZ_ASSERT(mWorkerFeature);
2047
2048   JSContext* cx = GetCurrentThreadJSContext();
2049   mWorkerPrivate->RemoveFeature(cx, mWorkerFeature);
2050   mWorkerFeature = nullptr;
2051   mWorkerPrivate = nullptr;
2052
2053 #ifdef DEBUG
2054   SetHasFeatureRegistered(false);
2055 #endif
2056 }
2057
2058 nsresult
2059 WebSocketImpl::UpdateURI()
2060 {
2061   AssertIsOnTargetThread();
2062
2063   // Check for Redirections
2064   nsRefPtr<BaseWebSocketChannel> channel;
2065   channel = static_cast<BaseWebSocketChannel*>(mChannel.get());
2066   MOZ_ASSERT(channel);
2067
2068   channel->GetEffectiveURL(mWebSocket->mEffectiveURL);
2069   mSecure = channel->IsEncrypted();
2070
2071   return NS_OK;
2072 }
2073
2074 void
2075 WebSocket::EventListenerAdded(nsIAtom* aType)
2076 {
2077   AssertIsOnMainThread();
2078   UpdateMustKeepAlive();
2079 }
2080
2081 void
2082 WebSocket::EventListenerRemoved(nsIAtom* aType)
2083 {
2084   AssertIsOnMainThread();
2085   UpdateMustKeepAlive();
2086 }
2087
2088 //-----------------------------------------------------------------------------
2089 // WebSocket - methods
2090 //-----------------------------------------------------------------------------
2091
2092 // webIDL: readonly attribute unsigned short readyState;
2093 uint16_t
2094 WebSocket::ReadyState()
2095 {
2096   MutexAutoLock lock(mMutex);
2097   return mReadyState;
2098 }
2099
2100 void
2101 WebSocket::SetReadyState(uint16_t aReadyState)
2102 {
2103   MutexAutoLock lock(mMutex);
2104   mReadyState = aReadyState;
2105 }
2106
2107 // webIDL: readonly attribute unsigned long bufferedAmount;
2108 uint32_t
2109 WebSocket::BufferedAmount() const
2110 {
2111   AssertIsOnTargetThread();
2112   return mOutgoingBufferedAmount;
2113 }
2114
2115 // webIDL: attribute BinaryType binaryType;
2116 dom::BinaryType
2117 WebSocket::BinaryType() const
2118 {
2119   AssertIsOnTargetThread();
2120   return mBinaryType;
2121 }
2122
2123 // webIDL: attribute BinaryType binaryType;
2124 void
2125 WebSocket::SetBinaryType(dom::BinaryType aData)
2126 {
2127   AssertIsOnTargetThread();
2128   mBinaryType = aData;
2129 }
2130
2131 // webIDL: readonly attribute DOMString url
2132 void
2133 WebSocket::GetUrl(nsAString& aURL)
2134 {
2135   AssertIsOnTargetThread();
2136
2137   if (mEffectiveURL.IsEmpty()) {
2138     aURL = mOriginalURL;
2139   } else {
2140     aURL = mEffectiveURL;
2141   }
2142 }
2143
2144 // webIDL: readonly attribute DOMString extensions;
2145 void
2146 WebSocket::GetExtensions(nsAString& aExtensions)
2147 {
2148   AssertIsOnTargetThread();
2149   CopyUTF8toUTF16(mEstablishedExtensions, aExtensions);
2150 }
2151
2152 // webIDL: readonly attribute DOMString protocol;
2153 void
2154 WebSocket::GetProtocol(nsAString& aProtocol)
2155 {
2156   AssertIsOnTargetThread();
2157   CopyUTF8toUTF16(mEstablishedProtocol, aProtocol);
2158 }
2159
2160 // webIDL: void send(DOMString data);
2161 void
2162 WebSocket::Send(const nsAString& aData,
2163                 ErrorResult& aRv)
2164 {
2165   AssertIsOnTargetThread();
2166
2167   NS_ConvertUTF16toUTF8 msgString(aData);
2168   Send(nullptr, msgString, msgString.Length(), false, aRv);
2169 }
2170
2171 void
2172 WebSocket::Send(File& aData, ErrorResult& aRv)
2173 {
2174   AssertIsOnTargetThread();
2175
2176   nsCOMPtr<nsIInputStream> msgStream;
2177   nsresult rv = aData.GetInternalStream(getter_AddRefs(msgStream));
2178   if (NS_FAILED(rv)) {
2179     aRv.Throw(rv);
2180     return;
2181   }
2182
2183   uint64_t msgLength;
2184   rv = aData.GetSize(&msgLength);
2185   if (NS_FAILED(rv)) {
2186     aRv.Throw(rv);
2187     return;
2188   }
2189
2190   if (msgLength > UINT32_MAX) {
2191     aRv.Throw(NS_ERROR_FILE_TOO_BIG);
2192     return;
2193   }
2194
2195   Send(msgStream, EmptyCString(), msgLength, true, aRv);
2196 }
2197
2198 void
2199 WebSocket::Send(const ArrayBuffer& aData,
2200                 ErrorResult& aRv)
2201 {
2202   AssertIsOnTargetThread();
2203
2204   aData.ComputeLengthAndData();
2205
2206   static_assert(sizeof(*aData.Data()) == 1, "byte-sized data required");
2207
2208   uint32_t len = aData.Length();
2209   char* data = reinterpret_cast<char*>(aData.Data());
2210
2211   nsDependentCSubstring msgString(data, len);
2212   Send(nullptr, msgString, len, true, aRv);
2213 }
2214
2215 void
2216 WebSocket::Send(const ArrayBufferView& aData,
2217                 ErrorResult& aRv)
2218 {
2219   AssertIsOnTargetThread();
2220
2221   aData.ComputeLengthAndData();
2222
2223   static_assert(sizeof(*aData.Data()) == 1, "byte-sized data required");
2224
2225   uint32_t len = aData.Length();
2226   char* data = reinterpret_cast<char*>(aData.Data());
2227
2228   nsDependentCSubstring msgString(data, len);
2229   Send(nullptr, msgString, len, true, aRv);
2230 }
2231
2232 void
2233 WebSocket::Send(nsIInputStream* aMsgStream,
2234                 const nsACString& aMsgString,
2235                 uint32_t aMsgLength,
2236                 bool aIsBinary,
2237                 ErrorResult& aRv)
2238 {
2239   AssertIsOnTargetThread();
2240
2241   int64_t readyState = ReadyState();
2242   if (readyState == CONNECTING) {
2243     aRv.Throw(NS_ERROR_DOM_INVALID_STATE_ERR);
2244     return;
2245   }
2246
2247   // Always increment outgoing buffer len, even if closed
2248   mOutgoingBufferedAmount += aMsgLength;
2249
2250   if (readyState == CLOSING ||
2251       readyState == CLOSED) {
2252     return;
2253   }
2254
2255   // We must have mImpl when connected.
2256   MOZ_ASSERT(mImpl);
2257   MOZ_ASSERT(readyState == OPEN, "Unknown state in WebSocket::Send");
2258
2259   nsresult rv;
2260   if (aMsgStream) {
2261     rv = mImpl->mChannel->SendBinaryStream(aMsgStream, aMsgLength);
2262   } else {
2263     if (aIsBinary) {
2264       rv = mImpl->mChannel->SendBinaryMsg(aMsgString);
2265     } else {
2266       rv = mImpl->mChannel->SendMsg(aMsgString);
2267     }
2268   }
2269
2270   if (NS_FAILED(rv)) {
2271     aRv.Throw(rv);
2272     return;
2273   }
2274
2275   UpdateMustKeepAlive();
2276 }
2277
2278 // webIDL: void close(optional unsigned short code, optional DOMString reason):
2279 void
2280 WebSocket::Close(const Optional<uint16_t>& aCode,
2281                  const Optional<nsAString>& aReason,
2282                  ErrorResult& aRv)
2283 {
2284   AssertIsOnTargetThread();
2285
2286   // the reason code is optional, but if provided it must be in a specific range
2287   uint16_t closeCode = 0;
2288   if (aCode.WasPassed()) {
2289     if (aCode.Value() != 1000 && (aCode.Value() < 3000 || aCode.Value() > 4999)) {
2290       aRv.Throw(NS_ERROR_DOM_INVALID_ACCESS_ERR);
2291       return;
2292     }
2293     closeCode = aCode.Value();
2294   }
2295
2296   nsCString closeReason;
2297   if (aReason.WasPassed()) {
2298     CopyUTF16toUTF8(aReason.Value(), closeReason);
2299
2300     // The API requires the UTF-8 string to be 123 or less bytes
2301     if (closeReason.Length() > 123) {
2302       aRv.Throw(NS_ERROR_DOM_SYNTAX_ERR);
2303       return;
2304     }
2305   }
2306
2307   int64_t readyState = ReadyState();
2308   if (readyState == CLOSING ||
2309       readyState == CLOSED) {
2310     return;
2311   }
2312
2313   // If the webSocket is not closed we MUST have a mImpl.
2314   MOZ_ASSERT(mImpl);
2315
2316   if (readyState == CONNECTING) {
2317     mImpl->FailConnection(closeCode, closeReason);
2318     return;
2319   }
2320
2321   MOZ_ASSERT(readyState == OPEN);
2322   mImpl->CloseConnection(closeCode, closeReason);
2323 }
2324
2325 //-----------------------------------------------------------------------------
2326 // WebSocketImpl::nsIObserver
2327 //-----------------------------------------------------------------------------
2328
2329 NS_IMETHODIMP
2330 WebSocketImpl::Observe(nsISupports* aSubject,
2331                        const char* aTopic,
2332                        const char16_t* aData)
2333 {
2334   AssertIsOnMainThread();
2335
2336   int64_t readyState = mWebSocket->ReadyState();
2337   if ((readyState == WebSocket::CLOSING) ||
2338       (readyState == WebSocket::CLOSED)) {
2339     return NS_OK;
2340   }
2341
2342   nsCOMPtr<nsPIDOMWindow> window = do_QueryInterface(aSubject);
2343   if (!mWebSocket->GetOwner() || window != mWebSocket->GetOwner()) {
2344     return NS_OK;
2345   }
2346
2347   if ((strcmp(aTopic, DOM_WINDOW_FROZEN_TOPIC) == 0) ||
2348       (strcmp(aTopic, DOM_WINDOW_DESTROYED_TOPIC) == 0))
2349   {
2350     CloseConnection(nsIWebSocketChannel::CLOSE_GOING_AWAY);
2351   }
2352
2353   return NS_OK;
2354 }
2355
2356 //-----------------------------------------------------------------------------
2357 // WebSocketImpl::nsIRequest
2358 //-----------------------------------------------------------------------------
2359
2360 NS_IMETHODIMP
2361 WebSocketImpl::GetName(nsACString& aName)
2362 {
2363   AssertIsOnMainThread();
2364
2365   CopyUTF16toUTF8(mWebSocket->mOriginalURL, aName);
2366   return NS_OK;
2367 }
2368
2369 NS_IMETHODIMP
2370 WebSocketImpl::IsPending(bool* aValue)
2371 {
2372   AssertIsOnTargetThread();
2373
2374   int64_t readyState = mWebSocket->ReadyState();
2375   *aValue = (readyState != WebSocket::CLOSED);
2376   return NS_OK;
2377 }
2378
2379 NS_IMETHODIMP
2380 WebSocketImpl::GetStatus(nsresult* aStatus)
2381 {
2382   AssertIsOnTargetThread();
2383
2384   *aStatus = NS_OK;
2385   return NS_OK;
2386 }
2387
2388 namespace {
2389
2390 class CancelRunnable MOZ_FINAL : public WorkerRunnable
2391 {
2392 public:
2393   CancelRunnable(WorkerPrivate* aWorkerPrivate, WebSocketImpl* aImpl)
2394     : WorkerRunnable(aWorkerPrivate, WorkerThreadUnchangedBusyCount)
2395     , mImpl(aImpl)
2396   {
2397   }
2398
2399   bool WorkerRun(JSContext* aCx, WorkerPrivate* aWorkerPrivate)
2400   {
2401     aWorkerPrivate->AssertIsOnWorkerThread();
2402     aWorkerPrivate->ModifyBusyCountFromWorker(aCx, true);
2403     return !NS_FAILED(mImpl->CancelInternal());
2404   }
2405
2406   void PostRun(JSContext* aCx, WorkerPrivate* aWorkerPrivate, bool aRunResult)
2407   {
2408     aWorkerPrivate->ModifyBusyCountFromWorker(aCx, false);
2409   }
2410
2411   bool
2412   PreDispatch(JSContext* aCx, WorkerPrivate* aWorkerPrivate)
2413   {
2414     return true;
2415   }
2416
2417   void
2418   PostDispatch(JSContext* aCx, WorkerPrivate* aWorkerPrivate,
2419                bool aDispatchResult)
2420   {
2421   }
2422
2423 private:
2424   nsRefPtr<WebSocketImpl> mImpl;
2425 };
2426
2427 } // anonymous namespace
2428
2429 // Window closed, stop/reload button pressed, user navigated away from page, etc.
2430 NS_IMETHODIMP
2431 WebSocketImpl::Cancel(nsresult aStatus)
2432 {
2433   AssertIsOnMainThread();
2434
2435   if (!mIsMainThread) {
2436     MOZ_ASSERT(mWorkerPrivate);
2437     nsRefPtr<CancelRunnable> runnable =
2438       new CancelRunnable(mWorkerPrivate, this);
2439     if (!runnable->Dispatch(nullptr)) {
2440       return NS_ERROR_FAILURE;
2441     }
2442
2443     return NS_OK;
2444   }
2445
2446   return CancelInternal();
2447 }
2448
2449 nsresult
2450 WebSocketImpl::CancelInternal()
2451 {
2452   AssertIsOnTargetThread();
2453
2454    // If CancelInternal is called by a runnable, we may already be disconnected
2455    // by the time it runs.
2456   if (mDisconnectingOrDisconnected) {
2457     return NS_OK;
2458   }
2459
2460   int64_t readyState = mWebSocket->ReadyState();
2461   if (readyState == WebSocket::CLOSING || readyState == WebSocket::CLOSED) {
2462     return NS_OK;
2463   }
2464
2465   ConsoleError();
2466
2467   return CloseConnection(nsIWebSocketChannel::CLOSE_GOING_AWAY);
2468 }
2469
2470 NS_IMETHODIMP
2471 WebSocketImpl::Suspend()
2472 {
2473   AssertIsOnMainThread();
2474   return NS_ERROR_NOT_IMPLEMENTED;
2475 }
2476
2477 NS_IMETHODIMP
2478 WebSocketImpl::Resume()
2479 {
2480   AssertIsOnMainThread();
2481   return NS_ERROR_NOT_IMPLEMENTED;
2482 }
2483
2484 NS_IMETHODIMP
2485 WebSocketImpl::GetLoadGroup(nsILoadGroup** aLoadGroup)
2486 {
2487   AssertIsOnMainThread();
2488
2489   *aLoadGroup = nullptr;
2490
2491   if (mIsMainThread) {
2492     nsresult rv;
2493     nsIScriptContext* sc = mWebSocket->GetContextForEventHandlers(&rv);
2494     nsCOMPtr<nsIDocument> doc =
2495       nsContentUtils::GetDocumentFromScriptContext(sc);
2496
2497     if (doc) {
2498       *aLoadGroup = doc->GetDocumentLoadGroup().take();
2499     }
2500
2501     return NS_OK;
2502   }
2503
2504   MOZ_ASSERT(mWorkerPrivate);
2505
2506   // Walk up to our containing page
2507   WorkerPrivate* wp = mWorkerPrivate;
2508   while (wp->GetParent()) {
2509     wp = wp->GetParent();
2510   }
2511
2512   nsPIDOMWindow* window = wp->GetWindow();
2513   if (!window) {
2514     return NS_OK;
2515   }
2516
2517   nsIDocument* doc = window->GetExtantDoc();
2518   if (doc) {
2519     *aLoadGroup = doc->GetDocumentLoadGroup().take();
2520   }
2521
2522   return NS_OK;
2523 }
2524
2525 NS_IMETHODIMP
2526 WebSocketImpl::SetLoadGroup(nsILoadGroup* aLoadGroup)
2527 {
2528   AssertIsOnMainThread();
2529   return NS_ERROR_UNEXPECTED;
2530 }
2531
2532 NS_IMETHODIMP
2533 WebSocketImpl::GetLoadFlags(nsLoadFlags* aLoadFlags)
2534 {
2535   AssertIsOnMainThread();
2536
2537   *aLoadFlags = nsIRequest::LOAD_BACKGROUND;
2538   return NS_OK;
2539 }
2540
2541 NS_IMETHODIMP
2542 WebSocketImpl::SetLoadFlags(nsLoadFlags aLoadFlags)
2543 {
2544   AssertIsOnMainThread();
2545
2546   // we won't change the load flags at all.
2547   return NS_OK;
2548 }
2549
2550 namespace {
2551
2552 class WorkerRunnableDispatcher MOZ_FINAL : public WorkerRunnable
2553 {
2554 public:
2555   WorkerRunnableDispatcher(WorkerPrivate* aWorkerPrivate, nsIRunnable* aEvent)
2556     : WorkerRunnable(aWorkerPrivate, WorkerThreadUnchangedBusyCount)
2557     , mEvent(aEvent)
2558   {
2559   }
2560
2561   bool WorkerRun(JSContext* aCx, WorkerPrivate* aWorkerPrivate)
2562   {
2563     aWorkerPrivate->AssertIsOnWorkerThread();
2564     aWorkerPrivate->ModifyBusyCountFromWorker(aCx, true);
2565     return !NS_FAILED(mEvent->Run());
2566   }
2567
2568   void PostRun(JSContext* aCx, WorkerPrivate* aWorkerPrivate, bool aRunResult)
2569   {
2570     aWorkerPrivate->ModifyBusyCountFromWorker(aCx, false);
2571   }
2572
2573   bool
2574   PreDispatch(JSContext* aCx, WorkerPrivate* aWorkerPrivate)
2575   {
2576     return true;
2577   }
2578
2579   void
2580   PostDispatch(JSContext* aCx, WorkerPrivate* aWorkerPrivate,
2581                bool aDispatchResult)
2582   {
2583   }
2584
2585 private:
2586   nsCOMPtr<nsIRunnable> mEvent;
2587 };
2588
2589 } // anonymous namespace
2590
2591 NS_IMETHODIMP
2592 WebSocketImpl::Dispatch(nsIRunnable* aEvent, uint32_t aFlags)
2593 {
2594   // If the target is the main-thread we can just dispatch the runnable.
2595   if (mIsMainThread) {
2596     return NS_DispatchToMainThread(aEvent);
2597   }
2598
2599   // No messages when disconnected.
2600   if (mDisconnectingOrDisconnected) {
2601     NS_WARNING("Dispatching a WebSocket event after the disconnection!");
2602     return NS_OK;
2603   }
2604
2605   if (mWorkerShuttingDown) {
2606     return NS_OK;
2607   }
2608
2609   MOZ_ASSERT(mWorkerPrivate);
2610
2611 #ifdef DEBUG
2612   MOZ_ASSERT(HasFeatureRegistered());
2613 #endif
2614
2615   // If the target is a worker, we have to use a custom WorkerRunnableDispatcher
2616   // runnable.
2617   nsRefPtr<WorkerRunnableDispatcher> event =
2618     new WorkerRunnableDispatcher(mWorkerPrivate, aEvent);
2619   if (!event->Dispatch(nullptr)) {
2620     return NS_ERROR_FAILURE;
2621   }
2622
2623   return NS_OK;
2624 }
2625
2626 NS_IMETHODIMP
2627 WebSocketImpl::IsOnCurrentThread(bool* aResult)
2628 {
2629   *aResult = IsTargetThread();
2630   return NS_OK;
2631 }
2632
2633 bool
2634 WebSocketImpl::IsTargetThread() const
2635 {
2636   return NS_IsMainThread() == mIsMainThread;
2637 }
2638
2639 void
2640 WebSocket::AssertIsOnTargetThread() const
2641 {
2642   MOZ_ASSERT(NS_IsMainThread() == mIsMainThread);
2643 }
2644
2645 } // dom namespace
2646 } // mozilla namespace