fix: toctou in wasm apply_patch causing UB (#5534)

packages/subsecond/subsecond/src/lib.rs

@@ -551,7 +551,7 @@ pub unsafe fn apply_patch(mut table: JumpTable) -> Result<(), PatchError> {
     wasm_bindgen_futures::spawn_local(async move {
         use js_sys::{
             ArrayBuffer, Object, Reflect,
-            WebAssembly::{self, Memory, Table},
+            WebAssembly::{self, Instance, Memory, Module, Table},
         };
         use wasm_bindgen::prelude::*;
         use wasm_bindgen::JsValue;
@@ -561,20 +561,19 @@ pub unsafe fn apply_patch(mut table: JumpTable) -> Result<(), PatchError> {
         let funcs: Table = wasm_bindgen::function_table().unchecked_into();
         let memory: Memory = wasm_bindgen::memory().unchecked_into();
         let exports: Object = wasm_bindgen::exports().unchecked_into();
-        let buffer: ArrayBuffer = memory.buffer().unchecked_into();
 
         let path = table.lib.to_str().unwrap();
         if !path.ends_with(".wasm") {
             return;
         }
 
-        // Start the fetch of the module
+        // Fetch + decode the patch wasm. Both awaits are pure I/O — they
-        let response = web_sys::window().unwrap_throw().fetch_with_str(&path);
+        // touch no shared state, so the future is safe to drop here.
-
+        let response: web_sys::Response =
-        // Wait for the fetch to complete - we need the wasm module size in bytes to reserve in the memory
+            JsFuture::from(web_sys::window().unwrap_throw().fetch_with_str(&path))
-        let response: web_sys::Response = JsFuture::from(response).await.unwrap().unchecked_into();
+                .await
-
+                .unwrap()
-        // If the status is not success, we bail
+                .unchecked_into();
         if !response.ok() {
             panic!(
                 "Failed to patch wasm module at {} - response failed with: {}",
@@ -582,100 +581,94 @@ pub unsafe fn apply_patch(mut table: JumpTable) -> Result<(), PatchError> {
                 response.status_text()
             );
         }
-
         let dl_bytes: ArrayBuffer = JsFuture::from(response.array_buffer().unwrap())
             .await
             .unwrap()
             .unchecked_into();
 
-        // Expand the memory and table size to accommodate the new data and functions
+        // Pre-compile. This is the slow part — V8 hands it to a worker —
+        // and yields the longest. The result is a Module with no side
+        // effects on host JS state, so cancelling here is also safe.
+        let module: Module = JsFuture::from(WebAssembly::compile(dl_bytes.unchecked_ref()))
+            .await
+            .unwrap()
+            .unchecked_into();
+
+        // ── SYNC SECTION — no `.await` past this point ────────────────
         //
-        // Normally we wouldn't be able to trust that we are allocating *enough* memory
+        // Everything below mutates shared JS state (linear memory, the
-        // for BSS segments, but ld emits them in the binary when using import-memory.
+        // indirect function table, the patch's data segment) and the
+        // jump-table pointer. We must not yield, because:
         //
-        // Make sure we align the memory base to the page size
+        //  * `memory.grow` detaches the previous `ArrayBuffer`. The old
+        //    code captured `memory.buffer()` BEFORE the awaits and then
+        //    read its `byteLength` to compute `__memory_base`; if any
+        //    other task ran a `memory.grow` during the awaits (e.g. an
+        //    allocator call from a concurrent render), the snapshot was
+        //    detached, `byteLength` returned 0, and the patch's data
+        //    segment ended up loaded over the host's stack/rodata.
+        //  * Two concurrent `apply_patch` tasks must not interleave
+        //    between memory/table grow and `commit_patch`. Doing all
+        //    side-effecting work in one synchronous turn guarantees that.
+        //  * If the future is cancelled before this point, host state
+        //    is untouched.
         const PAGE_SIZE: u32 = 64 * 1024;
-        let page_count = (buffer.byte_length() as f64 / PAGE_SIZE as f64).ceil() as u32;
+        let patch_pages = (dl_bytes.byte_length() as f64 / PAGE_SIZE as f64).ceil() as u32 + 1;
-        let memory_base = (page_count + 1) * PAGE_SIZE;
 
-        // We need to grow the memory to accommodate the new module
+        // Use grow's return value (the prior page count) to derive
-        memory.grow((dl_bytes.byte_length() as f64 / PAGE_SIZE as f64).ceil() as u32 + 1);
+        // memory_base. Atomic w.r.t. concurrent grows, unlike reading
+        // memory.buffer().byteLength().
+        let prev_pages = memory.grow(patch_pages);
+        let memory_base = (prev_pages + 1) * PAGE_SIZE;
 
-        // We grow the ifunc table to accommodate the new functions
+        // grow returns the prior table length, which is __table_base.
-        // In theory we could just put all the ifuncs in the jump map and use that for our count,
-        // but there's no guarantee from the jump table that it references "itself"
-        // We might need a sentinel value for each ifunc in the jump map to indicate that it is
         let table_base = funcs.grow(table.ifunc_count as u32).unwrap();
 
-        // Adjust the jump table to be relative to the new base address
+        // Rebase the jump table entries onto the patch's table slot range.
         for v in table.map.values_mut() {
             *v += table_base as u64;
         }
 
-        // Build up the import object. We copy everything over from the current exports, but then
+        // Build the env import object: copy every host export through and
-        // need to add in the memory and table base offsets for the relocations to work.
+        // add __memory_base / __table_base globals so the patch's PIC
-        //
+        // code resolves correctly.
-        // let imports = {
-        //     env: {
-        //         memory: base.memory,
-        //         __tls_base: base.__tls_base,
-        //         __stack_pointer: base.__stack_pointer,
-        //         __indirect_function_table: base.__indirect_function_table,
-        //         __memory_base: memory_base,
-        //         __table_base: table_base,
-        //        ..base_exports
-        //     },
-        // };
         let env = Object::new();
-
-        // Move memory, __tls_base, __stack_pointer, __indirect_function_table, and all exports over
         for key in Object::keys(&exports) {
             Reflect::set(&env, &key, &Reflect::get(&exports, &key).unwrap()).unwrap();
         }
-
-        // Set the memory and table in the imports
-        // Following this pattern: Global.new({ value: "i32", mutable: false }, value)
         for (name, value) in [("__table_base", table_base), ("__memory_base", memory_base)] {
             let descriptor = Object::new();
             Reflect::set(&descriptor, &"value".into(), &"i32".into()).unwrap();
             Reflect::set(&descriptor, &"mutable".into(), &false.into()).unwrap();
-            let value = WebAssembly::Global::new(&descriptor, &value.into()).unwrap();
+            let global = WebAssembly::Global::new(&descriptor, &value.into()).unwrap();
-            Reflect::set(&env, &name.into(), &value.into()).unwrap();
+            Reflect::set(&env, &name.into(), &global.into()).unwrap();
         }
-
-        // Set the memory and table in the imports
         let imports = Object::new();
         Reflect::set(&imports, &"env".into(), &env).unwrap();
 
-        // Download the module, returning { module, instance }
+        // Synchronous instantiation of the precompiled module — writes
-        // we unwrap here instead of using `?` since this whole thing is async
+        // the patch's data segment into memory at `memory_base`.
-        let result_object = JsFuture::from(WebAssembly::instantiate_module(
+        let instance = Instance::new(&module, &imports).unwrap();
-            dl_bytes.unchecked_ref(),
+        let inst_exports: Object = instance.exports();
-            &imports,
-        ))
-        .await
-        .unwrap();
 
-        // We need to run the data relocations and then fire off the constructors
+        // Run the patch's relocation thunks and constructors. Order
-        let res: Object = result_object.unchecked_into();
+        // matters: data relocs first (write memory_base- and table_base-
-        let instance: Object = Reflect::get(&res, &"instance".into())
+        // relative pointers into the patch's data segment), then global
-            .unwrap()
+        // relocs (adjust GOT.func.internal globals by __table_base —
-            .unchecked_into();
+        // wasm-ld synthesizes those as element-segment-relative offsets),
-        let exports: Object = Reflect::get(&instance, &"exports".into())
+        // then ctors. `dyn_into` instead of `unchecked_into` so missing
-            .unwrap()
+        // exports just no-op rather than throwing.
-            .unchecked_into();
+        for func_name in [
-        _ = Reflect::get(&exports, &"__wasm_apply_data_relocs".into())
+            "__wasm_apply_data_relocs",
-            .unwrap()
+            "__wasm_apply_global_relocs",
-            .unchecked_into::<js_sys::Function>()
+            "__wasm_call_ctors",
-            .call0(&JsValue::undefined());
+        ] {
-        _ = Reflect::get(&exports, &"__wasm_apply_global_relocs".into())
+            if let Ok(val) = Reflect::get(&inst_exports, &func_name.into()) {
-            .unwrap()
+                if let Ok(func) = val.dyn_into::<js_sys::Function>() {
-            .unchecked_into::<js_sys::Function>()
+                    _ = func.call0(&JsValue::undefined());
-            .call0(&JsValue::undefined());
+                }
-        _ = Reflect::get(&exports, &"__wasm_call_ctors".into())
+            }
-            .unwrap()
+        }
-            .unchecked_into::<js_sys::Function>()
-            .call0(&JsValue::undefined());
 
         unsafe { commit_patch(table) };
     });