From 4a4d8e8daf49ebe6a88d32d12cece0494fef4877 Mon Sep 17 00:00:00 2001
From: Joe Richey <joerichey@google.com>
Date: Tue, 5 Jan 2021 01:34:59 -0800
Subject: [PATCH] Add test for custom RNG handler

---
 .github/workflows/tests.yml |  4 +--
 tests/common/mod.rs         |  2 +-
 tests/custom.rs             | 50 +++++++++++++++++++++++++++++++++++++
 tests/normal.rs             |  8 ++++++
 4 files changed, 61 insertions(+), 3 deletions(-)
 create mode 100644 tests/custom.rs

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 24dd94249..163ca6a32 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -202,8 +202,8 @@ jobs:
         env:
           CHROMEDRIVER: /usr/bin/chromedriver
         run: cargo test --target=wasm32-unknown-unknown --features=js,test-in-browser
-      - name: Build Tests (with custom, without JS)
-        run: cargo test --no-run --target=wasm32-unknown-unknown --features=custom
+      - name: Test (custom getrandom)
+        run: cargo test --target=wasm32-unknown-unknown --features=custom
 
   wasi-tests:
     name: WASI test
diff --git a/tests/common/mod.rs b/tests/common/mod.rs
index 9c64c7578..006f230d7 100644
--- a/tests/common/mod.rs
+++ b/tests/common/mod.rs
@@ -1,6 +1,6 @@
 use super::getrandom_impl;
 
-#[cfg(all(target_arch = "wasm32", target_os = "unknown", not(cargo_web)))]
+#[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
 use wasm_bindgen_test::wasm_bindgen_test as test;
 
 #[cfg(feature = "test-in-browser")]
diff --git a/tests/custom.rs b/tests/custom.rs
new file mode 100644
index 000000000..62eae1d66
--- /dev/null
+++ b/tests/custom.rs
@@ -0,0 +1,50 @@
+// Test that a custom handler works on wasm32-unknown-unknown
+#![cfg(all(
+    target_arch = "wasm32",
+    target_os = "unknown",
+    feature = "custom",
+    not(feature = "js")
+))]
+
+use wasm_bindgen_test::wasm_bindgen_test as test;
+#[cfg(feature = "test-in-browser")]
+wasm_bindgen_test::wasm_bindgen_test_configure!(run_in_browser);
+
+use core::{
+    num::NonZeroU32,
+    sync::atomic::{AtomicU8, Ordering},
+};
+use getrandom::{getrandom, register_custom_getrandom, Error};
+
+fn len7_err() -> Error {
+    NonZeroU32::new(Error::INTERNAL_START + 7).unwrap().into()
+}
+
+fn super_insecure_rng(buf: &mut [u8]) -> Result<(), Error> {
+    // Length 7 buffers return a custom error
+    if buf.len() == 7 {
+        return Err(len7_err());
+    }
+    // Otherwise, increment an atomic counter
+    static COUNTER: AtomicU8 = AtomicU8::new(0);
+    for b in buf {
+        *b = COUNTER.fetch_add(1, Ordering::Relaxed);
+    }
+    Ok(())
+}
+
+register_custom_getrandom!(super_insecure_rng);
+
+#[test]
+fn custom_rng_output() {
+    let mut buf = [0u8; 4];
+    assert_eq!(getrandom(&mut buf), Ok(()));
+    assert_eq!(buf, [0, 1, 2, 3]);
+    assert_eq!(getrandom(&mut buf), Ok(()));
+    assert_eq!(buf, [4, 5, 6, 7]);
+}
+
+#[test]
+fn rng_err_output() {
+    assert_eq!(getrandom(&mut [0; 7]), Err(len7_err()));
+}
diff --git a/tests/normal.rs b/tests/normal.rs
index 7f3c37ab0..5fff13b38 100644
--- a/tests/normal.rs
+++ b/tests/normal.rs
@@ -1,3 +1,11 @@
+// Don't test on custom wasm32-unknown-unknown
+#![cfg(not(all(
+    target_arch = "wasm32",
+    target_os = "unknown",
+    feature = "custom",
+    not(feature = "js")
+)))]
+
 // Use the normal getrandom implementation on this architecture.
 use getrandom::getrandom as getrandom_impl;
 mod common;