Bug 1820434: Isolate AsyncInputStream from InputToReadableStream for workers r=asuth

nsIAsyncInputStreams keep a reference to the callback object, causing
possible leaks which will result in assertions in workers.  Also use a
StrongWorkerRef to ensure the worker remains around until the
OnInputStreamReady callback occurs.

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

dom/streams/UnderlyingSourceCallbackHelpers.cpp

@@ -9,6 +9,9 @@
 #include "mozilla/dom/ReadableStreamDefaultController.h"
 #include "mozilla/dom/UnderlyingSourceCallbackHelpers.h"
 #include "mozilla/dom/UnderlyingSourceBinding.h"
+#include "mozilla/dom/WorkerCommon.h"
+#include "mozilla/dom/WorkerPrivate.h"
+#include "mozilla/dom/WorkerRunnable.h"
 #include "js/experimental/TypedData.h"
 
 namespace mozilla::dom {
@@ -148,10 +151,67 @@ already_AddRefed<Promise> UnderlyingSourceAlgorithmsWrapper::CancelCallback(
       aRv);
 }
 
+NS_IMPL_ISUPPORTS(InputStreamHolder, nsIInputStreamCallback)
+
+InputStreamHolder::InputStreamHolder(JSContext* aCx,
+                                     InputToReadableStreamAlgorithms* aCallback,
+                                     nsIAsyncInputStream* aInput)
+    : mCallback(aCallback), mInput(aInput) {
+  if (!NS_IsMainThread()) {
+    // We're in a worker
+    WorkerPrivate* workerPrivate = GetWorkerPrivateFromContext(aCx);
+    MOZ_ASSERT(workerPrivate);
+
+    workerPrivate->AssertIsOnWorkerThread();
+
+    // Note, this will create a ref-cycle between the holder and the stream.
+    // The cycle is broken when the stream is closed or the worker begins
+    // shutting down.
+    mWorkerRef =
+        StrongWorkerRef::Create(workerPrivate, "InputStreamHolder",
+                                [self = RefPtr{this}]() { self->Shutdown(); });
+    if (NS_WARN_IF(!mWorkerRef)) {
+      return;
+    }
+  }
+}
+
+InputStreamHolder::~InputStreamHolder() { MOZ_ASSERT(!mCallback); }
+
+void InputStreamHolder::Shutdown() {
+  if (mCallback) {
+    mCallback = nullptr;
+    mInput->CloseWithStatus(NS_BASE_STREAM_CLOSED);
+    mWorkerRef = nullptr;
+    // If we have an AsyncWait running, we'll get a callback and clear
+    // the mAsyncWaitWorkerRef
+  }
+}
+
+nsresult InputStreamHolder::AsyncWait(uint32_t aFlags, uint32_t aRequestedCount,
+                                      nsIEventTarget* aEventTarget) {
+  nsresult rv = mInput->AsyncWait(this, aFlags, aRequestedCount, aEventTarget);
+  if (NS_SUCCEEDED(rv)) {
+    mAsyncWaitWorkerRef = mWorkerRef;
+  }
+  return rv;
+}
+
+NS_IMETHODIMP InputStreamHolder::OnInputStreamReady(
+    nsIAsyncInputStream* aStream) {
+  mAsyncWaitWorkerRef = nullptr;
+  // We may get called back after ::Shutdown()
+  if (mCallback) {
+    return mCallback->OnInputStreamReady(aStream);
+  }
+  return NS_ERROR_FAILURE;
+}
+
 NS_IMPL_ISUPPORTS_CYCLE_COLLECTION_INHERITED_0(
     InputToReadableStreamAlgorithms, UnderlyingSourceAlgorithmsWrapper)
 NS_IMPL_CYCLE_COLLECTION_INHERITED(InputToReadableStreamAlgorithms,
-                                   UnderlyingSourceAlgorithmsWrapper, mStream)
+                                   UnderlyingSourceAlgorithmsWrapper, mStream,
+                                   mOwningEventTarget)
 
 already_AddRefed<Promise> InputToReadableStreamAlgorithms::PullCallbackImpl(
     JSContext* aCx, ReadableStreamController& aController, ErrorResult& aRv) {
@@ -188,7 +248,7 @@ already_AddRefed<Promise> InputToReadableStreamAlgorithms::PullCallbackImpl(
 
   MOZ_DIAGNOSTIC_ASSERT(mInput);
 
-  nsresult rv = mInput->AsyncWait(this, 0, 0, mOwningEventTarget);
+  nsresult rv = mInput->AsyncWait(0, 0, mOwningEventTarget);
   if (NS_WARN_IF(NS_FAILED(rv))) {
     ErrorPropagation(aCx, stream, rv);
     return nullptr;
@@ -204,7 +264,7 @@ InputToReadableStreamAlgorithms::OnInputStreamReady(
   MOZ_DIAGNOSTIC_ASSERT(aStream);
 
   // Already closed. We have nothing else to do here.
-  if (mState == eClosed) {
+  if (mState == eClosed || !mStream) {
     return NS_OK;
   }
 
@@ -293,7 +353,7 @@ void InputToReadableStreamAlgorithms::WriteIntoReadRequestBuffer(
     return;
   }
 
-  rv = mInput->AsyncWait(this, 0, 0, mOwningEventTarget);
+  rv = mInput->AsyncWait(0, 0, mOwningEventTarget);
   if (NS_WARN_IF(NS_FAILED(rv))) {
     ErrorPropagation(aCx, aStream, rv);
     return;
@@ -360,7 +420,11 @@ void InputToReadableStreamAlgorithms::CloseAndReleaseObjects(
 }
 
 void InputToReadableStreamAlgorithms::ReleaseObjects() {
-  mInput->CloseWithStatus(NS_BASE_STREAM_CLOSED);
+  if (mInput) {
+    mInput->CloseWithStatus(NS_BASE_STREAM_CLOSED);
+    mInput->Shutdown();
+    mInput = nullptr;
+  }
 }
 
 void InputToReadableStreamAlgorithms::ErrorPropagation(JSContext* aCx,

dom/streams/UnderlyingSourceCallbackHelpers.h

@@ -27,6 +27,7 @@ enum class nsresult : uint32_t;
 
 namespace mozilla::dom {
 
+class StrongWorkerRef;
 class BodyStreamHolder;
 class ReadableStreamController;
 class ReadableStream;
@@ -159,6 +160,48 @@ class UnderlyingSourceAlgorithmsWrapper
   }
 };
 
+class InputToReadableStreamAlgorithms;
+
+// This class exists to isolate InputToReadableStreamAlgorithms from the
+// nsIAsyncInputStream.  If we call AsyncWait(this,...), it holds a
+// reference to 'this' which can't be cc'd, and we can leak the stream,
+// causing a Worker to assert with globalScopeAlive.  By isolating
+// ourselves from the inputstream, we can safely be CC'd if needed and
+// will inform the inputstream to shut down.
+class InputStreamHolder final : public nsIInputStreamCallback {
+ public:
+  NS_DECL_THREADSAFE_ISUPPORTS
+  NS_DECL_NSIINPUTSTREAMCALLBACK
+
+  InputStreamHolder(JSContext* aCx, InputToReadableStreamAlgorithms* aCallback,
+                    nsIAsyncInputStream* aInput);
+
+  // This MUST be called before we're destroyed
+  void Shutdown();
+
+  // These just proxy the calls to the nsIAsyncInputStream
+  nsresult AsyncWait(uint32_t aFlags, uint32_t aRequestedCount,
+                     nsIEventTarget* aEventTarget);
+  nsresult Available(uint64_t* aSize) { return mInput->Available(aSize); }
+  nsresult Read(char* aBuffer, uint32_t aLength, uint32_t* aWritten) {
+    return mInput->Read(aBuffer, aLength, aWritten);
+  }
+  nsresult CloseWithStatus(nsresult aStatus) {
+    return mInput->CloseWithStatus(aStatus);
+  }
+
+ private:
+  ~InputStreamHolder();
+
+  // Purposely rawptr to avoid cycles.  Must be cleared with Shutdown() before
+  // destruction
+  InputToReadableStreamAlgorithms* mCallback;
+  // To ensure the worker sticks around
+  RefPtr<StrongWorkerRef> mAsyncWaitWorkerRef;
+  RefPtr<StrongWorkerRef> mWorkerRef;
+  nsCOMPtr<nsIAsyncInputStream> mInput;
+};
+
 class InputToReadableStreamAlgorithms final
     : public UnderlyingSourceAlgorithmsWrapper,
       public nsIInputStreamCallback {
@@ -167,11 +210,11 @@ class InputToReadableStreamAlgorithms final
   NS_DECL_CYCLE_COLLECTION_CLASS_INHERITED(InputToReadableStreamAlgorithms,
                                            UnderlyingSourceAlgorithmsWrapper)
 
-  InputToReadableStreamAlgorithms(nsIAsyncInputStream* aInput,
+  InputToReadableStreamAlgorithms(JSContext* aCx, nsIAsyncInputStream* aInput,
                                   ReadableStream* aStream)
       : mState(eInitializing),
         mOwningEventTarget(GetCurrentSerialEventTarget()),
-        mInput(aInput),
+        mInput(new InputStreamHolder(aCx, this, aInput)),
         mStream(aStream) {}
 
   // Streams algorithms
@@ -183,7 +226,11 @@ class InputToReadableStreamAlgorithms final
   void ReleaseObjects() override;
 
  private:
-  ~InputToReadableStreamAlgorithms() override = default;
+  ~InputToReadableStreamAlgorithms() {
+    if (mInput) {
+      mInput->Shutdown();
+    }
+  }
 
   MOZ_CAN_RUN_SCRIPT_BOUNDARY void CloseAndReleaseObjects(
       JSContext* aCx, ReadableStream* aStream);
@@ -227,7 +274,7 @@ class InputToReadableStreamAlgorithms final
 
   nsCOMPtr<nsIEventTarget> mOwningEventTarget;
 
-  nsCOMPtr<nsIAsyncInputStream> mInput;
+  RefPtr<InputStreamHolder> mInput;
   RefPtr<ReadableStream> mStream;
 };
 

dom/webtransport/api/WebTransport.cpp

@@ -45,6 +45,8 @@ NS_IMPL_CYCLE_COLLECTION_TRAVERSE_END
 
 NS_IMPL_CYCLE_COLLECTION_UNLINK_BEGIN(WebTransport)
   NS_IMPL_CYCLE_COLLECTION_UNLINK(mGlobal)
+  NS_IMPL_CYCLE_COLLECTION_UNLINK(mUnidirectionalStreams)
+  NS_IMPL_CYCLE_COLLECTION_UNLINK(mBidirectionalStreams)
   NS_IMPL_CYCLE_COLLECTION_UNLINK(mIncomingUnidirectionalStreams)
   NS_IMPL_CYCLE_COLLECTION_UNLINK(mIncomingBidirectionalStreams)
   NS_IMPL_CYCLE_COLLECTION_UNLINK(mIncomingUnidirectionalAlgorithm)

dom/webtransport/api/WebTransportDatagramDuplexStream.cpp

@@ -15,7 +15,7 @@ namespace mozilla::dom {
 NS_IMPL_CYCLE_COLLECTION_WRAPPERCACHE(WebTransportDatagramDuplexStream, mGlobal,
                                       mReadable, mWritable, mWebTransport,
                                       mIncomingAlgorithms, mOutgoingAlgorithms)
-
+// mIncomingDatagramsQueue can't participate in a cycle
 NS_IMPL_CYCLE_COLLECTING_ADDREF(WebTransportDatagramDuplexStream)
 NS_IMPL_CYCLE_COLLECTING_RELEASE(WebTransportDatagramDuplexStream)
 NS_INTERFACE_MAP_BEGIN_CYCLE_COLLECTION(WebTransportDatagramDuplexStream)

dom/webtransport/api/WebTransportReceiveStream.cpp

@@ -41,7 +41,7 @@ already_AddRefed<WebTransportReceiveStream> WebTransportReceiveStream::Create(
 
   nsCOMPtr<nsIAsyncInputStream> inputStream = receiver;
   auto algorithms = MakeRefPtr<InputToReadableStreamAlgorithms>(
-      inputStream, (ReadableStream*)stream);
+      cx, inputStream, (ReadableStream*)stream);
 
   stream->SetUpByteNative(cx, *algorithms, Some(0.0), aRv);
   if (aRv.Failed()) {