stochastic-routing-extended 1.0.2

//! Deterministic PRNG seeded from the shared session seed.
//!
//! Draws are produced by SHA-256(`seed ‖ counter`) so both peers obtain the
//! same sequence if they call the same methods in the same order.

use sha2::{Digest, Sha256};

/// A deterministic RNG that both peers derive identically from the shared seed.
///
/// Used to make synchronized stochastic decisions (port hops, transport
/// selection, jitter values) without exchanging additional messages.
pub struct SeedRng {
    _seed: [u8; 32],
    _counter: u64,
}

impl SeedRng {
    /// Create a new RNG from a 32-byte seed.
    pub fn new(seed: [u8; 32]) -> Self {
        Self {
            _seed: seed,
            _counter: 0,
        }
    }

    /// Shared session seed (both peers hold the same value).
    #[must_use]
    pub fn seed_bytes(&self) -> [u8; 32] {
        self._seed
    }

    /// Replace the seed and reset the counter (e.g. after [`crate::client::SelfHealing::heal`]).
    pub fn reseed(&mut self, new_seed: [u8; 32]) {
        self._seed = new_seed;
        self._counter = 0;
    }

    /// Generate the next pseudo-random u64 (deterministic stream).
    pub fn next_u64(&mut self) -> u64 {
        self._counter = self._counter.wrapping_add(1);
        let mut hasher = Sha256::new();
        hasher.update(self._seed);
        hasher.update(self._counter.to_be_bytes());
        let digest = hasher.finalize();
        u64::from_be_bytes(digest[..8].try_into().expect("digest length"))
    }

    /// Generate a value in `[low, high)` (half-open).
    pub fn range(&mut self, low: u64, high: u64) -> u64 {
        if high <= low {
            return low;
        }
        low + self.next_u64() % (high - low)
    }

    /// Select a random element from a slice.
    pub fn choose<'a, T>(&mut self, items: &'a [T]) -> Option<&'a T> {
        if items.is_empty() {
            return None;
        }
        let idx = self.range(0, items.len() as u64) as usize;
        Some(&items[idx])
    }
}

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

    #[test]
    fn stream_is_deterministic() {
        let mut a = SeedRng::new([9u8; 32]);
        let mut b = SeedRng::new([9u8; 32]);
        assert_eq!(a.next_u64(), b.next_u64());
        assert_eq!(a.next_u64(), b.next_u64());
    }

    #[test]
    fn different_seeds_diverge() {
        let mut a = SeedRng::new([1u8; 32]);
        let mut b = SeedRng::new([2u8; 32]);
        assert_ne!(a.next_u64(), b.next_u64());
    }
}