use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::LazyLock;
use parking_lot::Mutex;
use rand::RngCore;
use crate::error::{HsmError, HsmResult};
const REPETITION_COUNT_CUTOFF: u32 = 41;
const ADAPTIVE_PROPORTION_WINDOW: usize = 1024;
const ADAPTIVE_PROPORTION_CUTOFF: u32 = 670;
const STARTUP_SAMPLES: usize = 1024;
pub struct EntropyHealthMonitor {
last_sample: u8,
repetition_count: u32,
first_sample_seen: bool,
window: [u8; ADAPTIVE_PROPORTION_WINDOW],
window_pos: usize,
window_value: u8, window_count: u32, window_initialized: bool,
startup_passed: bool,
}
impl EntropyHealthMonitor {
pub fn new() -> Self {
Self {
last_sample: 0,
repetition_count: 0,
first_sample_seen: false,
window: [0u8; ADAPTIVE_PROPORTION_WINDOW],
window_pos: 0,
window_value: 0,
window_count: 0,
window_initialized: false,
startup_passed: false,
}
}
pub fn startup_test(&mut self, entropy_bytes: &[u8]) -> HsmResult<()> {
if entropy_bytes.len() < STARTUP_SAMPLES {
tracing::error!(
"Startup health test: insufficient samples ({} < {})",
entropy_bytes.len(),
STARTUP_SAMPLES
);
return Err(HsmError::GeneralError);
}
self.repetition_count = 0;
self.first_sample_seen = false;
self.window_pos = 0;
self.window_initialized = false;
for &byte in &entropy_bytes[..STARTUP_SAMPLES] {
self.feed_sample_unchecked(byte)?;
}
self.startup_passed = true;
tracing::debug!("SP 800-90B startup health test passed");
Ok(())
}
pub fn feed_sample(&mut self, sample: u8) -> HsmResult<()> {
if !self.startup_passed {
return Err(HsmError::GeneralError);
}
self.feed_sample_unchecked(sample)
}
fn feed_sample_unchecked(&mut self, sample: u8) -> HsmResult<()> {
self.repetition_count_test(sample)?;
self.adaptive_proportion_test(sample)?;
Ok(())
}
pub fn feed_bytes(&mut self, bytes: &[u8]) -> HsmResult<()> {
if !self.startup_passed {
return Err(HsmError::GeneralError);
}
for &b in bytes {
self.feed_sample_unchecked(b)?;
}
Ok(())
}
pub fn observe(&mut self, bytes: &[u8]) -> HsmResult<()> {
self.feed_bytes(bytes)
}
pub fn startup_passed(&self) -> bool {
self.startup_passed
}
fn repetition_count_test(&mut self, sample: u8) -> HsmResult<()> {
if !self.first_sample_seen {
self.last_sample = sample;
self.repetition_count = 1;
self.first_sample_seen = true;
return Ok(());
}
if sample == self.last_sample {
self.repetition_count += 1;
if self.repetition_count >= REPETITION_COUNT_CUTOFF {
tracing::error!(
"SP 800-90B repetition count test FAILED: {} consecutive identical bytes (0x{:02X})",
self.repetition_count,
sample
);
return Err(HsmError::GeneralError);
}
} else {
self.last_sample = sample;
self.repetition_count = 1;
}
Ok(())
}
fn adaptive_proportion_test(&mut self, sample: u8) -> HsmResult<()> {
if !self.window_initialized {
self.window_value = sample;
self.window_count = 1;
self.window[0] = sample;
self.window_pos = 1;
self.window_initialized = true;
return Ok(());
}
self.window[self.window_pos] = sample;
self.window_pos += 1;
if sample == self.window_value {
self.window_count += 1;
}
if self.window_count >= ADAPTIVE_PROPORTION_CUTOFF {
tracing::error!(
"SP 800-90B adaptive proportion test FAILED: value 0x{:02X} appeared {} times in window of {}",
self.window_value,
self.window_count,
self.window_pos
);
return Err(HsmError::GeneralError);
}
if self.window_pos >= ADAPTIVE_PROPORTION_WINDOW {
self.window_pos = 0;
self.window_initialized = false;
}
Ok(())
}
}
impl Default for EntropyHealthMonitor {
fn default() -> Self {
Self::new()
}
}
static GLOBAL_ENTROPY_ERROR: AtomicBool = AtomicBool::new(false);
static GLOBAL_MONITOR: LazyLock<Mutex<EntropyHealthMonitor>> = LazyLock::new(|| {
use rand::rngs::OsRng;
let mut monitor = EntropyHealthMonitor::new();
let mut startup = [0u8; STARTUP_SAMPLES];
OsRng.fill_bytes(&mut startup);
if let Err(e) = monitor.startup_test(&startup) {
tracing::error!(
"SP 800-90B startup health test FAILED during global monitor init: {:?} \
— entropy source is unhealthy",
e
);
GLOBAL_ENTROPY_ERROR.store(true, Ordering::Release);
}
use zeroize::Zeroize;
startup.zeroize();
Mutex::new(monitor)
});
pub fn global_is_error_state() -> bool {
GLOBAL_ENTROPY_ERROR.load(Ordering::Acquire)
}
pub fn global_observe(bytes: &[u8]) -> HsmResult<()> {
if GLOBAL_ENTROPY_ERROR.load(Ordering::Acquire) {
return Err(HsmError::GeneralError);
}
let monitor = &*GLOBAL_MONITOR;
if GLOBAL_ENTROPY_ERROR.load(Ordering::Acquire) {
return Err(HsmError::GeneralError);
}
let mut guard = monitor.lock();
match guard.observe(bytes) {
Ok(()) => Ok(()),
Err(e) => {
tracing::error!(
"SP 800-90B continuous health test FAILED — global monitor entering error state"
);
GLOBAL_ENTROPY_ERROR.store(true, Ordering::Release);
Err(e)
}
}
}
#[cfg(test)]
pub fn global_reset_for_tests() {
GLOBAL_ENTROPY_ERROR.store(false, Ordering::Release);
let mut guard = GLOBAL_MONITOR.lock();
*guard = EntropyHealthMonitor::new();
use rand::rngs::OsRng;
let mut startup = [0u8; STARTUP_SAMPLES];
OsRng.fill_bytes(&mut startup);
let _ = guard.startup_test(&startup);
use zeroize::Zeroize;
startup.zeroize();
}
pub struct HealthMonitoredRng {
source: RngSource,
}
enum RngSource {
Os(OsEntropySource),
#[cfg(test)]
Injected(Box<dyn EntropySource + Send + Sync>),
}
impl RngSource {
#[inline]
fn fill(&mut self, dest: &mut [u8]) {
match self {
RngSource::Os(s) => s.fill(dest),
#[cfg(test)]
RngSource::Injected(s) => s.fill(dest),
}
}
}
pub trait EntropySource {
fn fill(&mut self, dest: &mut [u8]);
}
pub struct OsEntropySource;
impl EntropySource for OsEntropySource {
fn fill(&mut self, dest: &mut [u8]) {
use rand::rngs::OsRng;
OsRng.fill_bytes(dest);
}
}
impl HealthMonitoredRng {
pub fn new() -> Self {
Self {
source: RngSource::Os(OsEntropySource),
}
}
#[cfg(test)]
pub fn with_source(source: Box<dyn EntropySource + Send + Sync>) -> Self {
Self {
source: RngSource::Injected(source),
}
}
}
impl Default for HealthMonitoredRng {
fn default() -> Self {
Self::new()
}
}
impl RngCore for HealthMonitoredRng {
fn next_u32(&mut self) -> u32 {
let mut buf = [0u8; 4];
self.fill_bytes(&mut buf);
u32::from_le_bytes(buf)
}
fn next_u64(&mut self) -> u64 {
let mut buf = [0u8; 8];
self.fill_bytes(&mut buf);
u64::from_le_bytes(buf)
}
fn fill_bytes(&mut self, dest: &mut [u8]) {
self.source.fill(dest);
if let Err(e) = global_observe(dest) {
use zeroize::Zeroize;
dest.zeroize();
tracing::error!(
"HealthMonitoredRng: SP 800-90B health test failed ({:?}) — aborting",
e
);
std::process::abort();
}
}
fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), rand::Error> {
self.source.fill(dest);
global_observe(dest).map_err(|_| rand::Error::new("SP 800-90B health test failed"))
}
}
impl rand::CryptoRng for HealthMonitoredRng {}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_healthy_entropy() {
let mut monitor = EntropyHealthMonitor::new();
let mut bytes = vec![0u8; STARTUP_SAMPLES];
for (i, b) in bytes.iter_mut().enumerate() {
*b = (i.wrapping_mul(97).wrapping_add(31)) as u8;
}
assert!(monitor.startup_test(&bytes).is_ok());
assert!(monitor.startup_passed());
}
fn monitor_with_startup_passed() -> EntropyHealthMonitor {
let mut monitor = EntropyHealthMonitor::new();
let mut bytes = vec![0u8; STARTUP_SAMPLES];
for (i, b) in bytes.iter_mut().enumerate() {
*b = (i.wrapping_mul(97).wrapping_add(31)) as u8;
}
monitor.startup_test(&bytes).unwrap();
monitor
}
#[test]
fn test_feed_sample_fails_before_startup() {
let mut monitor = EntropyHealthMonitor::new();
assert!(monitor.feed_sample(0x42).is_err());
assert!(monitor.feed_bytes(&[1, 2, 3]).is_err());
}
#[test]
fn test_repetition_count_failure() {
let mut monitor = monitor_with_startup_passed();
for i in 0..REPETITION_COUNT_CUTOFF {
let result = monitor.feed_sample(0xAA);
if i < REPETITION_COUNT_CUTOFF - 1 {
assert!(result.is_ok());
} else {
assert!(result.is_err());
}
}
}
#[test]
fn test_repetition_count_reset_on_different_value() {
let mut monitor = monitor_with_startup_passed();
for _ in 0..40 {
assert!(monitor.feed_sample(0xBB).is_ok());
}
assert!(monitor.feed_sample(0xCC).is_ok()); }
#[test]
fn test_adaptive_proportion_failure() {
let mut monitor = monitor_with_startup_passed();
let mut failed = false;
for _ in 0..ADAPTIVE_PROPORTION_WINDOW {
if monitor.feed_sample(0x42).is_err() {
failed = true;
break;
}
}
assert!(failed, "Expected adaptive proportion test to fail");
}
#[test]
fn test_insufficient_startup_samples() {
let mut monitor = EntropyHealthMonitor::new();
let short = vec![0u8; 100];
assert!(monitor.startup_test(&short).is_err());
}
#[test]
fn test_first_byte_zero_no_off_by_one() {
let mut monitor = EntropyHealthMonitor::new();
let mut bytes = vec![0u8; STARTUP_SAMPLES];
for (i, b) in bytes.iter_mut().enumerate() {
*b = (i % 256) as u8; }
assert!(monitor.startup_test(&bytes).is_ok());
assert!(monitor.feed_sample(0x00).is_ok());
}
#[test]
fn test_observe_alias_of_feed_bytes() {
let mut monitor = monitor_with_startup_passed();
assert!(monitor.observe(&[1, 2, 3, 4, 5]).is_ok());
let mut fresh = EntropyHealthMonitor::new();
assert!(fresh.observe(&[1, 2, 3]).is_err());
}
#[test]
fn test_health_monitored_rng_healthy_source() {
global_reset_for_tests();
struct Pseudo(u32);
impl EntropySource for Pseudo {
fn fill(&mut self, dest: &mut [u8]) {
for b in dest.iter_mut() {
self.0 = self.0.wrapping_mul(1664525).wrapping_add(1013904223);
*b = (self.0 >> 24) as u8;
}
}
}
let mut rng = HealthMonitoredRng::with_source(Box::new(Pseudo(0xDEADBEEF)));
let mut buf = [0u8; 64];
rng.fill_bytes(&mut buf);
assert!(!global_is_error_state());
assert!(buf.iter().any(|&b| b != 0));
}
#[test]
fn test_health_monitored_rng_try_fill_detects_bad_source() {
global_reset_for_tests();
struct AllZero;
impl EntropySource for AllZero {
fn fill(&mut self, dest: &mut [u8]) {
for b in dest.iter_mut() {
*b = 0;
}
}
}
let mut rng = HealthMonitoredRng::with_source(Box::new(AllZero));
let mut buf = [0u8; 128];
let result = rng.try_fill_bytes(&mut buf);
assert!(
result.is_err(),
"all-zero source must trip SP 800-90B repetition count test"
);
assert!(
global_is_error_state(),
"global monitor should latch into error state after failure"
);
struct Pseudo(u32);
impl EntropySource for Pseudo {
fn fill(&mut self, dest: &mut [u8]) {
for b in dest.iter_mut() {
self.0 = self.0.wrapping_mul(1664525).wrapping_add(1013904223);
*b = (self.0 >> 24) as u8;
}
}
}
let mut rng2 = HealthMonitoredRng::with_source(Box::new(Pseudo(0xABCD)));
let mut buf2 = [0u8; 32];
assert!(
rng2.try_fill_bytes(&mut buf2).is_err(),
"monitor error state must persist across HealthMonitoredRng instances"
);
global_reset_for_tests();
}
#[test]
fn test_constant_nonzero_source_trips_monitor() {
global_reset_for_tests();
struct ConstantSource(u8);
impl EntropySource for ConstantSource {
fn fill(&mut self, dest: &mut [u8]) {
for b in dest.iter_mut() {
*b = self.0;
}
}
}
let mut rng = HealthMonitoredRng::with_source(Box::new(ConstantSource(0xAA)));
let mut buf = [0u8; 128];
assert!(
rng.try_fill_bytes(&mut buf).is_err(),
"constant 0xAA source must trip SP 800-90B continuous health tests"
);
assert!(global_is_error_state());
global_reset_for_tests();
}
#[test]
fn test_reset_clears_error_state() {
global_reset_for_tests();
struct AllZero;
impl EntropySource for AllZero {
fn fill(&mut self, dest: &mut [u8]) {
for b in dest.iter_mut() {
*b = 0;
}
}
}
let mut rng = HealthMonitoredRng::with_source(Box::new(AllZero));
let mut buf = [0u8; 128];
let _ = rng.try_fill_bytes(&mut buf);
assert!(global_is_error_state());
global_reset_for_tests();
assert!(
!global_is_error_state(),
"reset_for_tests must clear the latched error state"
);
struct Pseudo(u32);
impl EntropySource for Pseudo {
fn fill(&mut self, dest: &mut [u8]) {
for b in dest.iter_mut() {
self.0 = self.0.wrapping_mul(1664525).wrapping_add(1013904223);
*b = (self.0 >> 24) as u8;
}
}
}
let mut rng = HealthMonitoredRng::with_source(Box::new(Pseudo(42)));
let mut buf = [0u8; 64];
assert!(rng.try_fill_bytes(&mut buf).is_ok());
global_reset_for_tests();
}
}