From e832cff45edfc17dc8f1d5d9ecd39020e1f5e49b Mon Sep 17 00:00:00 2001
From: Andrea Marchesini <amarchesini@mozilla.com>
Date: Sat, 30 Dec 2017 12:47:12 +0100
Subject: [PATCH] Bug 1426887 - Avoid adding/removing streams to
 nsMultiplexInputStream when reading in Presentation API, r=jdm

---
 .../PresentationTCPSessionTransport.cpp       | 99 +++++++++----------
 .../PresentationTCPSessionTransport.h         |  9 +-
 2 files changed, 50 insertions(+), 58 deletions(-)

diff --git a/dom/presentation/PresentationTCPSessionTransport.cpp b/dom/presentation/PresentationTCPSessionTransport.cpp
index 965c5f066f9e..eb7a7e7e52f4 100644
--- a/dom/presentation/PresentationTCPSessionTransport.cpp
+++ b/dom/presentation/PresentationTCPSessionTransport.cpp
@@ -61,7 +61,7 @@ CopierCallbacks::OnStopRequest(nsIRequest* aRequest, nsISupports* aContext, nsre
 NS_IMPL_CYCLE_COLLECTION(PresentationTCPSessionTransport, mTransport,
                          mSocketInputStream, mSocketOutputStream,
                          mInputStreamPump, mInputStreamScriptable,
-                         mMultiplexStream, mMultiplexStreamCopier, mCallback)
+                         mCallback)
 
 NS_IMPL_CYCLE_COLLECTING_ADDREF(PresentationTCPSessionTransport)
 NS_IMPL_CYCLE_COLLECTING_RELEASE(PresentationTCPSessionTransport)
@@ -236,36 +236,6 @@ PresentationTCPSessionTransport::CreateStream()
     return rv;
   }
 
-  mMultiplexStream = do_CreateInstance("@mozilla.org/io/multiplex-input-stream;1", &rv);
-  if (NS_WARN_IF(NS_FAILED(rv))) {
-    return rv;
-  }
-  nsCOMPtr<nsIInputStream> stream = do_QueryInterface(mMultiplexStream);
-
-  mMultiplexStreamCopier = do_CreateInstance("@mozilla.org/network/async-stream-copier;1", &rv);
-  if (NS_WARN_IF(NS_FAILED(rv))) {
-    return rv;
-  }
-
-  nsCOMPtr<nsISocketTransportService> sts =
-    do_GetService(NS_SOCKETTRANSPORTSERVICE_CONTRACTID);
-  if (NS_WARN_IF(!sts)) {
-    return NS_ERROR_NOT_AVAILABLE;
-  }
-
-  nsCOMPtr<nsIEventTarget> target = do_QueryInterface(sts);
-  rv = mMultiplexStreamCopier->Init(stream,
-                                    mSocketOutputStream,
-                                    target,
-                                    true, /* source buffered */
-                                    false, /* sink buffered */
-                                    BUFFER_SIZE,
-                                    false, /* close source */
-                                    false); /* close sink */
-  if (NS_WARN_IF(NS_FAILED(rv))) {
-    return rv;
-  }
-
   return NS_OK;
 }
 
@@ -347,23 +317,59 @@ PresentationTCPSessionTransport::GetSelfAddress(nsINetAddr** aSelfAddress)
   return mTransport->GetScriptableSelfAddr(aSelfAddress);
 }
 
-void
+nsresult
 PresentationTCPSessionTransport::EnsureCopying()
 {
   if (mAsyncCopierActive) {
-    return;
+    return NS_OK;
   }
 
   mAsyncCopierActive = true;
+
+  nsresult rv;
+
+  nsCOMPtr<nsIMultiplexInputStream> multiplexStream =
+    do_CreateInstance("@mozilla.org/io/multiplex-input-stream;1", &rv);
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  nsCOMPtr<nsIInputStream> stream = do_QueryInterface(multiplexStream);
+
+  while (!mPendingData.IsEmpty()) {
+    nsCOMPtr<nsIInputStream> stream = mPendingData[0];
+    multiplexStream->AppendStream(stream);
+    mPendingData.RemoveElementAt(0);
+  }
+
+  nsCOMPtr<nsIAsyncStreamCopier> copier =
+    do_CreateInstance("@mozilla.org/network/async-stream-copier;1", &rv);
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  nsCOMPtr<nsISocketTransportService> sts =
+      do_GetService("@mozilla.org/network/socket-transport-service;1");
+
+  nsCOMPtr<nsIEventTarget> target = do_QueryInterface(sts);
+  rv = copier->Init(stream,
+                    mSocketOutputStream,
+                    target,
+                    true, /* source buffered */
+                    false, /* sink buffered */
+                    BUFFER_SIZE,
+                    false, /* close source */
+                    false); /* close sink */
+  NS_ENSURE_SUCCESS(rv, rv);
+
   RefPtr<CopierCallbacks> callbacks = new CopierCallbacks(this);
-  Unused << NS_WARN_IF(NS_FAILED(mMultiplexStreamCopier->AsyncCopy(callbacks, nullptr)));
+  rv = copier->AsyncCopy(callbacks, nullptr);
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  return NS_OK;
 }
 
 void
 PresentationTCPSessionTransport::NotifyCopyComplete(nsresult aStatus)
 {
   mAsyncCopierActive = false;
-  mMultiplexStream->RemoveStream(0);
+
   if (NS_WARN_IF(NS_FAILED(aStatus))) {
     if (mReadyState != ReadyState::CLOSED) {
       mCloseStatus = aStatus;
@@ -372,13 +378,7 @@ PresentationTCPSessionTransport::NotifyCopyComplete(nsresult aStatus)
     return;
   }
 
-  uint32_t count;
-  nsresult rv = mMultiplexStream->GetCount(&count);
-  if (NS_WARN_IF(NS_FAILED(rv))) {
-    return;
-  }
-
-  if (count) {
+  if (!mPendingData.IsEmpty()) {
     EnsureCopying();
     return;
   }
@@ -410,7 +410,7 @@ PresentationTCPSessionTransport::Send(const nsAString& aData)
     return NS_ERROR_DOM_INVALID_STATE_ERR;
   }
 
-  mMultiplexStream->AppendStream(stream);
+  mPendingData.AppendElement(stream);
 
   EnsureCopying();
 
@@ -441,9 +441,8 @@ PresentationTCPSessionTransport::Close(nsresult aReason)
   mCloseStatus = aReason;
   SetReadyState(ReadyState::CLOSING);
 
-  uint32_t count = 0;
-  mMultiplexStream->GetCount(&count);
-  if (!count) {
+  if (!mAsyncCopierActive) {
+    mPendingData.Clear();
     mSocketOutputStream->Close();
   }
 
@@ -540,15 +539,9 @@ PresentationTCPSessionTransport::OnStopRequest(nsIRequest* aRequest,
 
   MOZ_ASSERT(NS_IsMainThread());
 
-  uint32_t count;
-  nsresult rv = mMultiplexStream->GetCount(&count);
-  if (NS_WARN_IF(NS_FAILED(rv))) {
-    return rv;
-  }
-
   mInputStreamPump = nullptr;
 
-  if (count != 0 && NS_SUCCEEDED(aStatusCode)) {
+  if (mAsyncCopierActive && NS_SUCCEEDED(aStatusCode)) {
     // If we have some buffered output still, and status is not an error, the
     // other side has done a half-close, but we don't want to be in the close
     // state until we are done sending everything that was buffered. We also
diff --git a/dom/presentation/PresentationTCPSessionTransport.h b/dom/presentation/PresentationTCPSessionTransport.h
index f36b371a4b7f..40cd094ce6a1 100644
--- a/dom/presentation/PresentationTCPSessionTransport.h
+++ b/dom/presentation/PresentationTCPSessionTransport.h
@@ -64,7 +64,7 @@ private:
 
   nsresult CreateInputStreamPump();
 
-  void EnsureCopying();
+  nsresult EnsureCopying();
 
   enum class ReadyState {
     CONNECTING,
@@ -96,12 +96,11 @@ private:
   nsCOMPtr<nsIInputStreamPump> mInputStreamPump;
   nsCOMPtr<nsIScriptableInputStream> mInputStreamScriptable;
 
-  // Output stream machinery
-  nsCOMPtr<nsIMultiplexInputStream> mMultiplexStream;
-  nsCOMPtr<nsIAsyncStreamCopier> mMultiplexStreamCopier;
-
   nsCOMPtr<nsIPresentationSessionTransportCallback> mCallback;
   nsCOMPtr<nsIPresentationSessionTransportBuilderListener> mListener;
+
+  // The data to be sent.
+  nsTArray<nsCOMPtr<nsIInputStream>> mPendingData;
 };
 
 } // namespace dom