win32_rand.rs

use std::os::windows::ffi::OsStrExt;

type BCRYPT_ALG_HANDLE = usize;
type wchar_t = u16;
type WCHAR = wchar_t;
type LPCWSTR = *const WCHAR;
type c_uchar = u8;
type UCHAR = c_uchar;
type PUCHAR = *mut UCHAR;
type c_ulong = u32;
type ULONG = c_ulong;
type c_long = i32;
type LONG = c_long;
type NTSTATUS = LONG;
#[link(name = "Bcrypt")]
extern "C" {
    fn BCryptOpenAlgorithmProvider(
        _: *mut BCRYPT_ALG_HANDLE,
        _: LPCWSTR,
        _: LPCWSTR,
        _: ULONG,
    ) -> NTSTATUS;
    fn BCryptGenRandom(_: BCRYPT_ALG_HANDLE, _: PUCHAR, _: ULONG, _: ULONG) -> NTSTATUS;
    fn BCryptCloseAlgorithmProvider(_: BCRYPT_ALG_HANDLE, _: ULONG) -> NTSTATUS;
}
const STATUS_SUCCESS: NTSTATUS = 0;

pub struct RandState {
    handle: BCRYPT_ALG_HANDLE,
    buf: Box<[u8; 1024 * 1024]>,
    pos: usize,
}
impl RandState {
    fn reload_buf(&mut self) {
        let ret =
            unsafe { BCryptGenRandom(self.handle, self.buf.as_mut_ptr(), self.buf.len() as _, 0) };
        assert!(ret == STATUS_SUCCESS);
        self.pos = 0;
    }
    pub fn new() -> Self {
        let alg_name = std::ffi::OsStr::new("RNG")
            .encode_wide()
            .chain(Some(0))
            .collect::<Vec<u16>>();
        let mut alg = std::mem::MaybeUninit::<BCRYPT_ALG_HANDLE>::uninit();
        let mut ret: NTSTATUS = unsafe {
            BCryptOpenAlgorithmProvider(alg.as_mut_ptr(), alg_name.as_ptr(), std::ptr::null(), 0)
        };
        assert!(ret == STATUS_SUCCESS);
        let alg = unsafe { alg.assume_init() };

        use std::convert::TryInto;
        let buf: Box<[u8; 1024 * 1024]> =
            vec![0; 1024 * 1024].into_boxed_slice().try_into().unwrap();
        let mut ret = Self {
            buf,
            pos: 0,
            handle: alg,
        };
        ret.reload_buf();
        ret
    }

    pub fn gen_random_bytes(&mut self, buf: &mut [u8]) {
        let mut to_copy = buf.len();
        let remaining = self.buf.len() - self.pos;
        if to_copy <= remaining {
            buf.copy_from_slice(&self.buf[self.pos..(self.pos + to_copy)]);
            self.pos += to_copy;
        } else {
            (&mut buf[..remaining]).copy_from_slice(&self.buf[self.pos..]);
            self.reload_buf();
            self.gen_random_bytes(&mut buf[remaining..]);
        }
    }
}

impl Drop for RandState {
    fn drop(&mut self) {
        let ret = unsafe { BCryptCloseAlgorithmProvider(self.handle, 0) };
        assert!(ret == STATUS_SUCCESS);
    }
}