fix: lock onmessage to prevent duplicate message handlers

src/webworker/js.rs

@@ -4,25 +4,46 @@ use wasm_bindgen::prelude::wasm_bindgen;
 /// The initialization code for the worker,
 /// which will be loaded as a blob.
 ///
+/// All wasmworker traffic (init handshake and task dispatch) runs over a
+/// dedicated `MessageChannel` port, which is transferred with the first
+/// message to the worker. This keeps the worker's global message channel
+/// free for the embedded module, so message handlers installed by module
+/// code (e.g. in a `#[wasm_bindgen(start)]` function) never interfere
+/// with task dispatch, and messages posted by the module on the global
+/// scope never reach wasmworker's response callback.
+///
 /// `{{wasm}}` will be replaced later by an actual path.
 pub(crate) const WORKER_JS: &str = r#"
 console.debug('Initializing worker');
 
+// Capture the dedicated task port before any module code can run.
+const portPromise = new Promise(resolve => {
+    const initListener = event => {
+        if (event.data && event.data.type === 'init_port') {
+            self.removeEventListener('message', initListener);
+            resolve(event.ports[0]);
+        }
+    };
+    self.addEventListener('message', initListener);
+});
+
 (async () => {
+    const port = await portPromise;
+
     let mod;
     try {
         mod = await import('{{wasm}}');
     } catch (e) {
         console.error('Unable to import module {{wasm}}', e);
-        self.postMessage({ success: false, message: e.toString() });
+        port.postMessage({ success: false, message: e.toString() });
         return;
     }
 
     await mod.default({{wasm_bg}});
-    self.postMessage({ success: true });
+    port.postMessage({ success: true });
     console.debug('Worker started');
 
-    self.addEventListener('message', async event => {
+    port.onmessage = async event => {
         console.debug('Received worker event');
         const { id, func_name, is_channel, arg } = event.data;
 
@@ -31,16 +52,16 @@ console.debug('Initializing worker');
         const fn = mod[webworker_func_name];
         if (!fn) {
             console.error(`Function '${func_name}' is not exported.`);
-            self.postMessage({ id: id, response: null });
+            port.postMessage({ id: id, response: null });
             return;
         }
 
         const worker_result = await fn(arg, event.ports[0]);
 
         // Send response back to be handled by callback in main thread.
         console.debug('Send worker result');
-        self.postMessage({ id: id, response: worker_result });
-    });
+        port.postMessage({ id: id, response: worker_result });
+    };
 })();
 "#;
 
@@ -56,39 +77,42 @@ pub(crate) fn main_js() -> JsString {
     URL.with(Clone::clone)
 }
 
-/// The initialization code for workers that receive a pre-compiled WASM module
+/// The initialization code for workers that receive a pre-compiled WASM module.
+///
+/// Like [`WORKER_JS`], all wasmworker traffic runs over a dedicated
+/// `MessageChannel` port, which arrives with the `wasm_module` init message.
 pub(crate) const WORKER_JS_WITH_PRECOMPILED: &str = r#"
 console.debug('Initializing worker with pre-compiled WASM');
 
-let wasmModule = null;
 let mod = null;
 let initHandler = null;
 
-// Listen for the pre-compiled WASM module
+// Listen for the pre-compiled WASM module and the dedicated task port
 initHandler = async function(event) {
     const data = event.data;
 
     if (data.type === 'wasm_module') {
         console.debug('Received pre-compiled WASM module');
-        wasmModule = data.module;
+        const port = event.ports[0];
+
+        // Remove this listener before running module code, so wasmworker
+        // no longer listens on the global scope at all.
+        self.removeEventListener('message', initHandler);
 
         // Now initialize with the pre-compiled module
         try {
             mod = await import('{{wasm}}');
-            await mod.default({ module_or_path: wasmModule });
-            self.postMessage({ success: true });
+            await mod.default({ module_or_path: data.module });
+            port.postMessage({ success: true });
             console.debug('Worker started with pre-compiled WASM');
         } catch (e) {
             console.error('Unable to initialize with pre-compiled WASM', e);
-            self.postMessage({ success: false, message: e.toString() });
+            port.postMessage({ success: false, message: e.toString() });
             return;
         }
 
-        // Remove this listener and add the task handler
-        self.removeEventListener('message', initHandler);
-
         // Add the main message handler for tasks
-        self.addEventListener('message', async event => {
+        port.onmessage = async event => {
             console.debug('Received worker event');
             const { id, func_name, is_channel, arg } = event.data;
 
@@ -97,16 +121,16 @@ initHandler = async function(event) {
             const fn = mod[webworker_func_name];
             if (!fn) {
                 console.error(`Function '${func_name}' is not exported.`);
-                self.postMessage({ id: id, response: null });
+                port.postMessage({ id: id, response: null });
                 return;
             }
 
             const worker_result = await fn(arg, event.ports[0]);
 
             // Send response back to be handled by callback in main thread.
             console.debug('Send worker result');
-            self.postMessage({ id: id, response: worker_result });
-        });
+            port.postMessage({ id: id, response: worker_result });
+        };
     }
 };
 

src/webworker/worker.rs

@@ -51,6 +51,12 @@ type Callback = dyn FnMut(MessageEvent);
 pub struct WebWorker {
     /// The underlying web worker.
     worker: Worker,
+    /// The dedicated port for all wasmworker traffic (init handshake and
+    /// task dispatch). Using a `MessageChannel` instead of the worker's
+    /// global message channel keeps the latter free for the embedded module,
+    /// so message handlers installed by module code never interfere with
+    /// task dispatch and vice versa.
+    port: MessagePort,
     /// An optional limit on the number of tasks queued at the same time.
     task_limit: Option<Semaphore>,
     /// The current task id, which is used to reidentify responses.
@@ -146,9 +152,17 @@ impl WebWorker {
         let worker = Worker::new_with_options(&script_url, &worker_options)
             .map_err(InitError::WebWorkerCreation)?;
 
-        // Send pre-compiled WASM module if provided
+        // Create the dedicated channel for all wasmworker traffic. One port
+        // stays on the main thread, the other is transferred to the worker
+        // with the init message.
+        let channel = MessageChannel::new().map_err(InitError::ChannelCreation)?;
+        let port = channel.port1();
+        let worker_port = channel.port2();
+
+        // Send the init message with the task port
+        // (and the pre-compiled WASM module if provided).
+        let init_msg = js_sys::Object::new();
         if let Some(module) = wasm_module {
-            let init_msg = js_sys::Object::new();
             js_sys::Reflect::set(
                 &init_msg,
                 &JsValue::from_str("type"),
@@ -157,12 +171,21 @@ impl WebWorker {
             .expect_throw("Could not set type");
             js_sys::Reflect::set(&init_msg, &JsValue::from_str("module"), &module)
                 .expect_throw("Could not set module");
-
-            worker
-                .post_message(&init_msg)
-                .expect_throw("Could not send WASM module to worker");
+        } else {
+            js_sys::Reflect::set(
+                &init_msg,
+                &JsValue::from_str("type"),
+                &JsValue::from_str("init_port"),
+            )
+            .expect_throw("Could not set type");
         }
 
+        let transfer = Array::new();
+        transfer.push(&worker_port);
+        worker
+            .post_message_with_transfer(&init_msg, &transfer)
+            .expect_throw("Could not send init message to worker");
+
         // Wait until worker is initialized.
         let (tx, rx) = oneshot::channel();
         let handler = Closure::once(move |event: MessageEvent| {
@@ -171,7 +194,7 @@ impl WebWorker {
                 .expect_throw("Error deserializing post init data");
             let _ = tx.send(post_init);
         });
-        worker.set_onmessage(Some(handler.as_ref().unchecked_ref()));
+        port.set_onmessage(Some(handler.as_ref().unchecked_ref()));
         let post_init = rx.await.expect_throw("WebWorker init sender dropped");
 
         // Handle errors in webworker init
@@ -187,10 +210,11 @@ impl WebWorker {
         let last_active = Rc::new(Cell::new(js_sys::Date::now()));
 
         let callback_handle = Self::callback(Rc::clone(&tasks), Rc::clone(&last_active));
-        worker.set_onmessage(Some(callback_handle.as_ref().unchecked_ref()));
+        port.set_onmessage(Some(callback_handle.as_ref().unchecked_ref()));
 
         Ok(WebWorker {
             worker,
+            port,
             task_limit: task_limit.map(|limit| Semaphore::new(limit)),
             current_task: AtomicU32::new(0),
             open_tasks: tasks,
@@ -447,15 +471,15 @@ impl WebWorker {
             let transfer = Array::new();
             transfer.push(&port);
 
-            self.worker
-                .post_message_with_transfer(
+            self.port
+                .post_message_with_transferable(
                     &serde_wasm_bindgen::to_value(&request)
                         .expect_throw("Could not serialize request"),
                     &transfer,
                 )
                 .expect_throw("WebWorker gone");
         } else {
-            self.worker
+            self.port
                 .post_message(
                     &serde_wasm_bindgen::to_value(&request)
                         .expect_throw("Could not serialize request"),
@@ -497,8 +521,8 @@ impl WebWorker {
         let transfer = Array::new();
         transfer.push(&port);
 
-        self.worker
-            .post_message_with_transfer(
+        self.port
+            .post_message_with_transferable(
                 &serde_wasm_bindgen::to_value(&request).expect_throw("Could not serialize request"),
                 &transfer,
             )
@@ -534,6 +558,7 @@ impl WebWorker {
 
 impl Drop for WebWorker {
     fn drop(&mut self) {
+        self.port.close();
         self.worker.terminate();
     }
 }

test/Cargo.toml

@@ -21,7 +21,7 @@ serde = { version = "1.0", features = ["derive"] }
 wasm-bindgen = "0.2"
 js-sys = "0.3"
 wasm-bindgen-futures = "0.4"
-web-sys = { version = "0.3", features = ["MessagePort", "Window"] }
+web-sys = { version = "0.3", features = ["DedicatedWorkerGlobalScope", "MessageEvent", "MessagePort", "Window"] }
 wasmworker = { workspace = true }
 
 [features]

test/src/lib.rs

@@ -1,10 +1,12 @@
 use channel::*;
 use convert::*;
+use onmessage::*;
 use raw::*;
 use wasm_bindgen::prelude::wasm_bindgen;
 
 pub(crate) mod channel;
 pub(crate) mod convert;
+pub(crate) mod onmessage;
 pub(crate) mod raw;
 
 #[macro_export]
@@ -46,4 +48,7 @@ pub async fn run_tests() {
 
     // Idle timeout test
     can_use_idle_timeout().await;
+
+    // Conflicting onmessage regression test
+    can_run_task_with_conflicting_onmessage().await;
 }