qs-drbg 0.1.0

#![forbid(unsafe_code)]

//! SP 800-90A Rev.1 HMAC_DRBG(SHA-512) implementation.

extern crate alloc;

/// Pure Rust, zeroizes state on drop, and enforces a reseed interval.
pub mod rand_adapter;

use alloc::vec::Vec;
use core::{cmp::min, fmt};
use getrandom::getrandom;
use hmac::{Hmac, Mac};
use sha2::Sha512;
use subtle::ConstantTimeEq;
use zeroize::{Zeroize, ZeroizeOnDrop};

/// Errors that can occur during DRBG operation.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Error {
    /// Generated output exceeds SP 800-90A per-request cap (64 KiB).
    RequestTooLarge,
    /// Reseed interval exhausted and new entropy is required.
    ReseedRequired,
    /// OS entropy source failed.
    EntropyUnavailable,
    /// Entropy health test failed (SP 800-90B).
    EntropyHealthFailed,
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{:?}", self)
    }
}

#[cfg(feature = "std")]
impl std::error::Error for Error {}

/// Default reseed interval recommended by SP 800-90A (Section 10.1.2).
const DEFAULT_RESEED_INTERVAL: u64 = 1u64 << 48;
const BLOCK_LEN: usize = 64;
const MAX_REQUEST: usize = 65536; // 64 KiB cap per generate call.
const DEFAULT_MAX_BYTES_BETWEEN_RESEED: u128 = 1u128 << 20;

/// HMAC_DRBG with HMAC-SHA-512 backbone.
#[derive(Zeroize, ZeroizeOnDrop)]
pub struct HmacDrbg {
    k: [u8; BLOCK_LEN],
    v: [u8; BLOCK_LEN],
    reseed_counter: u64,
    reseed_interval: u64,
    generated_bytes: u128,
    max_bytes_between_reseed: u128,
    last_entropy: Vec<u8>,
}

impl HmacDrbg {
    /// Instantiate DRBG from entropy, nonce, and optional personalization string.
    pub fn new(
        entropy: &[u8],
        nonce: &[u8],
        personalization: Option<&[u8]>,
    ) -> Result<Self, Error> {
        Self::validate_entropy(entropy)?;

        let mut seed = Vec::with_capacity(
            entropy.len() + nonce.len() + personalization.map_or(0, |p| p.len()),
        );
        seed.extend_from_slice(entropy);
        seed.extend_from_slice(nonce);
        if let Some(pers) = personalization {
            seed.extend_from_slice(pers);
        }

        let mut drbg = Self {
            k: [0u8; BLOCK_LEN],
            v: [0x01u8; BLOCK_LEN],
            reseed_counter: 1,
            reseed_interval: DEFAULT_RESEED_INTERVAL,
            generated_bytes: 0,
            max_bytes_between_reseed: DEFAULT_MAX_BYTES_BETWEEN_RESEED,
            last_entropy: entropy.to_vec(),
        };
        drbg.update(Some(&seed));
        seed.zeroize();
        Ok(drbg)
    }

    /// Instantiate by sampling entropy+nonce from the operating system.
    pub fn from_os(personalization: Option<&[u8]>) -> Result<Self, Error> {
        let mut entropy = [0u8; 48];
        let mut nonce = [0u8; 16];
        getrandom(&mut entropy).map_err(|_| Error::EntropyUnavailable)?;
        getrandom(&mut nonce).map_err(|_| Error::EntropyUnavailable)?;
        Self::validate_entropy(&nonce)?;
        let drbg = Self::new(&entropy, &nonce, personalization)?;
        entropy.zeroize();
        nonce.zeroize();
        Ok(drbg)
    }

    /// Override the reseed interval (useful for tests or hardened policies).
    pub fn set_reseed_interval(&mut self, interval: u64) {
        self.reseed_interval = interval.max(1);
    }

    /// Override the byte budget requiring reseed.
    pub fn set_max_bytes_between_reseed(&mut self, max_bytes: u128) {
        self.max_bytes_between_reseed = max_bytes.max(1);
    }

    /// Reseed with fresh entropy and optional additional input.
    pub fn reseed(&mut self, entropy: &[u8], additional_input: Option<&[u8]>) -> Result<(), Error> {
        self.check_new_entropy(entropy)?;
        let mut seed = Vec::with_capacity(entropy.len() + additional_input.map_or(0, |a| a.len()));
        seed.extend_from_slice(entropy);
        if let Some(ai) = additional_input {
            seed.extend_from_slice(ai);
        }
        self.update(Some(&seed));
        seed.zeroize();
        self.last_entropy.clear();
        self.last_entropy.extend_from_slice(entropy);
        self.reseed_counter = 1;
        self.generated_bytes = 0;
        Ok(())
    }

    /// Generate pseudorandom bytes. Optional `additional_input` mixes per-request data.
    pub fn generate(
        &mut self,
        out: &mut [u8],
        additional_input: Option<&[u8]>,
    ) -> Result<(), Error> {
        if out.len() > MAX_REQUEST {
            return Err(Error::RequestTooLarge);
        }
        if self.reseed_counter > self.reseed_interval {
            return Err(Error::ReseedRequired);
        }
        if self.generated_bytes + out.len() as u128 >= self.max_bytes_between_reseed {
            return Err(Error::ReseedRequired);
        }

        if let Some(ai) = additional_input {
            if !ai.is_empty() {
                self.update(Some(ai));
            }
        }

        let mut produced = 0usize;
        while produced < out.len() {
            self.v = self.hmac(&self.v);
            let take = min(BLOCK_LEN, out.len() - produced);
            out[produced..produced + take].copy_from_slice(&self.v[..take]);
            produced += take;
        }

        // Post-generation update (Section 10.1.2.5).
        if let Some(ai) = additional_input {
            if !ai.is_empty() {
                self.update(Some(ai));
            } else {
                self.update(None);
            }
        } else {
            self.update(None);
        }

        self.reseed_counter = self.reseed_counter.saturating_add(1);
        self.generated_bytes = self.generated_bytes.saturating_add(out.len() as u128);
        Ok(())
    }

    /// Compare internal state with another DRBG (useful for tests).
    pub fn ct_eq(&self, other: &Self) -> subtle::Choice {
        let mut acc = self.k.ct_eq(&other.k);
        acc &= self.v.ct_eq(&other.v);
        if self.reseed_counter == other.reseed_counter
            && self.reseed_interval == other.reseed_interval
            && self.generated_bytes == other.generated_bytes
            && self.max_bytes_between_reseed == other.max_bytes_between_reseed
        {
            acc &= subtle::Choice::from(1);
        } else {
            acc &= subtle::Choice::from(0);
        }

        if self.last_entropy.len() == other.last_entropy.len() {
            acc &= self
                .last_entropy
                .as_slice()
                .ct_eq(other.last_entropy.as_slice());
        } else {
            acc &= subtle::Choice::from(0);
        }
        acc
    }

    fn validate_entropy(entropy: &[u8]) -> Result<(), Error> {
        if entropy.is_empty() {
            return Err(Error::EntropyHealthFailed);
        }
        if entropy.iter().all(|&b| b == entropy[0]) {
            return Err(Error::EntropyHealthFailed);
        }
        Ok(())
    }

    fn check_new_entropy(&self, entropy: &[u8]) -> Result<(), Error> {
        Self::validate_entropy(entropy)?;
        if !self.last_entropy.is_empty() && entropy == self.last_entropy.as_slice() {
            return Err(Error::EntropyHealthFailed);
        }
        Ok(())
    }
    fn hmac(&self, data: &[u8]) -> [u8; BLOCK_LEN] {
        let mut mac = Hmac::<Sha512>::new_from_slice(&self.k).expect("HMAC key length valid");
        mac.update(data);
        let out = mac.finalize().into_bytes();
        let mut buf = [0u8; BLOCK_LEN];
        buf.copy_from_slice(&out);
        buf
    }

    fn update(&mut self, provided_data: Option<&[u8]>) {
        // Step (a): K = HMAC(K, V || 0x00 || provided)
        // Step (b): V = HMAC(K, V)
        let mut buf = Vec::with_capacity(self.v.len() + 1 + provided_data.map_or(0, |d| d.len()));
        buf.extend_from_slice(&self.v);
        buf.push(0x00);
        if let Some(data) = provided_data {
            buf.extend_from_slice(data);
        }
        let new_k = {
            let mut mac = Hmac::<Sha512>::new_from_slice(&self.k).expect("HMAC key length valid");
            mac.update(&buf);
            mac.finalize().into_bytes()
        };
        self.k.copy_from_slice(&new_k);
        self.v = self.hmac(&self.v);
        buf.zeroize();

        if let Some(data) = provided_data {
            let mut buf = Vec::with_capacity(self.v.len() + 1 + data.len());
            buf.extend_from_slice(&self.v);
            buf.push(0x01);
            buf.extend_from_slice(data);
            let new_k = {
                let mut mac =
                    Hmac::<Sha512>::new_from_slice(&self.k).expect("HMAC key length valid");
                mac.update(&buf);
                mac.finalize().into_bytes()
            };
            self.k.copy_from_slice(&new_k);
            self.v = self.hmac(&self.v);
            buf.zeroize();
        }
    }
}

#[cfg(feature = "aes")]
mod aes_ctr_drbg {
    use super::Error;
    use aes::Aes256;
    use ctr::cipher::{KeyIvInit, StreamCipher};
    use zeroize::{Zeroize, ZeroizeOnDrop};

    type Ctr128 = ctr::Ctr128BE<Aes256>;

    /// AES-256-CTR_DRBG scaffold (update/derivation per SP 800-90A still TODO).
    #[derive(Zeroize, ZeroizeOnDrop)]
    pub struct AesCtrDrbg {
        key: [u8; 32],
        v: [u8; 16],
        reseed_counter: u64,
        reseed_interval: u64,
    }

    impl AesCtrDrbg {
        pub fn new(seed_key: &[u8; 32], seed_v: &[u8; 16]) -> Self {
            Self {
                key: *seed_key,
                v: *seed_v,
                reseed_counter: 1,
                reseed_interval: super::DEFAULT_RESEED_INTERVAL,
            }
        }

        pub fn generate(&mut self, out: &mut [u8]) -> Result<(), Error> {
            if self.reseed_counter > self.reseed_interval {
                return Err(Error::ReseedRequired);
            }
            let mut cipher = Ctr128::new((&self.key).into(), (&self.v).into());
            cipher.apply_keystream(out);
            // Increment counter by number of blocks consumed.
            let blocks = (out.len() as u128 + 15) / 16;
            let ctr = u128::from_be_bytes(self.v).wrapping_add(blocks);
            self.v = ctr.to_be_bytes();
            self.reseed_counter = self.reseed_counter.saturating_add(1);
            Ok(())
        }
    }
}

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

    fn entropy(seed: u8) -> [u8; 48] {
        let mut out = [0u8; 48];
        for (i, byte) in out.iter_mut().enumerate() {
            *byte = seed.wrapping_add(i as u8);
        }
        out
    }

    fn nonce(seed: u8) -> [u8; 16] {
        let mut out = [0u8; 16];
        for (i, byte) in out.iter_mut().enumerate() {
            *byte = seed.wrapping_add((i * 3) as u8);
        }
        out
    }

    #[test]
    fn identical_seed_produces_identical_stream() {
        let mut a = HmacDrbg::new(&entropy(1), &nonce(2), Some(b"p")).unwrap();
        let mut b = HmacDrbg::new(&entropy(1), &nonce(2), Some(b"p")).unwrap();
        let mut out_a = [0u8; 96];
        let mut out_b = [0u8; 96];
        a.generate(&mut out_a, None).unwrap();
        b.generate(&mut out_b, None).unwrap();
        assert_eq!(out_a, out_b);
        assert!(bool::from(a.ct_eq(&b)));
    }

    #[test]
    fn reseed_changes_output() {
        let mut drbg = HmacDrbg::new(&entropy(3), &nonce(4), None).unwrap();
        let mut first = [0u8; 64];
        drbg.generate(&mut first, None).unwrap();
        drbg.reseed(&entropy(9), None).unwrap();
        let mut second = [0u8; 64];
        drbg.generate(&mut second, None).unwrap();
        assert_ne!(first, second);
    }

    #[test]
    fn additional_input_affects_stream() {
        let mut drbg = HmacDrbg::new(&entropy(5), &nonce(6), None).unwrap();
        let mut buf1 = [0u8; 64];
        let mut buf2 = [0u8; 64];
        drbg.generate(&mut buf1, Some(b"ai1")).unwrap();
        drbg.generate(&mut buf2, Some(b"ai2")).unwrap();
        assert_ne!(buf1, buf2);
    }

    #[test]
    fn request_too_large_fails() {
        let mut drbg = HmacDrbg::new(&entropy(7), &nonce(8), None).unwrap();
        let mut buf = vec![0u8; MAX_REQUEST + 1];
        assert_eq!(drbg.generate(&mut buf, None), Err(Error::RequestTooLarge));
    }

    #[test]
    fn reseed_interval_enforced() {
        let mut drbg = HmacDrbg::new(&entropy(9), &nonce(10), None).unwrap();
        drbg.set_reseed_interval(1);
        let mut buf = [0u8; 32];
        drbg.generate(&mut buf, None).unwrap();
        assert_eq!(drbg.generate(&mut buf, None), Err(Error::ReseedRequired));
    }
    #[test]
    fn repeated_entropy_fails_health() {
        let mut drbg = HmacDrbg::new(&entropy(1), &nonce(2), None).unwrap();
        assert!(matches!(
            drbg.reseed(&entropy(1), None),
            Err(Error::EntropyHealthFailed)
        ));
    }

    #[test]
    fn byte_budget_enforced() {
        let mut drbg = HmacDrbg::new(&entropy(11), &nonce(12), None).unwrap();
        drbg.set_max_bytes_between_reseed(64);
        let mut buf = [0u8; 32];
        drbg.generate(&mut buf, None).unwrap();
        assert!(matches!(
            drbg.generate(&mut buf, None),
            Err(Error::ReseedRequired)
        ));
    }
}