Bug 1451731 - Synchronize access to various stream classes' async wait callback reference - part 5 - FileSnapshot StreamWrapper, r=mayhemer

dom/indexedDB/FileSnapshot.cpp

@@ -10,6 +10,7 @@
 #include "IDBFileHandle.h"
 #include "IDBMutableFile.h"
 #include "mozilla/Assertions.h"
+#include "mozilla/Mutex.h"
 #include "nsIAsyncInputStream.h"
 #include "nsICloneableInputStream.h"
 #include "nsIIPCSerializableInputStream.h"
@@ -36,8 +37,11 @@ class StreamWrapper final
   bool mFinished;
 
   // This is needed to call OnInputStreamReady() with the correct inputStream.
+  // It is protected by mutex.
   nsCOMPtr<nsIInputStreamCallback> mAsyncWaitCallback;
 
+  Mutex mMutex;
+
 public:
   StreamWrapper(nsIInputStream* aInputStream,
                 IDBFileHandle* aFileHandle)
@@ -45,6 +49,7 @@ public:
     , mInputStream(aInputStream)
     , mFileHandle(aFileHandle)
     , mFinished(false)
+    , mMutex("StreamWrapper::mMutex")
   {
     AssertIsOnOwningThread();
     MOZ_ASSERT(aInputStream);
@@ -355,17 +360,18 @@ StreamWrapper::AsyncWait(nsIInputStreamCallback* aCallback,
     return NS_ERROR_NO_INTERFACE;
   }
 
-  if (mAsyncWaitCallback && aCallback) {
-    return NS_ERROR_FAILURE;
+  nsCOMPtr<nsIInputStreamCallback> callback = aCallback ? this : nullptr;
+  {
+    MutexAutoLock lock(mMutex);
+
+    if (mAsyncWaitCallback && aCallback) {
+      return NS_ERROR_FAILURE;
+    }
+
+    mAsyncWaitCallback = aCallback;
   }
 
-  mAsyncWaitCallback = aCallback;
-
-  if (!mAsyncWaitCallback) {
-    return NS_OK;
-  }
-
-  return stream->AsyncWait(this, aFlags, aRequestedCount, aEventTarget);
+  return stream->AsyncWait(callback, aFlags, aRequestedCount, aEventTarget);
 }
 
 // nsIInputStreamCallback
@@ -378,14 +384,19 @@ StreamWrapper::OnInputStreamReady(nsIAsyncInputStream* aStream)
     return NS_ERROR_NO_INTERFACE;
   }
 
-  // We have been canceled in the meanwhile.
-  if (!mAsyncWaitCallback) {
-    return NS_OK;
+  nsCOMPtr<nsIInputStreamCallback> callback;
+  {
+    MutexAutoLock lock(mMutex);
+
+    // We have been canceled in the meanwhile.
+    if (!mAsyncWaitCallback) {
+      return NS_OK;
+    }
+
+    callback.swap(mAsyncWaitCallback);
   }
 
-  nsCOMPtr<nsIInputStreamCallback> callback;
-  callback.swap(mAsyncWaitCallback);
-
+  MOZ_ASSERT(callback);
   return callback->OnInputStreamReady(this);
 }