Guard MessagePortChannel state with a per-channel Lock

src/bun.js/bindings/webcore/MessagePortChannel.cpp

@@ -29,6 +29,7 @@
 // #include "Logging.h"
 #include "MessagePortChannelRegistry.h"
 #include <wtf/CompletionHandler.h>
+#include <wtf/Locker.h>
 #include <wtf/MainThread.h>
 
 namespace WebCore {
@@ -42,12 +43,13 @@ MessagePortChannel::MessagePortChannel(MessagePortChannelRegistry& registry, con
     : m_ports { port1, port2 }
     , m_registry(registry)
 {
-    relaxAdoptionRequirement();
-
-    m_processes[0] = port1.processIdentifier;
-    m_entangledToProcessProtectors[0] = this;
-    m_processes[1] = port2.processIdentifier;
-    m_entangledToProcessProtectors[1] = this;
+    {
+        Locker locker { m_lock };
+        m_processes[0] = port1.processIdentifier;
+        m_entangledToProcessProtectors[0] = this;
+        m_processes[1] = port2.processIdentifier;
+        m_entangledToProcessProtectors[1] = this;
+    }
 
     m_registry.messagePortChannelCreated(*this);
 }
@@ -61,6 +63,8 @@ std::optional<ProcessIdentifier> MessagePortChannel::processForPort(const Messag
 {
     ASSERT(port == m_ports[0] || port == m_ports[1]);
     size_t i = port == m_ports[0] ? 0 : 1;
+
+    Locker locker { m_lock };
     return m_processes[i];
 }
 
@@ -76,6 +80,7 @@ void MessagePortChannel::entanglePortWithProcess(const MessagePortIdentifier& po
 
     // LOG(MessagePorts, "MessagePortChannel %s (%p) entangling port %s (that port has %zu messages available)", logString().utf8().data(), this, port.logString().utf8().data(), m_pendingMessages[i].size());
 
+    Locker locker { m_lock };
     ASSERT(!m_processes[i] || *m_processes[i] == process);
     m_processes[i] = process;
     m_entangledToProcessProtectors[i] = this;
@@ -89,38 +94,49 @@ void MessagePortChannel::disentanglePort(const MessagePortIdentifier& port)
     ASSERT(port == m_ports[0] || port == m_ports[1]);
     size_t i = port == m_ports[0] ? 0 : 1;
 
-    ASSERT(m_processes[i] || m_isClosed[i]);
-    m_processes[i] = std::nullopt;
-    m_pendingMessagePortTransfers[i].add(this);
-
     // This set of steps is to guarantee that the lock is unlocked before the
     // last ref to this object is released.
-    auto protectedThis = WTF::move(m_entangledToProcessProtectors[i]);
+    RefPtr<MessagePortChannel> protectedThis;
+    {
+        Locker locker { m_lock };
+        ASSERT(m_processes[i] || m_isClosed[i]);
+        m_processes[i] = std::nullopt;
+        m_pendingMessagePortTransfers[i].add(this);
+        protectedThis = WTF::move(m_entangledToProcessProtectors[i]);
+    }
 }
 
 void MessagePortChannel::closePort(const MessagePortIdentifier& port)
 {
     ASSERT(port == m_ports[0] || port == m_ports[1]);
     size_t i = port == m_ports[0] ? 0 : 1;
 
-    m_processes[i] = std::nullopt;
-    m_isClosed[i] = true;
-
     // This set of steps is to guarantee that the lock is unlocked before the
     // last ref to this object is released.
     Ref protectedThis { *this };
-
-    m_pendingMessages[i].clear();
-    m_pendingMessagePortTransfers[i].clear();
-    m_pendingMessageProtectors[i] = nullptr;
-    m_entangledToProcessProtectors[i] = nullptr;
+    Vector<MessageWithMessagePorts> pendingMessages;
+    UncheckedKeyHashSet<RefPtr<MessagePortChannel>> pendingMessagePortTransfers;
+    RefPtr<MessagePortChannel> pendingMessageProtector;
+    RefPtr<MessagePortChannel> entangledToProcessProtector;
+    {
+        Locker locker { m_lock };
+        m_processes[i] = std::nullopt;
+        m_isClosed[i] = true;
+
+        pendingMessages = WTF::move(m_pendingMessages[i]);
+        pendingMessagePortTransfers = WTF::move(m_pendingMessagePortTransfers[i]);
+        pendingMessageProtector = WTF::move(m_pendingMessageProtectors[i]);
+        entangledToProcessProtector = WTF::move(m_entangledToProcessProtectors[i]);
+    }
 }
 
 bool MessagePortChannel::postMessageToRemote(MessageWithMessagePorts&& message, const MessagePortIdentifier& remoteTarget)
 {
     ASSERT(remoteTarget == m_ports[0] || remoteTarget == m_ports[1]);
     size_t i = remoteTarget == m_ports[0] ? 0 : 1;
 
+    Locker locker { m_lock };
+
     if (m_isClosed[i])
         return false;
 
@@ -143,22 +159,30 @@ void MessagePortChannel::takeAllMessagesForPort(const MessagePortIdentifier& por
     ASSERT(port == m_ports[0] || port == m_ports[1]);
     size_t i = port == m_ports[0] ? 0 : 1;
 
-    if (m_pendingMessages[i].isEmpty()) {
-        callback({}, [] {});
-        return;
-    }
+    Vector<MessageWithMessagePorts> result;
+    RefPtr<MessagePortChannel> protectedThis;
+    {
+        Locker locker { m_lock };
 
-    ASSERT(m_pendingMessageProtectors[i]);
+        if (m_pendingMessages[i].isEmpty()) {
+            locker.unlockEarly();
+            callback({}, [] {});
+            return;
+        }
 
-    Vector<MessageWithMessagePorts> result;
-    result.swap(m_pendingMessages[i]);
+        ASSERT(m_pendingMessageProtectors[i]);
 
-    ++m_messageBatchesInFlight;
+        result.swap(m_pendingMessages[i]);
+        protectedThis = WTF::move(m_pendingMessageProtectors[i]);
+
+        ++m_messageBatchesInFlight;
+    }
 
     // LOG(MessagePorts, "There are %zu messages to take for port %s. Taking them now, messages in flight is now %" PRIu64, result.size(), port.logString().utf8().data(), m_messageBatchesInFlight);
 
-    callback(WTF::move(result), [this, port, protectedThis = WTF::move(m_pendingMessageProtectors[i])] {
+    callback(WTF::move(result), [this, port, protectedThis = WTF::move(protectedThis)] {
         UNUSED_PARAM(port);
+        Locker locker { m_lock };
         --m_messageBatchesInFlight;
         // LOG(MessagePorts, "Message port channel %s was notified that a batch of %zu message port messages targeted for port %s just completed dispatch, in flight is now %" PRIu64, logString().utf8().data(), size, port.logString().utf8().data(), m_messageBatchesInFlight);
     });
@@ -169,6 +193,8 @@ std::optional<MessageWithMessagePorts> MessagePortChannel::tryTakeMessageForPort
     ASSERT(port == m_ports[0] || port == m_ports[1]);
     size_t i = port == m_ports[0] ? 0 : 1;
 
+    Locker locker { m_lock };
+
     if (m_pendingMessages[i].isEmpty())
         return std::nullopt;
 

src/bun.js/bindings/webcore/MessagePortChannel.h

@@ -29,16 +29,17 @@
 #include "MessagePortIdentifier.h"
 #include "MessageWithMessagePorts.h"
 #include "ProcessIdentifier.h"
+#include <wtf/CanMakeWeakPtr.h>
 #include <wtf/HashSet.h>
-#include <wtf/RefCounted.h>
+#include <wtf/Lock.h>
+#include <wtf/ThreadSafeRefCounted.h>
 #include <wtf/text/WTFString.h>
-#include <wtf/RefCountedAndCanMakeWeakPtr.h>
 
 namespace WebCore {
 
 class MessagePortChannelRegistry;
 
-class MessagePortChannel : public RefCountedAndCanMakeWeakPtr<MessagePortChannel> {
+class MessagePortChannel : public ThreadSafeRefCounted<MessagePortChannel>, public CanMakeWeakPtr<MessagePortChannel> {
 public:
     static Ref<MessagePortChannel> create(MessagePortChannelRegistry&, const MessagePortIdentifier& port1, const MessagePortIdentifier& port2);
 
@@ -71,6 +72,11 @@ class MessagePortChannel : public RefCountedAndCanMakeWeakPtr<MessagePortChannel
 private:
     MessagePortChannel(MessagePortChannelRegistry&, const MessagePortIdentifier& port1, const MessagePortIdentifier& port2);
 
+    // Upstream WebKit serializes all access to this class onto the main thread. Bun calls into
+    // it from both the main thread and worker threads (see MessagePortChannelRegistry.cpp), so
+    // all mutable state below must be guarded by m_lock.
+    mutable Lock m_lock;
+
     MessagePortIdentifier m_ports[2];
     bool m_isClosed[2] { false, false };
     std::optional<ProcessIdentifier> m_processes[2];

test/js/web/workers/message-channel.test.ts

@@ -1,3 +1,5 @@
+import { bunEnv, bunExe, isASAN, isDebug, tempDir } from "harness";
+
 test("simple usage", done => {
   const channel = new MessageChannel();
   const port1 = channel.port1;
@@ -323,3 +325,79 @@
   mc.port2.postMessage(blocklist);
   await promise;
 });
+
+// MessagePortChannel::m_pendingMessages is appended on the sender thread (postMessageToRemote)
+// and swapped/drained on the receiver thread (takeAllMessagesForPort). Without a per-channel
+// lock, Vector::append can reallocate the backing buffer while the other thread is reading it,
+// which ASAN reports as container-overflow / heap-use-after-free and the non-atomic RefCounted
+// refcount is corrupted. This test hammers that path from both directions.
+// Debug builds have ASAN enabled; release builds race silently, so skip there.
+test.skipIf(!(isDebug || isASAN))(
+  "concurrent MessagePort postMessage/onmessage across threads does not race",
+  async () => {
+    using dir = tempDir("message-port-race", {
+      "worker.js": `
+        self.onmessage = (e) => {
+          const port = e.data;
+          let got = 0;
+          port.onmessage = (ev) => {
+            if (ev.data === "done") {
+              port.postMessage("worker-done");
+              port.close();
+            } else {
+              got++;
+            }
+          };
+          for (let i = 0; i < 20000; i++) port.postMessage(i);
+          port.postMessage("flood-done");
+        };
+      `,
+      "main.js": `
+        const worker = new Worker(new URL("./worker.js", import.meta.url).href);
+        const { port1, port2 } = new MessageChannel();
+
+        let received = 0;
+
+        port1.onmessage = (e) => {
+          if (e.data === "flood-done") {
+            port1.postMessage("done");
+          } else if (e.data === "worker-done") {
+            console.log("received=" + received);
+            worker.terminate();
+            port1.close();
+          } else {
+            received++;
+            // Echo back while the worker is still flooding us so both threads are appending
+            // to and draining from the shared MessagePortChannel concurrently.
+            port1.postMessage(e.data);
+          }
+        };
+
+        worker.onerror = (e) => {
+          console.error("worker error", e.message);
+          process.exit(1);
+        };
+
+        worker.postMessage(port2, [port2]);
+      `,
+    });
+
+    // The race is probabilistic; three attempts brings the false-pass rate from ~10% to ~0.1%.
+    for (let attempt = 0; attempt < 3; attempt++) {
+      await using proc = Bun.spawn({
+        cmd: [bunExe(), "main.js"],
+        env: bunEnv,
+        cwd: String(dir),
+        stdout: "pipe",
+        stderr: "pipe",
+      });
+
+      const [stdout, stderr, exitCode] = await Promise.all([proc.stdout.text(), proc.stderr.text(), proc.exited]);
+
+      expect(stderr).toBe("");
+      expect(stdout.trim()).toBe("received=20000");
+      expect(exitCode).toBe(0);
+    }
+  },
+  120_000,
+);