nist-rand 0.1.0

//! # NIST SP 800-90B Entropy Source & Health Tests
//!
//! This module provides two components:
//!
//! 1. **[`HealthTests`]** — the mandatory continuous health tests required by SP 800-90B § 4.4
//!    and FIPS 140-3.  These are run on every raw byte emitted by a noise source.
//!
//! 2. **[`OsRng`]** — a zero-dependency reader of the OS entropy pool (`/dev/urandom`)
//!    that subjects every byte to the health tests before passing it upstream.
//!
//! ## FIPS 140-3 Fail-Safe Requirement
//!
//! Both tests **panic** on failure rather than returning an error. This is the
//! correct FIPS 140-3 behaviour for an entropy source failure: the module must
//! enter an *error state* and cease generating output. A panic in Rust satisfies
//! this requirement by unwinding (or aborting) the calling thread.
//!
//! ## Test Parameters & Derivation
//!
//! ### OS RNG (H = 7.9 bits/sample, α = 2⁻³⁰)
//!
//! | Test | Formula | Value |
//! |------|---------|-------|
//! | RCT cutoff  | `1 + ⌈−log₂(α) / H⌉` | **5**  |
//! | APT cutoff  | binomial CDF⁻¹(1 − α; n = 511, p = 2⁻⁷·⁹) | **20** |
//!
//! ### CPU Jitter (H = 2 bits/sample, α = 2⁻³⁰)
//!
//! | Test | Formula | Value |
//! |------|---------|-------|
//! | RCT cutoff  | `1 + ⌈−log₂(α) / H⌉` | **16**  |
//! | APT cutoff  | binomial CDF⁻¹(1 − α; n = 511, p = 2⁻²) | **250** |
//!
//! The `advance` feature additionally enables **Autocorrelation** (Markov transition
//! flip-count) and **Symbol Diversity** tests that catch subtle periodicity and bias
//! that the mandatory FIPS tests alone would miss.

use std::fs::File;
use std::io::{self, Read};

#[cfg(feature = "zeroize")]
use zeroize::Zeroizing;

// ────────────────────────────────────────────────────────────────────────────────
// Constants — OS RNG (H ≈ 7.9 bits/sample)
// ────────────────────────────────────────────────────────────────────────────────

/// **RCT cutoff** for the OS entropy source.
///
/// Derivation (SP 800-90B § 4.4.1, H = 7.9, α = 2⁻³⁰):
/// `C = 1 + ⌈30 / 7.9⌉ = 1 + 4 = 5`
pub const RCT_CUTOFF: usize = 5;

/// **APT cutoff** for the OS entropy source.
///
/// Derivation (SP 800-90B § 4.4.2, n = 511, p = 2⁻⁷·⁹, α = 2⁻³⁰):
/// Approximate binomial CDF inverse; value confirmed against SP 800-90B Table 1.
pub const APT_CUTOFF: usize = 20;

/// **Autocorrelation (Markov flip) lower bound** for the OS entropy source
/// over a 512-sample window.  A healthy `/dev/urandom` byte stream should
/// exhibit approximately `256 × 4 = 1024` bit-flips per 512 bytes; we allow
/// a generous ±45 % margin (1856 – 2240) to absorb natural variance.
///
/// Only active when the `advance` feature is enabled.
pub const AUTO_LOWER_BOUND: usize = 1856;

/// **Autocorrelation (Markov flip) upper bound** for the OS entropy source.
///
/// Only active when the `advance` feature is enabled.
pub const AUTO_UPPER_BOUND: usize = 2240;

/// **Minimum distinct byte values** observed in a 512-sample window (OS RNG).
///
/// A healthy source should cover most of the 0–255 alphabet; 180 distinct
/// symbols per 512 samples is a conservative lower bound.
///
/// Only active when the `advance` feature is enabled.
pub const SYMBOL_MIN_UNIQUE: usize = 180;

// ────────────────────────────────────────────────────────────────────────────────
// Constants — CPU Jitter (H ≈ 2 bits/sample)
// ────────────────────────────────────────────────────────────────────────────────

/// **RCT cutoff** for the CPU jitter entropy source.
///
/// Derivation (SP 800-90B § 4.4.1, H = 2, α = 2⁻³⁰):
/// `C = 1 + ⌈30 / 2⌉ = 16`
pub const JITTER_RCT_CUTOFF: usize = 16;

/// **APT cutoff** for the CPU jitter entropy source.
///
/// Derivation (SP 800-90B § 4.4.2, n = 511, p = 2⁻², α = 2⁻³⁰):
/// Approximate binomial CDF inverse.
pub const JITTER_APT_CUTOFF: usize = 250;

/// **Autocorrelation lower bound** for the jitter source (very loose;
/// natural timing distributions can be far from uniform).
///
/// Only active when the `advance` feature is enabled.
pub const JITTER_AUTO_LOWER_BOUND: usize = 200;

/// **Autocorrelation upper bound** for the jitter source.
///
/// Only active when the `advance` feature is enabled.
pub const JITTER_AUTO_UPPER_BOUND: usize = 3800;

/// **Minimum distinct byte values** observed per 512-sample jitter window.
///
/// Only active when the `advance` feature is enabled.
pub const JITTER_SYMBOL_MIN_UNIQUE: usize = 32;

// ────────────────────────────────────────────────────────────────────────────────
// Health Test State Machine
// ────────────────────────────────────────────────────────────────────────────────

/// Continuous health test state for an entropy noise source (SP 800-90B § 4.4).
///
/// Instantiate with calibrated cutoffs for the target source, then call
/// [`process_sample`](Self::process_sample) for every raw byte before use.
///
/// ## Mandatory FIPS Tests (always active)
///
/// - **Repetition Count Test (RCT)** — detects a "stuck" source that emits the
///   same byte more than `rct_cutoff` consecutive times.
/// - **Adaptive Proportion Test (APT)** — detects a heavily biased source where
///   a single value dominates a 512-sample window.
///
/// ## Extended Tests (`advance` feature only)
///
/// - **Autocorrelation / Markov flip-count** — catches highly periodic byte
///   sequences that pass RCT/APT but exhibit suspiciously low or high transition
///   rates between successive byte values.
/// - **Symbol Diversity** — counts distinct byte values in each 512-sample window;
///   too few distinct values indicates a structurally restricted source.
#[derive(Debug)]
pub struct HealthTests {
    // ── Repetition Count Test ─────────────────────────────────────────────────
    rct_previous_sample: Option<u8>,
    rct_count: usize,
    rct_cutoff: usize,

    // ── Adaptive Proportion Test ──────────────────────────────────────────────
    /// Ring buffer of the current 512-sample APT window.
    apt_window: [u8; 512],
    apt_index: usize,
    apt_base_sample: Option<u8>,
    apt_count: usize,
    apt_cutoff: usize,

    // ── Advanced: Autocorrelation (Markov flip count) ─────────────────────────
    #[cfg(feature = "advance")]
    auto_flips: usize,
    #[cfg(feature = "advance")]
    auto_lower_bound: usize,
    #[cfg(feature = "advance")]
    auto_upper_bound: usize,

    // ── Advanced: Symbol Diversity ────────────────────────────────────────────
    /// 256-bit bitmap: bit `i` is set when byte value `i` has been seen.
    #[cfg(feature = "advance")]
    symbol_seen: [u64; 4],
    #[cfg(feature = "advance")]
    symbol_min_unique: usize,
}

impl HealthTests {
    /// Constructs a `HealthTests` instance with extended (non-FIPS) checks.
    ///
    /// Available only when the `advance` Cargo feature is enabled.
    ///
    /// # Parameters
    ///
    /// - `rct_cutoff` — maximum run-length of a repeated byte before panic.
    /// - `apt_cutoff` — maximum count of a base value in a 512-sample window.
    /// - `auto_lower_bound` / `auto_upper_bound` — acceptable flip-count range.
    /// - `symbol_min_unique` — minimum distinct byte values per window.
    #[cfg(feature = "advance")]
    pub fn new(
        rct_cutoff: usize,
        apt_cutoff: usize,
        auto_lower_bound: usize,
        auto_upper_bound: usize,
        symbol_min_unique: usize,
    ) -> Self {
        Self {
            rct_previous_sample: None,
            rct_count: 0,
            rct_cutoff,
            apt_window: [0; 512],
            apt_index: 0,
            apt_base_sample: None,
            apt_count: 0,
            apt_cutoff,
            auto_flips: 0,
            auto_lower_bound,
            auto_upper_bound,
            symbol_seen: [0u64; 4],
            symbol_min_unique,
        }
    }

    /// Constructs a strictly FIPS SP 800-90B compliant `HealthTests` instance.
    ///
    /// Only the mandatory RCT and APT tests are included.  Available when the
    /// `advance` feature is **not** enabled.
    ///
    /// # Parameters
    ///
    /// - `rct_cutoff` — maximum run-length of a repeated byte before panic.
    /// - `apt_cutoff` — maximum count of a base value in a 512-sample window.
    #[cfg(not(feature = "advance"))]
    pub fn new(rct_cutoff: usize, apt_cutoff: usize) -> Self {
        Self {
            rct_previous_sample: None,
            rct_count: 0,
            rct_cutoff,
            apt_window: [0; 512],
            apt_index: 0,
            apt_base_sample: None,
            apt_count: 0,
            apt_cutoff,
        }
    }

    /// Feeds `sample` into all active health tests.
    ///
    /// # Panics
    ///
    /// Panics immediately if any test detects a failure condition.  This is the
    /// correct FIPS 140-3 behaviour: halt and refuse to produce output rather
    /// than silently degrading entropy quality.
    pub fn process_sample(&mut self, sample: u8) {
        self.advanced_tests(sample);
        self.repetition_count_test(sample);
        self.adaptive_proportion_test(sample);
    }

    // ── Private: Repetition Count Test (SP 800-90B § 4.4.1) ──────────────────

    fn repetition_count_test(&mut self, sample: u8) {
        match self.rct_previous_sample {
            Some(prev) if sample == prev => {
                self.rct_count += 1;
                if self.rct_count >= self.rct_cutoff {
                    panic!(
                        "NIST SP 800-90B Repetition Count Test failed! \
                         Entropy source is stuck."
                    );
                }
            }
            _ => {
                self.rct_previous_sample = Some(sample);
                self.rct_count = 1;
            }
        }
    }

    // ── Private: Advanced tests (compile-time gated) ──────────────────────────

    /// Records per-sample state for the advanced autocorrelation and symbol
    /// diversity tests.  No-op when the `advance` feature is disabled.
    #[cfg(feature = "advance")]
    fn advanced_tests(&mut self, sample: u8) {
        // Symbol diversity: mark this byte value as seen in the 256-bit bitmap.
        let word_idx = sample as usize / 64;
        let bit_idx = sample as usize % 64;
        self.symbol_seen[word_idx] |= 1u64 << bit_idx;

        // Autocorrelation: count bit-transitions between consecutive bytes.
        if let Some(prev) = self.rct_previous_sample {
            self.auto_flips += (sample ^ prev).count_ones() as usize;
        }
    }

    #[cfg(not(feature = "advance"))]
    #[inline(always)]
    fn advanced_tests(&mut self, _sample: u8) {}

    // ── Private: Adaptive Proportion Test (SP 800-90B § 4.4.2) ───────────────

    fn adaptive_proportion_test(&mut self, sample: u8) {
        self.apt_window[self.apt_index] = sample;

        match self.apt_base_sample {
            None => {
                // First byte of a new window becomes the base sample.
                self.apt_base_sample = Some(sample);
                self.apt_count = 1;
            }
            Some(base) if sample == base => {
                self.apt_count += 1;
                if self.apt_count >= self.apt_cutoff {
                    panic!(
                        "NIST SP 800-90B Adaptive Proportion Test failed! \
                         Entropy source lacks variation."
                    );
                }
            }
            _ => {}
        }

        self.apt_index += 1;

        if self.apt_index == 512 {
            // ── End of 512-sample window ──────────────────────────────────────
            #[cfg(feature = "advance")]
            self.evaluate_advanced_tests();

            // Reset APT for the next window.
            self.apt_base_sample = None;
            self.apt_count = 0;
            self.apt_index = 0;
        }
    }

    /// Evaluates the autocorrelation and symbol diversity tests at window
    /// boundaries, then resets their accumulators.
    #[cfg(feature = "advance")]
    fn evaluate_advanced_tests(&mut self) {
        // Autocorrelation (Markov flip-count)
        if self.auto_lower_bound > 0
            && self.auto_upper_bound > 0
            && (self.auto_flips < self.auto_lower_bound
                || self.auto_flips > self.auto_upper_bound)
        {
            panic!(
                "Advanced Autocorrelation Test failed! \
                 Source shows dangerous periodicity."
            );
        }

        // Symbol diversity
        let unique_symbols: usize = self
            .symbol_seen
            .iter()
            .map(|&word| word.count_ones() as usize)
            .sum();

        if unique_symbols < self.symbol_min_unique {
            panic!(
                "Advanced Symbol Diversity Test failed! \
                 Source is heavily biased."
            );
        }

        // Reset per-window accumulators.
        self.auto_flips = 0;
        self.symbol_seen = [0u64; 4];
    }
}

// ────────────────────────────────────────────────────────────────────────────────
// OS Entropy Source
// ────────────────────────────────────────────────────────────────────────────────

/// A health-tested reader of the OS entropy pool.
///
/// Reads from `/dev/urandom` (Linux) and subjects every raw byte to the
/// mandatory SP 800-90B continuous health tests before forwarding it to the
/// caller.  This satisfies the FIPS 140-3 requirement that cryptographic
/// modules perform continuous monitoring on their entropy sources.
///
/// # Design Note
///
/// `/dev/urandom` is chosen over `/dev/random` because the latter can block
/// on kernels < 5.6.  Since Linux 5.6 both are equivalent in behaviour;
/// `/dev/urandom` provides the same quality without the risk of blocking.
#[derive(Debug)]
pub struct OsRng {
    health_tests: HealthTests,
}

impl OsRng {
    /// Opens `/dev/urandom` and initialises the health test state.
    ///
    /// # Errors
    ///
    /// Returns an `io::Error` if `/dev/urandom` cannot be opened (e.g., inside
    /// a highly restricted sandbox).
    pub fn new() -> io::Result<Self> {
        // Just verify we can open it once
        File::open("/dev/urandom")?;

        #[cfg(feature = "advance")]
        let health_tests = HealthTests::new(
            RCT_CUTOFF,
            APT_CUTOFF,
            AUTO_LOWER_BOUND,
            AUTO_UPPER_BOUND,
            SYMBOL_MIN_UNIQUE,
        );
        #[cfg(not(feature = "advance"))]
        let health_tests = HealthTests::new(RCT_CUTOFF, APT_CUTOFF);

        Ok(Self { health_tests })
    }

    /// Fills `dest` with OS entropy, running every byte through the SP 800-90B
    /// continuous health tests.
    ///
    /// # Errors
    ///
    /// Returns an `io::Error` if the read from `/dev/urandom` fails.
    ///
    /// # Panics
    ///
    /// Panics (FIPS fail-safe) if the health tests detect a degenerate source.
    pub fn fill_bytes(&mut self, dest: &mut [u8]) -> io::Result<()> {
        // Read into a temporary buffer first so that `dest` is never
        // partially filled if an I/O error occurs mid-read.
        // `Zeroizing` erases the OS entropy from the heap even on panic.
        #[cfg(feature = "zeroize")]
        let mut tmp = Zeroizing::new(vec![0u8; dest.len()]);
        #[cfg(not(feature = "zeroize"))]
        let mut tmp = vec![0u8; dest.len()];
        let mut file = File::open("/dev/urandom")?;
        file.read_exact(&mut tmp)?;

        for (i, &byte) in tmp.iter().enumerate() {
            self.health_tests.process_sample(byte);
            dest[i] = byte;
        }

        Ok(())
    }
}

// ────────────────────────────────────────────────────────────────────────────────
// Tests
// ────────────────────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;

    // Helper: build a HealthTests with the standard OS RNG parameters.
    fn os_rng_health() -> HealthTests {
        #[cfg(feature = "advance")]
        return HealthTests::new(
            RCT_CUTOFF,
            APT_CUTOFF,
            AUTO_LOWER_BOUND,
            AUTO_UPPER_BOUND,
            SYMBOL_MIN_UNIQUE,
        );
        #[cfg(not(feature = "advance"))]
        return HealthTests::new(RCT_CUTOFF, APT_CUTOFF);
    }

    // ── Mandatory FIPS failure paths ──────────────────────────────────────────

    #[test]
    #[should_panic(
        expected = "NIST SP 800-90B Repetition Count Test failed! Entropy source is stuck."
    )]
    fn rct_failure_on_stuck_source() {
        let mut h = os_rng_health();
        // Feed the exact same byte RCT_CUTOFF times — simulating a stuck source.
        for _ in 0..RCT_CUTOFF {
            h.process_sample(0x42);
        }
    }

    #[test]
    #[should_panic(
        expected = "NIST SP 800-90B Adaptive Proportion Test failed! Entropy source lacks variation."
    )]
    fn apt_failure_on_skewed_distribution() {
        let mut h = os_rng_health();
        let base = 0x77u8;
        // Interleave `base` with other values to evade RCT, but APT catches
        // that `base` appears ≥ APT_CUTOFF times in the 512-byte window.
        for i in 0..512usize {
            h.process_sample(if i % 2 == 0 { base } else { (i % 250) as u8 });
        }
    }

    /// Simulates a hypervisor feeding constant execution times to the jitter
    /// RNG, causing its samples to collapse to a single constant value.
    #[test]
    #[should_panic(
        expected = "NIST SP 800-90B Repetition Count Test failed! Entropy source is stuck."
    )]
    fn rct_catches_constant_time_jitter_attack() {
        eprintln!("\n--- SIMULATING CONSTANT TIME (ZERO JITTER) ATTACK ---");

        #[cfg(feature = "advance")]
        let mut h = HealthTests::new(
            JITTER_RCT_CUTOFF,
            JITTER_APT_CUTOFF,
            JITTER_AUTO_LOWER_BOUND,
            JITTER_AUTO_UPPER_BOUND,
            JITTER_SYMBOL_MIN_UNIQUE,
        );
        #[cfg(not(feature = "advance"))]
        let mut h = HealthTests::new(JITTER_RCT_CUTOFF, JITTER_APT_CUTOFF);

        // A constant sample will exceed JITTER_RCT_CUTOFF well within 100 iterations.
        for _ in 0..100 {
            h.process_sample(123);
        }
    }

    // ── Advanced test failure paths (advance feature only) ────────────────────

    /// Simulates a repeating 4-byte pattern that evades RCT and APT but is caught
    /// by the Symbol Diversity test (only 4 unique symbols < JITTER_SYMBOL_MIN_UNIQUE).
    #[test]
    #[cfg(feature = "advance")]
    #[should_panic(expected = "Advanced Symbol Diversity Test failed! Source is heavily biased.")]
    fn symbol_diversity_catches_low_complexity_pattern() {
        eprintln!("\n--- SIMULATING MARKOV PERIODICITY ATTACK ---");

        let mut h = HealthTests::new(
            JITTER_RCT_CUTOFF,
            JITTER_APT_CUTOFF,
            JITTER_AUTO_LOWER_BOUND,
            JITTER_AUTO_UPPER_BOUND,
            JITTER_SYMBOL_MIN_UNIQUE,
        );

        // Pattern [0, 1, 2, 3] repeated 128×:
        //   - Evades RCT (no consecutive repeats)
        //   - Evades APT (each value appears 128/512 ≈ 25 % < APT_CUTOFF)
        //   - Flip count ≈ 768, within autocorrelation bounds
        //   - Fails symbol diversity (only 4 distinct values < 32)
        let pattern = [0u8, 1, 2, 3];
        for i in 0..512 {
            h.process_sample(pattern[i % 4]);
        }
    }

    // ── Happy-path / integration tests ───────────────────────────────────────

    #[test]
    fn rct_passes_on_healthy_drbg_output() {
        let mut h = os_rng_health();
        let mut drbg = crate::drbg::HashDrbg::instantiate(
            b"test_seed_for_rct_success_with_sufficient_length",
        );
        let mut buf = [0u8; 1024];
        drbg.generate(&mut buf);
        for byte in buf {
            h.process_sample(byte);
        }
    }

    #[test]
    fn apt_passes_on_healthy_drbg_output() {
        let mut h = os_rng_health();
        let mut drbg = crate::drbg::HashDrbg::instantiate(
            b"test_seed_for_apt_success_with_sufficient_length",
        );
        let mut buf = [0u8; 1024];
        drbg.generate(&mut buf);
        for byte in buf {
            h.process_sample(byte);
        }
    }

    #[test]
    fn os_rng_reads_and_passes_health_tests() {
        let mut rng = OsRng::new().expect("OsRng init failed");
        // 1024 bytes triggers at least one full 512-sample APT window evaluation.
        let mut buf = [0u8; 1024];
        assert!(rng.fill_bytes(&mut buf).is_ok());
        assert!(buf.iter().any(|&b| b != 0), "should not be all-zero");
    }
}