use super::{intrinsics, FeatureNotSupported, Port};
use core::sync::atomic::{AtomicU16, Ordering};
use mycelium_util::sync::spin;
use raw_cpuid::CpuId;
#[cfg(feature = "rand_core")]
use crate::time::Rdtsc;
#[cfg(feature = "rand_core")]
use mycelium_util::sync::Lazy;
#[cfg(feature = "rand_core")]
use rand_core::{RngCore, SeedableRng};
/// Construct a new [`SeedableRng`] instance from all available entropy sources.
///
/// This function constructs a new seed value for an `R`-typed [`SeedableRng`]
/// using entropy from the following sources:
///
/// - [`Rdrand`], if it is available on this CPU
/// - [`Rdtsc`], if timestamp counters are available on this CPU
/// - [PIT channels 1, 2, and 3](PitEntropy)
///
#[cfg(feature = "rand_core")]
pub fn seed_rng<R: SeedableRng>() -> R {
let mut seed = R::Seed::default();
static RDRAND: Lazy<Option<Rdrand>> = Lazy::new(|| Rdrand::new().ok());
static RDTSC: Lazy<Option<Rdtsc>> = Lazy::new(|| Rdtsc::new().ok());
for byte in seed.as_mut() {
// PIT entropy is always available. When RDRAND and RDTSC are available,
// we'll use the PIT entropy to determine which byte of the RDRAND and
// RDTSC values to use, as well.
let pit = {
let mut dst = [0u8];
PitEntropy::read_bytes(&mut dst[..]);
dst[0]
};
// If the CPU supports RDRAND, read a random value using that
// instruction. This provides high-quality entropy.
//
// RDRAND returns either a 16, 32, or 64-bit value. Since we only need a
// byte currently, we'll use the PIT entropy to determine which byte of
// the RDRAND value to XOR with the current byte.
if let Some(rdrand) = RDRAND.as_ref() {
*byte ^= rdrand.next_u32().to_ne_bytes()[pit as usize % 4];
}
// If the CPU supports RDTSC, read from the timestamp counter as an
// additional entropy source. This isn't a very good source of entropy,
// since it's a counter, but XOR-ing additional entropy is always
// better.
//
// Also, since we select which byte of the timestamp to use based on the
// PIT timestamp, it's at least a little bit fucked up, which helps.
if let Some(rdtsc) = RDTSC.as_ref() {
*byte ^= rdtsc.read_timestamp().to_ne_bytes()[pit as usize % 8];
}
// Finally, XOR the current byte with the PIT timestamp. This is always
// available. A seed generated only from PIT timestamps isn't a great
// source of entropy, but if the CPU doesn't support RDRAND or RDTSC,
// it's better than returning all zeros.
*byte ^= pit;
}
R::from_seed(seed)
}
#[derive(Debug, Copy, Clone)]
pub struct Rdrand(());
#[derive(Debug, Copy, Clone, Default)]
pub struct PitEntropy(());
// === impl Rdrand ===
impl Rdrand {
pub const MAX_RETRIES: usize = 64;
/// Returns `true` if the CPU supports the `rdrand` instruction.
///
/// This is determined using the `cpuid` instruction.
pub fn is_supported() -> bool {
CpuId::new()
.get_feature_info()
.map(|features| features.has_rdrand())
.unwrap_or(false)
}
pub fn new() -> Result<Self, FeatureNotSupported> {
if Self::is_supported() {
Ok(Self(()))
} else {
Err(FeatureNotSupported::new("rdrand"))
}
}
pub fn try_next_u32(&self) -> Option<u32> {
let mut res: u32 = 0;
unsafe {
match intrinsics::rdrand32_step(&mut res) {
1 => Some(res),
0 => None,
x => unreachable!("rdrand32_step should only return 1 or 0, but got {}", x),
}
}
}
pub fn next_u32(&self) -> u32 {
self.with_retry(Self::try_next_u32)
}
pub fn try_next_u64(&self) -> Option<u64> {
let mut res: u64 = 0;
unsafe {
match intrinsics::rdrand64_step(&mut res) {
1 => Some(res),
0 => None,
x => unreachable!("rdrand64_step should only return 1 or 0, but got {}", x),
}
}
}
pub fn next_u64(&self) -> u64 {
self.with_retry(Self::try_next_u64)
}
/// Retry a rdrand instruction until it succeeds, spinning with an
/// exponential backoff if entropy is unavailable.
fn with_retry<T: Default>(&self, f: impl Fn(&Self) -> Option<T>) -> T {
let mut backoff = spin::Backoff::new();
// Eventually, we do need to bail, since we don't want to be stuck in an
// infinite loop...
for _ in 0..Self::MAX_RETRIES {
if let Some(res) = f(self) {
return res;
}
backoff.spin();
}
T::default()
}
}
#[cfg(feature = "rand_core")]
impl RngCore for Rdrand {
fn next_u32(&mut self) -> u32 {
Rdrand::next_u32(&*self)
}
#[inline]
fn next_u64(&mut self) -> u64 {
Rdrand::next_u64(&*self)
}
fn fill_bytes(&mut self, dest: &mut [u8]) {
rand_core::impls::fill_bytes_via_next(self, dest)
}
fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), rand_core::Error> {
use core::num::NonZeroU32;
let mut left = dest;
while left.len() >= 8 {
let (l, r) = { left }.split_at_mut(8);
left = r;
let chunk: [u8; 8] = self
.try_next_u64()
.ok_or_else(|| rand_core::Error::from(NonZeroU32::new(1).expect("1 is non-zero")))?
.to_ne_bytes();
l.copy_from_slice(&chunk);
}
let n = left.len();
if n > 4 {
let chunk: [u8; 8] = self
.try_next_u64()
.ok_or_else(|| rand_core::Error::from(NonZeroU32::new(1).expect("1 is non-zero")))?
.to_ne_bytes();
left.copy_from_slice(&chunk[..n]);
} else if n > 0 {
let chunk: [u8; 4] = self
.try_next_u32()
.ok_or_else(|| rand_core::Error::from(NonZeroU32::new(1).expect("1 is non-zero")))?
.to_ne_bytes();
left.copy_from_slice(&chunk[..n]);
}
Ok(())
}
}
#[cfg(feature = "rand_core")]
impl RngCore for Rdtsc {
fn next_u32(&mut self) -> u32 {
self.read_timestamp() as u32
}
#[inline]
fn next_u64(&mut self) -> u64 {
self.read_timestamp()
}
fn fill_bytes(&mut self, dest: &mut [u8]) {
rand_core::impls::fill_bytes_via_next(self, dest)
}
fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), rand_core::Error> {
self.fill_bytes(dest);
Ok(())
}
}
impl PitEntropy {
pub const fn new() -> Self {
Self(())
}
pub fn next_u32() -> u32 {
let mut bytes = [0u8; 4];
Self::read_bytes(&mut bytes[..]);
u32::from_be_bytes(bytes)
}
pub fn next_u64() -> u64 {
let mut bytes = [0u8; 8];
Self::read_bytes(&mut bytes[..]);
u64::from_be_bytes(bytes)
}
fn read_bytes(dest: &mut [u8]) {
const PIT_BASE: u16 = 0x40;
// making this global state probably *increases* the amount of entropy,
// since the next channel to read from will depend on other uses of the
// entropy source in an unpredictable way...
static NEXT_CHANNEL: AtomicU16 = AtomicU16::new(0);
for byte in dest {
let channel = NEXT_CHANNEL.fetch_add(1, Ordering::Relaxed) % 3;
*byte = unsafe { Port::at(PIT_BASE + channel).readb() };
}
}
}
#[cfg(feature = "rand_core")]
impl RngCore for PitEntropy {
#[inline]
fn next_u32(&mut self) -> u32 {
Self::next_u32()
}
#[inline]
fn next_u64(&mut self) -> u64 {
Self::next_u64()
}
fn fill_bytes(&mut self, dest: &mut [u8]) {
Self::read_bytes(dest);
}
fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), rand_core::Error> {
Self::read_bytes(dest);
Ok(())
}
}