Fix race condition in WASM crypto worker (#70185)

When sending a message between LibraryChannel and ChannelWorker, there is a race condition where both threads are reading/writing to shared memory at the same time. This can cause message pages to be skipped.

To fix this, add a shared mutex lock so only one side is reading/writing to shared memory at the same time.

Fix #69806

2 files changed, 62 insertions, 3 deletions

diff --git a/src/mono/wasm/runtime/crypto-worker.ts b/src/mono/wasm/runtime/crypto-worker.ts
index ea7bd9e6fce..f17a455892d 100644
--- a/src/mono/wasm/runtime/crypto-worker.ts
+++ b/src/mono/wasm/runtime/crypto-worker.ts
@@ -65,10 +65,15 @@ class LibraryChannel {
     private comm: Int32Array;
     private msg: Uint16Array;
 
+    // LOCK states
+    private get LOCK_UNLOCKED(): number { return 0; }  // 0 means the lock is unlocked
+    private get LOCK_OWNED(): number { return 1; } // 1 means the LibraryChannel owns the lock
+
     // Index constants for the communication buffer.
     private get STATE_IDX(): number { return 0; }
     private get MSG_SIZE_IDX(): number { return 1; }
-    private get COMM_LAST_IDX(): number { return this.MSG_SIZE_IDX; }
+    private get LOCK_IDX(): number { return 2; }
+    private get COMM_LAST_IDX(): number { return this.LOCK_IDX; }
 
     // Communication states.
     private get STATE_SHUTDOWN(): number { return -1; } // Shutdown
@@ -125,6 +130,8 @@ class LibraryChannel {
         let msg_written = 0;
 
         for (; ;) {
+            this.acquire_lock();
+
             // Write the message and return how much was written.
             const wrote = this.write_to_msg(msg, msg_written, msg_len);
             msg_written += wrote;
@@ -137,6 +144,9 @@ class LibraryChannel {
 
             // Notify webworker
             Atomics.store(this.comm, this.STATE_IDX, state);
+
+            this.release_lock();
+
             Atomics.notify(this.comm, this.STATE_IDX);
 
             // The send message is complete.
@@ -172,6 +182,8 @@ class LibraryChannel {
                 state = Atomics.load(this.comm, this.STATE_IDX);
             } while (state !== this.STATE_RESP && state !== this.STATE_RESP_P);
 
+            this.acquire_lock();
+
             const size_to_read = Atomics.load(this.comm, this.MSG_SIZE_IDX);
 
             // Append the latest part of the message.
@@ -179,12 +191,16 @@ class LibraryChannel {
 
             // The response is complete.
             if (state === this.STATE_RESP) {
+                this.release_lock();
                 break;
             }
 
             // Reset the size and transition to await state.
             Atomics.store(this.comm, this.MSG_SIZE_IDX, 0);
             Atomics.store(this.comm, this.STATE_IDX, this.STATE_AWAIT);
+
+            this.release_lock();
+
             Atomics.notify(this.comm, this.STATE_IDX);
         }
 
@@ -202,6 +218,19 @@ class LibraryChannel {
         return String.fromCharCode.apply(null, slicedMessage);
     }
 
+    private acquire_lock() {
+        while (Atomics.compareExchange(this.comm, this.LOCK_IDX, this.LOCK_UNLOCKED, this.LOCK_OWNED) !== this.LOCK_UNLOCKED) {
+            // empty
+        }
+    }
+
+    private release_lock() {
+        const result = Atomics.compareExchange(this.comm, this.LOCK_IDX, this.LOCK_OWNED, this.LOCK_UNLOCKED);
+        if (result !== this.LOCK_OWNED) {
+            throw "CRYPTO: LibraryChannel tried to release a lock that wasn't acquired: " + result;
+        }
+    }
+
     public static create(msg_char_len: number): LibraryChannel {
         if (msg_char_len === undefined) {
             msg_char_len = 1024; // Default size is arbitrary but is in 'char' units (i.e. UTF-16 code points).
diff --git a/src/mono/wasm/runtime/dotnet-crypto-worker.js b/src/mono/wasm/runtime/dotnet-crypto-worker.js
index c6416492a71..0110693a1a4 100644
--- a/src/mono/wasm/runtime/dotnet-crypto-worker.js
+++ b/src/mono/wasm/runtime/dotnet-crypto-worker.js
@@ -3,9 +3,14 @@
 
 var ChannelWorker = {
     _impl: class {
+        // LOCK states
+        get LOCK_UNLOCKED() { return 0; } // 0 means the lock is unlocked
+        get LOCK_OWNED() { return 2; } // 2 means the ChannelWorker owns the lock
+
         // BEGIN ChannelOwner contract - shared constants.
         get STATE_IDX() { return 0; }
         get MSG_SIZE_IDX() { return 1; }
+        get LOCK_IDX() { return 2; }
 
         // Communication states.
         get STATE_SHUTDOWN() { return -1; } // Shutdown
@@ -51,6 +56,8 @@ var ChannelWorker = {
         _read_request() {
             var request = "";
             for (;;) {
+                this._acquire_lock();
+
                 // Get the current state and message size
                 var state = Atomics.load(this.comm, this.STATE_IDX);
                 var size_to_read = Atomics.load(this.comm, this.MSG_SIZE_IDX);
@@ -59,16 +66,22 @@ var ChannelWorker = {
                 request += this._read_from_msg(0, size_to_read);
 
                 // The request is complete.
-                if (state === this.STATE_REQ)
+                if (state === this.STATE_REQ) {
+                    this._release_lock();
                     break;
+                }
 
                 // Shutdown the worker.
-                if (state === this.STATE_SHUTDOWN)
+                if (state === this.STATE_SHUTDOWN) {
+                    this._release_lock();
                     return this.STATE_SHUTDOWN;
+                }
 
                 // Reset the size and transition to await state.
                 Atomics.store(this.comm, this.MSG_SIZE_IDX, 0);
                 Atomics.store(this.comm, this.STATE_IDX, this.STATE_AWAIT);
+                this._release_lock();
+
                 Atomics.wait(this.comm, this.STATE_IDX, this.STATE_AWAIT);
             }
 
@@ -88,6 +101,8 @@ var ChannelWorker = {
             var msg_written = 0;
 
             for (;;) {
+                this._acquire_lock();
+
                 // Write the message and return how much was written.
                 var wrote = this._write_to_msg(msg, msg_written, msg_len);
                 msg_written += wrote;
@@ -101,6 +116,8 @@ var ChannelWorker = {
                 // Update the state
                 Atomics.store(this.comm, this.STATE_IDX, state);
 
+                this._release_lock();
+
                 // Wait for the transition to know the main thread has
                 // received the response by moving onto a new state.
                 Atomics.wait(this.comm, this.STATE_IDX, state);
@@ -121,6 +138,19 @@ var ChannelWorker = {
             }
             return ii - start;
         }
+
+        _acquire_lock() {
+            while (Atomics.compareExchange(this.comm, this.LOCK_IDX, this.LOCK_UNLOCKED, this.LOCK_OWNED) !== this.LOCK_UNLOCKED) {
+                // empty
+            }
+        }
+
+        _release_lock() {
+            const result = Atomics.compareExchange(this.comm, this.LOCK_IDX, this.LOCK_OWNED, this.LOCK_UNLOCKED);
+            if (result !== this.LOCK_OWNED) {
+                throw "CRYPTO: ChannelWorker tried to release a lock that wasn't acquired: " + result;
+            }
+        }
     },
 
     create: function (comm_buf, msg_buf, msg_char_len) {