From 43749b7e1b2ebc8fcfbbaadda4703fc2930f961e Mon Sep 17 00:00:00 2001
From: Joe Richey <joerichey@google.com>
Date: Tue, 11 Jun 2024 17:53:26 -0700
Subject: [PATCH] Test getrandom_uninit

Signed-off-by: Joe Richey <joerichey@google.com>
---
 src/tests.rs | 24 +++++++++++++++++++++++-
 1 file changed, 23 insertions(+), 1 deletion(-)

diff --git a/src/tests.rs b/src/tests.rs
index 8c99832e..82308fc3 100644
--- a/src/tests.rs
+++ b/src/tests.rs
@@ -2,7 +2,7 @@
 extern crate std;
 
 use crate::Error;
-use std::{sync::mpsc, thread, vec, vec::Vec};
+use std::{mem::MaybeUninit, sync::mpsc, thread, vec, vec::Vec};
 
 #[cfg(feature = "test-in-browser")]
 wasm_bindgen_test::wasm_bindgen_test_configure!(run_in_browser);
@@ -29,6 +29,15 @@ impl Byte for u8 {
         v
     }
 }
+impl Byte for MaybeUninit<u8> {
+    fn make_vec(len: usize, fill: FillFn<MaybeUninit<u8>>) -> Vec<u8> {
+        // Ensure that we always get uninitialized memory.
+        let mut v = Vec::with_capacity(len);
+        fill(v.spare_capacity_mut()).unwrap();
+        unsafe { v.set_len(len) }
+        v
+    }
+}
 
 // For calls of size `len`, count the number of bits which differ between calls
 // and check that between 3 and 5 bits per byte differ. Probability of failure:
@@ -112,3 +121,16 @@ macro_rules! define_tests {
 pub(crate) use define_tests;
 
 define_tests!(crate::getrandom);
+mod uninit {
+    use super::*;
+
+    fn wrapper(dest: &mut [MaybeUninit<u8>]) -> Result<(), Error> {
+        let dest_ptr = dest.as_ptr().cast::<u8>();
+        let res = crate::getrandom_uninit(dest)?;
+        // Ensure that the output points to the same bytes as the input.
+        assert_eq!(res.as_ptr(), dest_ptr);
+        assert_eq!(res.len(), dest.len());
+        Ok(())
+    }
+    super::define_tests!(wrapper);
+}