Bug 1352681 - Make SAB rawBuffer refcounting for structured clone more sophisticated. r=sfink

--HG--
extra : rebase_source : f5c97970013daab78075e2fe68c9704cda7064cd

Author: Lars T Hansen

js/public/StructuredClone.h

@@ -17,6 +17,7 @@
 #include "js/RootingAPI.h"
 #include "js/TypeDecls.h"
 #include "js/Value.h"
+#include "js/Vector.h"
 
 struct JSRuntime;
 struct JSStructuredCloneReader;
@@ -188,6 +189,28 @@ enum OwnTransferablePolicy {
     NoTransferables
 };
 
+namespace js
+{
+    class SharedArrayRawBuffer;
+
+    class SharedArrayRawBufferRefs
+    {
+      public:
+        SharedArrayRawBufferRefs() = default;
+        SharedArrayRawBufferRefs(SharedArrayRawBufferRefs&& other) = default;
+        SharedArrayRawBufferRefs& operator=(SharedArrayRawBufferRefs&& other);
+        ~SharedArrayRawBufferRefs();
+
+        MOZ_MUST_USE bool acquire(JSContext* cx, SharedArrayRawBuffer* rawbuf);
+        MOZ_MUST_USE bool acquireAll(JSContext* cx, const SharedArrayRawBufferRefs& that);
+        void takeOwnership(SharedArrayRawBufferRefs&&);
+        void releaseAll();
+
+      private:
+        js::Vector<js::SharedArrayRawBuffer*, 0, js::SystemAllocPolicy> refs_;
+    };
+}
+
 class MOZ_NON_MEMMOVABLE JS_PUBLIC_API(JSStructuredCloneData) :
     public mozilla::BufferList<js::SystemAllocPolicy>
 {
@@ -201,6 +224,7 @@ class MOZ_NON_MEMMOVABLE JS_PUBLIC_API(JSStructuredCloneData) :
     const JSStructuredCloneCallbacks* callbacks_;
     void* closure_;
     OwnTransferablePolicy ownTransferables_;
+    js::SharedArrayRawBufferRefs refsHeld_;
 
     void setOptionalCallbacks(const JSStructuredCloneCallbacks* callbacks,
                               void* closure,
@@ -279,7 +303,8 @@ class JS_PUBLIC_API(JSAutoStructuredCloneBuffer) {
     void clear(const JSStructuredCloneCallbacks* optionalCallbacks=nullptr, void* closure=nullptr);
 
     /** Copy some memory. It will be automatically freed by the destructor. */
-    bool copy(const JSStructuredCloneData& data, uint32_t version=JS_STRUCTURED_CLONE_VERSION,
+    bool copy(JSContext* cx, const JSStructuredCloneData& data,
+              uint32_t version=JS_STRUCTURED_CLONE_VERSION,
               const JSStructuredCloneCallbacks* callbacks=nullptr, void* closure=nullptr);
 
     /**

js/src/vm/StructuredClone.cpp

@@ -229,6 +229,69 @@ struct BufferIterator {
     typename BufferList::IterImpl mIter;
 };
 
+SharedArrayRawBufferRefs&
+SharedArrayRawBufferRefs::operator=(SharedArrayRawBufferRefs&& other)
+{
+    takeOwnership(Move(other));
+    return *this;
+}
+
+SharedArrayRawBufferRefs::~SharedArrayRawBufferRefs()
+{
+    releaseAll();
+}
+
+bool
+SharedArrayRawBufferRefs::acquire(JSContext* cx, SharedArrayRawBuffer* rawbuf)
+{
+    if (!refs_.append(rawbuf)) {
+        ReportOutOfMemory(cx);
+        return false;
+    }
+
+    if (!rawbuf->addReference()) {
+        refs_.popBack();
+        JS_ReportErrorNumberASCII(cx, GetErrorMessage, nullptr, JSMSG_SC_SAB_REFCNT_OFLO);
+        return false;
+    }
+
+    return true;
+}
+
+bool
+SharedArrayRawBufferRefs::acquireAll(JSContext* cx, const SharedArrayRawBufferRefs& that)
+{
+    if (!refs_.reserve(refs_.length() + that.refs_.length())) {
+        ReportOutOfMemory(cx);
+        return false;
+    }
+
+    for (auto ref : that.refs_) {
+        if (!ref->addReference()) {
+            JS_ReportErrorNumberASCII(cx, GetErrorMessage, nullptr, JSMSG_SC_SAB_REFCNT_OFLO);
+            return false;
+        }
+        MOZ_ALWAYS_TRUE(refs_.append(ref));
+    }
+
+    return true;
+}
+
+void
+SharedArrayRawBufferRefs::takeOwnership(SharedArrayRawBufferRefs&& other)
+{
+    MOZ_ASSERT(refs_.empty());
+    refs_ = Move(other.refs_);
+}
+
+void
+SharedArrayRawBufferRefs::releaseAll()
+{
+    for (auto ref : refs_)
+        ref->dropReference();
+    refs_.clear();
+}
+
 struct SCOutput {
   public:
     using Iter = BufferIterator<uint64_t, TempAllocPolicy>;
@@ -408,6 +471,9 @@ struct JSStructuredCloneWriter {
     bool extractBuffer(JSStructuredCloneData* data) {
         bool success = out.extractBuffer(data);
         if (success) {
+            // Move the SharedArrayRawBuf references here, SCOutput::extractBuffer
+            // moves the serialized data.
+            data->refsHeld_.takeOwnership(Move(refsHeld));
             data->setOptionalCallbacks(callbacks, closure,
                                        OwnTransferablePolicy::OwnsTransferablesIfAny);
         }
@@ -486,6 +552,9 @@ struct JSStructuredCloneWriter {
 
     const JS::CloneDataPolicy cloneDataPolicy;
 
+    // SharedArrayRawBuffers whose reference counts we have incremented.
+    SharedArrayRawBufferRefs refsHeld;
+
     friend bool JS_WriteString(JSStructuredCloneWriter* w, HandleString str);
     friend bool JS_WriteTypedArray(JSStructuredCloneWriter* w, HandleValue v);
     friend bool JS_ObjectNotWritten(JSStructuredCloneWriter* w, HandleObject obj);
@@ -1161,12 +1230,8 @@ JSStructuredCloneWriter::writeSharedArrayBuffer(HandleObject obj)
     Rooted<SharedArrayBufferObject*> sharedArrayBuffer(context(), &CheckedUnwrap(obj)->as<SharedArrayBufferObject>());
     SharedArrayRawBuffer* rawbuf = sharedArrayBuffer->rawBufferObject();
 
-    // Avoids a race condition where the parent thread frees the buffer
-    // before the child has accepted the transferable.
-    if (!rawbuf->addReference()) {
-        JS_ReportErrorNumberASCII(context(), GetErrorMessage, nullptr, JSMSG_SC_SAB_REFCNT_OFLO);
+    if (!refsHeld.acquire(context(), rawbuf))
         return false;
-    }
 
     intptr_t p = reinterpret_cast<intptr_t>(rawbuf);
     return out.writePair(SCTAG_SHARED_ARRAY_BUFFER_OBJECT, static_cast<uint32_t>(sizeof(p))) &&
@@ -1894,18 +1959,24 @@ JSStructuredCloneReader::readSharedArrayBuffer(uint32_t nbytes, MutableHandleVal
     // in any case.  Just fail at the receiving end if we can't handle it.
 
     if (!context()->compartment()->creationOptions().getSharedMemoryAndAtomicsEnabled()) {
-        // The sending side performed a reference increment before sending.
-        // Account for that here before leaving.
-        if (rawbuf)
-            rawbuf->dropReference();
-
         JS_ReportErrorNumberASCII(context(), GetErrorMessage, nullptr, JSMSG_SC_SAB_DISABLED);
         return false;
     }
 
-    // The constructor absorbs the reference count increment performed by the sender.
+    // The new object will have a new reference to the rawbuf.
+
+    if (!rawbuf->addReference()) {
+        JS_ReportErrorNumberASCII(context(), GetErrorMessage, nullptr, JSMSG_SC_SAB_REFCNT_OFLO);
+        return false;
+    }
+
     JSObject* obj = SharedArrayBufferObject::New(context(), rawbuf);
 
+    if (!obj) {
+        rawbuf->dropReference();
+        return false;
+    }
+
     vp.setObject(*obj);
     return true;
 }
@@ -2591,13 +2662,14 @@ JSAutoStructuredCloneBuffer::clear(const JSStructuredCloneCallbacks* optionalCal
     if (data_.ownTransferables_ == OwnTransferablePolicy::OwnsTransferablesIfAny)
         DiscardTransferables(data_, callbacks, closure);
     data_.ownTransferables_ = OwnTransferablePolicy::NoTransferables;
+    data_.refsHeld_.releaseAll();
     data_.Clear();
     version_ = 0;
 }
 
 bool
-JSAutoStructuredCloneBuffer::copy(const JSStructuredCloneData& srcData, uint32_t version,
-                                  const JSStructuredCloneCallbacks* callbacks,
+JSAutoStructuredCloneBuffer::copy(JSContext* cx, const JSStructuredCloneData& srcData,
+                                  uint32_t version, const JSStructuredCloneCallbacks* callbacks,
                                   void* closure)
 {
     // transferable objects cannot be copied
@@ -2608,11 +2680,16 @@ JSAutoStructuredCloneBuffer::copy(const JSStructuredCloneData& srcData, uint32_t
 
     auto iter = srcData.Iter();
     while (!iter.Done()) {
-            data_.WriteBytes(iter.Data(), iter.RemainingInSegment());
-            iter.Advance(srcData, iter.RemainingInSegment());
+        if (!data_.WriteBytes(iter.Data(), iter.RemainingInSegment()))
+            return false;
+        iter.Advance(srcData, iter.RemainingInSegment());
     }
 
     version_ = version;
+
+    if (!data_.refsHeld_.acquireAll(cx, srcData.refsHeld_))
+        return false;
+
     data_.setOptionalCallbacks(callbacks, closure, OwnTransferablePolicy::NoTransferables);
     return true;
 }