stochastic-routing-extended 1.0.2

//! Anti-replay sliding window — prevents processing of duplicate or replayed packets.
//!
//! Uses a bitmap-based sliding window (RFC 6479 style) that tracks which packet
//! counters have been seen. Rejects duplicates and counters that fall behind the
//! window.

/// Size of the replay window in bits.
const WINDOW_SIZE: u64 = 1024;

/// Number of `u64` words in the bitmap.
const BITMAP_WORDS: usize = (WINDOW_SIZE / 64) as usize;

/// Sliding window anti-replay detector.
///
/// Tracks which packet counters have been received. A counter is rejected if:
/// - It has already been seen (duplicate), or
/// - It falls before the window (too old).
///
/// The window advances automatically when a counter beyond the current
/// right edge is accepted.
pub struct ReplayWindow {
    /// Highest counter-value seen so far.
    top: u64,
    /// Bitmap of received counters in `[top - WINDOW_SIZE + 1 ... top]`.
    bitmap: [u64; BITMAP_WORDS],
}

/// Serializable-ish snapshot of replay-window state.
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct ReplayState {
    /// Highest counter-value seen.
    pub top: u64,
    /// Bitmap words (length must match window word count).
    pub bitmap_words: Vec<u64>,
}

impl ReplayWindow {
    /// Create a new empty replay window.
    pub fn new() -> Self {
        Self {
            top: 0,
            bitmap: [0u64; BITMAP_WORDS],
        }
    }

    /// Check if a counter is acceptable (not replayed, not too old).
    ///
    /// Returns `true` if the counter can be accepted, `false` if it should
    /// be rejected (duplicate or too old).
    ///
    /// **Does not mark the counter as seen** — call [`accept`] after
    /// successful decryption.
    pub fn check(&self, counter: u64) -> bool {
        if counter == 0 {
            // Counter 0 is reserved / never sent.
            return false;
        }
        if self.top == 0 {
            // No packets seen yet — accept anything > 0.
            return true;
        }
        if counter > self.top {
            // Ahead of the window — always acceptable.
            return true;
        }
        let diff = self.top - counter;
        if diff >= WINDOW_SIZE {
            // Too old — behind the window.
            return false;
        }
        // Check bitmap for duplicate.
        let word_idx = (diff as usize) / 64;
        let bit_idx = (diff as usize) % 64;
        (self.bitmap[word_idx] >> bit_idx) & 1 == 0
    }

    /// Mark a counter as received after successful authentication.
    ///
    /// Call this only after `check()` returned `true` AND the packet was
    /// successfully decrypted/authenticated. This prevents an attacker from
    /// poisoning the window with invalid packets.
    pub fn accept(&mut self, counter: u64) {
        if counter == 0 {
            return;
        }
        if self.top == 0 || counter > self.top {
            // Advance the window.
            let shift = if self.top == 0 { 0 } else { counter - self.top };
            self.shift_window(shift);
            self.top = counter;
            // Bit 0 (offset 0 from top) is the new counter — mark it.
            self.bitmap[0] |= 1;
        } else {
            // Within window — set the bit.
            let diff = self.top - counter;
            if diff < WINDOW_SIZE {
                let word_idx = (diff as usize) / 64;
                let bit_idx = (diff as usize) % 64;
                self.bitmap[word_idx] |= 1 << bit_idx;
            }
        }
    }

    /// Combined check + accept. Returns `true` if accepted, `false` if rejected.
    ///
    /// Use this only when authentication is guaranteed before calling (e.g. after
    /// AEAD decryption succeeds).
    pub fn check_and_accept(&mut self, counter: u64) -> bool {
        if !self.check(counter) {
            return false;
        }
        self.accept(counter);
        true
    }

    /// Current top (highest seen counter).
    pub fn top(&self) -> u64 {
        self.top
    }

    /// Window size in bits.
    pub fn window_size(&self) -> u64 {
        WINDOW_SIZE
    }

    /// Export current replay-window state.
    pub fn snapshot(&self) -> ReplayState {
        ReplayState {
            top: self.top,
            bitmap_words: self.bitmap.to_vec(),
        }
    }

    /// Restore replay-window state.
    ///
    /// Returns `false` if the snapshot shape is invalid.
    pub fn restore(&mut self, state: &ReplayState) -> bool {
        if state.bitmap_words.len() != BITMAP_WORDS {
            return false;
        }
        self.top = state.top;
        self.bitmap.copy_from_slice(&state.bitmap_words);
        true
    }

    /// Shift the bitmap right by `shift` positions (advancing the window).
    fn shift_window(&mut self, shift: u64) {
        if shift == 0 {
            return;
        }
        if shift >= WINDOW_SIZE {
            // The entire window is outdated — clear everything.
            self.bitmap = [0u64; BITMAP_WORDS];
            return;
        }

        let word_shift = (shift / 64) as usize;
        let bit_shift = (shift % 64) as u32;

        if word_shift > 0 {
            // Shift whole words.
            for i in (word_shift..BITMAP_WORDS).rev() {
                self.bitmap[i] = self.bitmap[i - word_shift];
            }
            for i in 0..word_shift.min(BITMAP_WORDS) {
                self.bitmap[i] = 0;
            }
        }

        if bit_shift > 0 {
            // Shift remaining bits within words (from high index to low).
            for i in (1..BITMAP_WORDS).rev() {
                self.bitmap[i] =
                    (self.bitmap[i] << bit_shift) | (self.bitmap[i - 1] >> (64 - bit_shift));
            }
            self.bitmap[0] <<= bit_shift;
        }
    }
}

impl Default for ReplayWindow {
    fn default() -> Self {
        Self::new()
    }
}

/// Monotonic merge for two replay snapshots.
///
/// The merged state preserves all counters seen in either input within the
/// resulting sliding window.
pub fn merge_replay_states(a: &ReplayState, b: &ReplayState) -> Option<ReplayState> {
    if a.bitmap_words.len() != BITMAP_WORDS || b.bitmap_words.len() != BITMAP_WORDS {
        return None;
    }
    let top = a.top.max(b.top);
    if top == 0 {
        return Some(ReplayState {
            top: 0,
            bitmap_words: vec![0u64; BITMAP_WORDS],
        });
    }

    let mut merged = ReplayWindow::new();
    merged.top = top;

    let left = top.saturating_sub(WINDOW_SIZE - 1);
    let mut counter = left;
    while counter <= top {
        if state_contains(a, counter) || state_contains(b, counter) {
            let diff = top - counter;
            let word_idx = (diff as usize) / 64;
            let bit_idx = (diff as usize) % 64;
            merged.bitmap[word_idx] |= 1 << bit_idx;
        }
        if counter == u64::MAX {
            break;
        }
        counter += 1;
    }
    Some(merged.snapshot())
}

fn state_contains(state: &ReplayState, counter: u64) -> bool {
    if counter == 0 || state.top == 0 || counter > state.top {
        return false;
    }
    let diff = state.top - counter;
    if diff >= WINDOW_SIZE {
        return false;
    }
    let word_idx = (diff as usize) / 64;
    let bit_idx = (diff as usize) % 64;
    (state.bitmap_words[word_idx] >> bit_idx) & 1 == 1
}

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

    #[test]
    fn fresh_window_accepts_first_packet() {
        let mut w = ReplayWindow::new();
        assert!(w.check_and_accept(1));
        assert_eq!(w.top(), 1);
    }

    #[test]
    fn rejects_counter_zero() {
        let w = ReplayWindow::new();
        assert!(!w.check(0));
    }

    #[test]
    fn rejects_duplicate() {
        let mut w = ReplayWindow::new();
        assert!(w.check_and_accept(5));
        assert!(!w.check_and_accept(5));
    }

    #[test]
    fn accepts_sequential_counters() {
        let mut w = ReplayWindow::new();
        for i in 1..=100 {
            assert!(w.check_and_accept(i), "should accept {i}");
        }
        assert_eq!(w.top(), 100);
    }

    #[test]
    fn accepts_out_of_order_within_window() {
        let mut w = ReplayWindow::new();
        w.check_and_accept(10);
        w.check_and_accept(8);
        w.check_and_accept(12);
        w.check_and_accept(9);
        // All should be marked.
        assert!(!w.check(8));
        assert!(!w.check(9));
        assert!(!w.check(10));
        assert!(!w.check(12));
        // 11 was never received.
        assert!(w.check(11));
    }

    #[test]
    fn rejects_too_old() {
        let mut w = ReplayWindow::new();
        w.check_and_accept(WINDOW_SIZE + 100);
        // Counter 1 is now behind the window.
        assert!(!w.check(1));
        // Counter just at the edge.
        assert!(!w.check(100));
    }

    #[test]
    fn window_edge_accepted() {
        let mut w = ReplayWindow::new();
        w.check_and_accept(WINDOW_SIZE);
        // Counter 1 is exactly at the left edge of the window.
        assert!(w.check_and_accept(1));
    }

    #[test]
    fn large_jump_clears_window() {
        let mut w = ReplayWindow::new();
        for i in 1..=50 {
            w.check_and_accept(i);
        }
        // Jump far ahead — all old counters should be behind the window.
        w.check_and_accept(WINDOW_SIZE + 200);
        assert!(!w.check(50));
        assert!(!w.check(1));
    }

    #[test]
    fn sequential_high_volume() {
        let mut w = ReplayWindow::new();
        for i in 1..=5000 {
            assert!(w.check_and_accept(i), "should accept {i}");
        }
        // Recent packets should be rejected as duplicates.
        assert!(!w.check(5000));
        assert!(!w.check(4999));
        // Very old should be behind the window.
        assert!(!w.check(1));
        // The next packet should work.
        assert!(w.check_and_accept(5001));
    }

    #[test]
    fn check_without_accept_does_not_mark() {
        let mut w = ReplayWindow::new();
        w.check_and_accept(10);
        assert!(w.check(11)); // check but don't accept
        assert!(w.check(11)); // should still pass — not marked
    }

    #[test]
    fn window_size_is_1024() {
        let w = ReplayWindow::new();
        assert_eq!(w.window_size(), 1024);
    }

    #[test]
    fn interleaved_gaps_handled() {
        let mut w = ReplayWindow::new();
        // Receive only even counters.
        for i in (2..=200).step_by(2) {
            assert!(w.check_and_accept(i));
        }
        // Odd counters within the window should still be acceptable.
        assert!(w.check_and_accept(199));
        assert!(w.check_and_accept(197));
        // But even ones should be rejected (duplicates).
        assert!(!w.check(200));
        assert!(!w.check(198));
    }

    #[test]
    fn snapshot_restore_roundtrip() {
        let mut w = ReplayWindow::new();
        for i in [10, 11, 20, 18, 19] {
            assert!(w.check_and_accept(i));
        }

        let snap = w.snapshot();
        let mut restored = ReplayWindow::new();
        assert!(restored.restore(&snap));
        assert_eq!(restored.top(), w.top());
        assert_eq!(restored.snapshot(), snap);

        // Seen packets remain blocked.
        assert!(!restored.check(20));
        assert!(!restored.check(19));
        // Fresh packet still accepted.
        assert!(restored.check_and_accept(21));
    }

    #[test]
    fn restore_rejects_invalid_bitmap_size() {
        let mut w = ReplayWindow::new();
        let bad = ReplayState {
            top: 42,
            bitmap_words: vec![0; 3],
        };
        assert!(!w.restore(&bad));
    }

    #[test]
    fn merge_replay_states_keeps_union_within_window() {
        let mut w1 = ReplayWindow::new();
        let mut w2 = ReplayWindow::new();
        for c in [10, 11, 13, 21] {
            w1.check_and_accept(c);
        }
        for c in [12, 20, 21, 30] {
            w2.check_and_accept(c);
        }
        let merged = merge_replay_states(&w1.snapshot(), &w2.snapshot()).expect("merge");
        let mut rw = ReplayWindow::new();
        assert!(rw.restore(&merged));

        for seen in [10, 11, 12, 13, 20, 21, 30] {
            assert!(!rw.check(seen), "merged snapshot must contain {seen}");
        }
        assert!(rw.check(31), "next unseen packet must be acceptable");
    }
}