Don't use Uint8Array when not necessary

rust-random · Dec 5, 2024 · f6ec0df · f6ec0df
1 parent 00530ff
commit f6ec0df
Show file tree

Hide file tree

Showing 2 changed files with 82 additions and 33 deletions.
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
@@ -231,9 +231,17 @@ jobs:
   web:
     name: Web
     runs-on: ubuntu-24.04
+    strategy:
+      fail-fast: false
+      matrix:
+        rust:
+          - { version: stable }
+          - { version: nightly, flags: "-Ctarget-feature=+atomics,+bulk-memory", args: -- -Zbuild-std=panic_abort,std }
     steps:
       - uses: actions/checkout@v4
-      - uses: dtolnay/rust-toolchain@stable
+      - uses: dtolnay/rust-toolchain@master
+        with:
+          toolchain: ${{ matrix.rust.version }}
       - name: Install precompiled wasm-pack
         shell: bash
         run: |
@@ -244,34 +252,34 @@ jobs:
       - uses: Swatinem/rust-cache@v2
       - name: Test (Node)
         env:
-          RUSTFLAGS: -Dwarnings --cfg getrandom_backend="wasm_js"
-        run: wasm-pack test --node
+          RUSTFLAGS: -Dwarnings --cfg getrandom_backend="wasm_js" ${{ matrix.rust.flags }}
+        run: wasm-pack test --node ${{ matrix.rust.args }}
       - name: Test (Firefox)
         env:
           WASM_BINDGEN_USE_BROWSER: 1
-          RUSTFLAGS: -Dwarnings --cfg getrandom_backend="wasm_js"
-        run: wasm-pack test --headless --firefox
+          RUSTFLAGS: -Dwarnings --cfg getrandom_backend="wasm_js" ${{ matrix.rust.flags }}
+        run: wasm-pack test --headless --firefox ${{ matrix.rust.args }}
       - name: Test (Chrome)
         env:
           WASM_BINDGEN_USE_BROWSER: 1
-          RUSTFLAGS: -Dwarnings --cfg getrandom_backend="wasm_js"
-        run: wasm-pack test --headless --chrome
+          RUSTFLAGS: -Dwarnings --cfg getrandom_backend="wasm_js" ${{ matrix.rust.flags }}
+        run: wasm-pack test --headless --chrome ${{ matrix.rust.args }}
       - name: Test (dedicated worker)
         env:
           WASM_BINDGEN_USE_DEDICATED_WORKER: 1
-          RUSTFLAGS: -Dwarnings --cfg getrandom_backend="wasm_js"
-        run: wasm-pack test --headless --firefox
+          RUSTFLAGS: -Dwarnings --cfg getrandom_backend="wasm_js" ${{ matrix.rust.flags }}
+        run: wasm-pack test --headless --firefox ${{ matrix.rust.args }}
       - name: Test (shared worker)
         env:
           WASM_BINDGEN_USE_SHARED_WORKER: 1
-          RUSTFLAGS: -Dwarnings --cfg getrandom_backend="wasm_js"
-        run: wasm-pack test --headless --firefox
+          RUSTFLAGS: -Dwarnings --cfg getrandom_backend="wasm_js" ${{ matrix.rust.flags }}
+        run: wasm-pack test --headless --firefox ${{ matrix.rust.args }}
       - name: Test (service worker)
         env:
           WASM_BINDGEN_USE_SERVICE_WORKER: 1
-          RUSTFLAGS: -Dwarnings --cfg getrandom_backend="wasm_js"
+          RUSTFLAGS: -Dwarnings --cfg getrandom_backend="wasm_js" ${{ matrix.rust.flags }}
         # Firefox doesn't support module service workers and therefor can't import scripts
-        run: wasm-pack test --headless --chrome
+        run: wasm-pack test --headless --chrome ${{ matrix.rust.args }}
 
   wasi:
     name: WASI

diff --git a/src/backends/wasm_js.rs b/src/backends/wasm_js.rs
@@ -5,43 +5,82 @@
 extern crate std;
 
 use crate::Error;
-use core::mem::MaybeUninit;
+use core::{
+    mem::MaybeUninit,
+    sync::atomic::{AtomicU8, Ordering},
+};
 use std::thread_local;
 
 pub use crate::util::{inner_u32, inner_u64};
 
 #[cfg(not(all(target_arch = "wasm32", any(target_os = "unknown", target_os = "none"))))]
 compile_error!("`wasm_js` backend can be enabled only for OS-less WASM targets!");
 
-use js_sys::Uint8Array;
-use wasm_bindgen::{prelude::wasm_bindgen, JsValue};
+use js_sys::{SharedArrayBuffer, Uint8Array, WebAssembly::Memory};
+use wasm_bindgen::{prelude::wasm_bindgen, JsCast, JsValue};
 
 // Size of our temporary Uint8Array buffer used with WebCrypto methods
 // Maximum is 65536 bytes see https://developer.mozilla.org/en-US/docs/Web/API/Crypto/getRandomValues
 const CRYPTO_BUFFER_SIZE: u16 = 256;
 
+const MEMORY_KIND_UNINIT: u8 = 0;
+const MEMORY_KIND_NOT_SHARED: u8 = 1;
+const MEMORY_KIND_SHARED: u8 = 2;
+
+static MEMORY_KIND: AtomicU8 = AtomicU8::new(0);
+
 pub fn fill_inner(dest: &mut [MaybeUninit<u8>]) -> Result<(), Error> {
     CRYPTO.with(|crypto| {
         let crypto = crypto.as_ref().ok_or(Error::WEB_CRYPTO)?;
 
-        // getRandomValues does not work with all types of WASM memory,
-        // so we initially write to browser memory to avoid exceptions.
-        let buf = Uint8Array::new_with_length(CRYPTO_BUFFER_SIZE.into());
-        for chunk in dest.chunks_mut(CRYPTO_BUFFER_SIZE.into()) {
-            let chunk_len: u32 = chunk
-                .len()
-                .try_into()
-                .expect("chunk length is bounded by CRYPTO_BUFFER_SIZE");
-            // The chunk can be smaller than buf's length, so we call to
-            // JS to create a smaller view of buf without allocation.
-            let sub_buf = buf.subarray(0, chunk_len);
-
-            if crypto.get_random_values(&sub_buf).is_err() {
-                return Err(Error::WEB_GET_RANDOM_VALUES);
-            }
+        let shared = loop {
+            break match MEMORY_KIND.load(Ordering::Relaxed) {
+                MEMORY_KIND_NOT_SHARED => false,
+                MEMORY_KIND_SHARED => true,
+                MEMORY_KIND_UNINIT => {
+                    let memory: Memory = wasm_bindgen::memory().unchecked_into();
+                    let val = if memory.buffer().is_instance_of::<SharedArrayBuffer>() {
+                        MEMORY_KIND_SHARED
+                    } else {
+                        MEMORY_KIND_NOT_SHARED
+                    };
+                    MEMORY_KIND.store(val, Ordering::Relaxed);
+                    continue;
+                }
+                _ => unreachable!(),
+            };
+        };
+
+        if shared {
+            // getRandomValues does not work with all types of WASM memory,
+            // so we initially write to browser memory to avoid exceptions.
+            let buf = Uint8Array::new_with_length(CRYPTO_BUFFER_SIZE.into());
+            for chunk in dest.chunks_mut(CRYPTO_BUFFER_SIZE.into()) {
+                let chunk_len: u32 = chunk
+                    .len()
+                    .try_into()
+                    .expect("chunk length is bounded by CRYPTO_BUFFER_SIZE");
+                // The chunk can be smaller than buf's length, so we call to
+                // JS to create a smaller view of buf without allocation.
+                let sub_buf = buf.subarray(0, chunk_len);
 
-            // SAFETY: `sub_buf`'s length is the same length as `chunk`
-            unsafe { sub_buf.raw_copy_to_ptr(chunk.as_mut_ptr().cast::<u8>()) };
+                if crypto.get_random_values(&sub_buf).is_err() {
+                    return Err(Error::WEB_GET_RANDOM_VALUES);
+                }
+
+                // SAFETY: `sub_buf`'s length is the same length as `chunk`
+                unsafe { sub_buf.raw_copy_to_ptr(chunk.as_mut_ptr().cast::<u8>()) };
+            }
+        } else {
+            for chunk in dest.chunks_mut(CRYPTO_BUFFER_SIZE.into()) {
+                // SAFETY: this is only safe because on Wasm the issues with unitialized data don't exist
+                let buf = unsafe {
+                    core::slice::from_raw_parts_mut(chunk.as_mut_ptr().cast::<u8>(), chunk.len())
+                };
+                if crypto.get_random_values_ref(buf).is_err() {
+                    return Err(Error::WEB_GET_RANDOM_VALUES);
+                }
+            }
         }
         Ok(())
     })
@@ -57,4 +96,6 @@ extern "C" {
     // Crypto.getRandomValues()
     #[wasm_bindgen(method, js_name = getRandomValues, catch)]
     fn get_random_values(this: &Crypto, buf: &Uint8Array) -> Result<(), JsValue>;
+    #[wasm_bindgen(method, js_name = getRandomValues, catch)]
+    fn get_random_values_ref(this: &Crypto, buf: &mut [u8]) -> Result<(), JsValue>;
 }