stochastic-routing-extended 1.0.2

//! Routing mask generation from shared seed + session counter.
//!
//! The routing mask determines which transport(s) carry each frame
//! or fragment, enabling multi-transport splitting of a single logical packet.

use sha2::{Digest, Sha256};

use crate::seed::SeedRng;
use crate::transport::TransportKind;

/// A routing mask that maps a frame to one or more transports.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RoutingMask {
    /// Raw mask value.
    pub value: u32,
    /// Selected transports for this frame.
    pub transports: Vec<TransportKind>,
}

impl RoutingMask {
    /// Generate a routing mask for the given frame counter.
    ///
    /// Deterministic from `(seed, frame_counter, available)` so peers agree without extra messages.
    pub fn generate(
        rng: &mut SeedRng,
        available_transports: &[TransportKind],
        frame_counter: u64,
    ) -> Self {
        if available_transports.is_empty() {
            return Self {
                value: 0,
                transports: Vec::new(),
            };
        }

        let seed = rng.seed_bytes();
        let mut hasher = Sha256::new();
        hasher.update(seed);
        hasher.update(frame_counter.to_be_bytes());
        let digest = hasher.finalize();

        let value = u32::from_be_bytes(digest[0..4].try_into().expect("digest prefix"));

        let mut selected = Vec::new();
        for (i, &t) in available_transports.iter().enumerate() {
            let byte = digest[(i % 32) + 4];
            if (byte & 1) != 0 {
                selected.push(t);
            }
        }

        if selected.is_empty() {
            let idx = (value as usize) % available_transports.len();
            selected.push(available_transports[idx]);
        }

        Self {
            value,
            transports: selected,
        }
    }
}

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

    #[test]
    fn peers_agree() {
        let kinds = [TransportKind::Tcp, TransportKind::Udp, TransportKind::Quic];
        let mut a = SeedRng::new([0xAAu8; 32]);
        let mut b = SeedRng::new([0xAAu8; 32]);
        let ma = RoutingMask::generate(&mut a, &kinds, 42);
        let mb = RoutingMask::generate(&mut b, &kinds, 42);
        assert_eq!(ma, mb);
    }

    #[test]
    fn empty_available() {
        let mut rng = SeedRng::new([1u8; 32]);
        let m = RoutingMask::generate(&mut rng, &[], 0);
        assert!(m.transports.is_empty());
    }
}