From cc5510d785769edc97e8e6f7b3283daf4bf00aa7 Mon Sep 17 00:00:00 2001
From: Joe Richey <joerichey@google.com>
Date: Tue, 11 Jun 2024 14:57:01 -0700
Subject: [PATCH] Cleanup Custom Tests

This marks the `tests/custom.rs` test as requiring the `"custom"`
feature. It also:
  - Simplifies the registered custom RNG
  - Makes sure the custom RNG _is not_ used on supported platforms
  - Makes sure the custom RNG _is_ used on unsupported platforms

Signed-off-by: Joe Richey <joerichey@google.com>
---
 .github/workflows/tests.yml |  6 ++--
 Cargo.toml                  |  4 +++
 tests/custom.rs             | 71 ++++++++++++++++++-------------------
 3 files changed, 43 insertions(+), 38 deletions(-)

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 1d86ee27..26f0760c 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -38,9 +38,9 @@ jobs:
           toolchain: ${{ matrix.toolchain }}
       - uses: Swatinem/rust-cache@v2
       - run: cargo test
-      - run: cargo test --features=std
+      # Make sure enabling the std and custom features don't break anything
+      - run: cargo test --features=std,custom
       - run: cargo test --features=linux_disable_fallback
-      - run: cargo test --features=custom # custom should do nothing here
       - if: ${{ matrix.toolchain == 'nightly' }}
         run: cargo test --benches
 
@@ -260,6 +260,8 @@ jobs:
         run: wasm-pack test --headless --safari --features=js,test-in-browser
       - name: Test (custom getrandom)
         run: wasm-pack test --node --features=custom
+      - name: Test (JS overrides custom)
+        run: wasm-pack test --node --features=custom,js
 
   wasm64-tests:
     name: wasm64 Build/Link
diff --git a/Cargo.toml b/Cargo.toml
index 47c7fd28..032ae083 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -53,6 +53,10 @@ rustc-dep-of-std = [
 # Unstable/test-only feature to run wasm-bindgen tests in a browser
 test-in-browser = []
 
+[[test]]
+name = "custom"
+required-features = ["custom"]
+
 [package.metadata.docs.rs]
 features = ["std", "custom"]
 rustdoc-args = ["--cfg", "docsrs"]
diff --git a/tests/custom.rs b/tests/custom.rs
index b085094b..fcc1b4cd 100644
--- a/tests/custom.rs
+++ b/tests/custom.rs
@@ -1,54 +1,53 @@
-// 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;
-
 use core::num::NonZeroU32;
 use getrandom::{getrandom, register_custom_getrandom, Error};
 
-fn len7_err() -> Error {
-    NonZeroU32::new(Error::INTERNAL_START + 7).unwrap().into()
-}
+#[cfg(all(target_family = "wasm", target_os = "unknown"))]
+use wasm_bindgen_test::wasm_bindgen_test as test;
 
-fn super_insecure_rng(buf: &mut [u8]) -> Result<(), Error> {
+const LEN7_CODE: u32 = Error::CUSTOM_START + 7;
+// Returns a custom error if input is length 7, otherwise fills with 0x55.
+fn mock_rng(buf: &mut [u8]) -> Result<(), Error> {
     // `getrandom` guarantees it will not call any implementation if the output
     // buffer is empty.
     assert!(!buf.is_empty());
-    // Length 7 buffers return a custom error
     if buf.len() == 7 {
-        return Err(len7_err());
-    }
-    // Otherwise, fill bytes based on input length
-    let mut start = buf.len() as u8;
-    for b in buf {
-        *b = start;
-        start = start.wrapping_mul(3);
+        return Err(NonZeroU32::new(LEN7_CODE).unwrap().into());
     }
+    buf.fill(0x55);
     Ok(())
 }
 
-register_custom_getrandom!(super_insecure_rng);
-
-use getrandom::getrandom as getrandom_impl;
-mod common;
+// Test registering a custom implementation, even on supported platforms.
+register_custom_getrandom!(mock_rng);
 
+// Invoking with an empty buffer should never call the custom implementation.
 #[test]
-fn custom_rng_output() {
-    let mut buf = [0u8; 4];
-    assert_eq!(getrandom(&mut buf), Ok(()));
-    assert_eq!(buf, [4, 12, 36, 108]);
-
-    let mut buf = [0u8; 3];
-    assert_eq!(getrandom(&mut buf), Ok(()));
-    assert_eq!(buf, [3, 9, 27]);
+fn custom_empty() {
+    getrandom(&mut []).unwrap();
 }
 
+// On a supported platform, make sure the custom implementation isn't used. We
+// test on a few common platfroms, rather than duplicating the lib.rs logic.
+#[cfg(any(
+    target_os = "linux",
+    target_os = "windows",
+    target_os = "macos",
+    target_os = "espidf",
+    target_os = "wasi",
+    all(target_family = "wasm", target_os = "unknown", feature = "js"),
+))]
 #[test]
-fn rng_err_output() {
-    assert_eq!(getrandom(&mut [0; 7]), Err(len7_err()));
+fn custom_not_used() {
+    getrandom(&mut [0; 7]).unwrap();
+}
+// On an unsupported platform, make sure the custom implementation is used.
+#[cfg(all(target_family = "wasm", target_os = "unknown", not(feature = "js")))]
+#[test]
+fn custom_used() {
+    let err = getrandom(&mut [0; 7]).unwrap_err();
+    assert_eq!(err.code().get(), LEN7_CODE);
+
+    let mut buf = [0; 12];
+    getrandom(&mut buf).unwrap();
+    assert_eq!(buf, [0x55; 12]);
 }