Files
wz-phone/crates/wzp-crypto/src/anti_replay.rs

263 lines
8.2 KiB
Rust

//! Sliding window replay protection.
//!
//! Tracks seen sequence numbers using a bitmap. Window size is configurable
//! at construction time. Sequence numbers that are too old (more than
//! `window_size` behind the highest seen) are rejected.
use wzp_proto::CryptoError;
/// Sliding window anti-replay detector.
///
/// Uses a bitmap to track which sequence numbers have been seen within
/// the current window. Handles `u32` wrapping correctly.
pub struct AntiReplayWindow {
/// Window size in packets.
window_size: u32,
/// Highest sequence number seen so far.
highest: u32,
/// Bitmap of seen packets. Bit i corresponds to (highest - i).
bitmap: Vec<u64>,
/// Whether any packet has been received yet.
initialized: bool,
}
impl AntiReplayWindow {
/// Create a new anti-replay window with the default size of 1024 packets.
pub fn new() -> Self {
Self::with_window(1024)
}
/// Create a new anti-replay window with a custom size.
pub fn with_window(size: usize) -> Self {
let window_size = size as u32;
let bitmap_words = (size + 63) / 64;
Self {
window_size,
highest: 0,
bitmap: vec![0u64; bitmap_words],
initialized: false,
}
}
/// Check if a sequence number is valid (not a replay, not too old).
/// If valid, marks it as seen.
pub fn check_and_update(&mut self, seq: u32) -> Result<(), CryptoError> {
if !self.initialized {
self.initialized = true;
self.highest = seq;
self.set_bit(0);
return Ok(());
}
let diff = seq.wrapping_sub(self.highest);
if diff == 0 {
// Duplicate of highest
return Err(CryptoError::ReplayDetected { seq });
}
if diff < 0x8000_0000 {
// seq is ahead of highest (wrapping-aware: diff in [1, 0x7FFF_FFFF])
let shift = diff as usize;
self.advance_window(shift);
self.highest = seq;
self.set_bit(0);
Ok(())
} else {
// seq is behind highest (wrapping-aware: diff in [0x8000_0000, 0xFFFF_FFFF])
let behind = self.highest.wrapping_sub(seq) as usize;
if behind >= self.window_size as usize {
return Err(CryptoError::ReplayDetected { seq });
}
if self.get_bit(behind) {
return Err(CryptoError::ReplayDetected { seq });
}
self.set_bit(behind);
Ok(())
}
}
/// Advance the window by `shift` positions (shift left = new bits at position 0).
fn advance_window(&mut self, shift: usize) {
let window_size = self.window_size as usize;
if shift >= window_size {
for word in &mut self.bitmap {
*word = 0;
}
return;
}
// We need to shift the entire bitmap right by `shift` bits.
// Bit 0 of word 0 is the most recent. Shifting right means
// old entries move to higher bit positions.
let word_shift = shift / 64;
let bit_shift = shift % 64;
// Move words
let len = self.bitmap.len();
for i in (0..len).rev() {
let mut val = 0u64;
if i >= word_shift {
val = self.bitmap[i - word_shift] << bit_shift;
if bit_shift > 0 && i > word_shift {
val |= self.bitmap[i - word_shift - 1] >> (64 - bit_shift);
}
}
self.bitmap[i] = val;
}
// Clear the lower words that shifted in
for word in &mut self.bitmap[..word_shift.min(len)] {
*word = 0;
}
// Clear the lower bits of the first non-shifted word
if word_shift < len && bit_shift > 0 {
self.bitmap[word_shift] &= !((1u64 << bit_shift) - 1);
}
}
fn set_bit(&mut self, offset: usize) {
let word = offset / 64;
let bit = offset % 64;
if word < self.bitmap.len() {
self.bitmap[word] |= 1u64 << bit;
}
}
fn get_bit(&self, offset: usize) -> bool {
let word = offset / 64;
let bit = offset % 64;
if word < self.bitmap.len() {
(self.bitmap[word] >> bit) & 1 == 1
} else {
false
}
}
}
impl Default for AntiReplayWindow {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn first_packet_accepted() {
let mut w = AntiReplayWindow::new();
assert!(w.check_and_update(0).is_ok());
}
#[test]
fn duplicate_rejected() {
let mut w = AntiReplayWindow::new();
assert!(w.check_and_update(100).is_ok());
assert!(w.check_and_update(100).is_err());
}
#[test]
fn sequential_accepted() {
let mut w = AntiReplayWindow::new();
for i in 0..200 {
assert!(
w.check_and_update(i).is_ok(),
"seq {} should be accepted",
i
);
}
}
#[test]
fn out_of_order_within_window() {
let mut w = AntiReplayWindow::new();
assert!(w.check_and_update(100).is_ok());
assert!(w.check_and_update(95).is_ok());
assert!(w.check_and_update(98).is_ok());
assert!(w.check_and_update(102).is_ok());
assert!(w.check_and_update(99).is_ok());
}
#[test]
fn old_packet_rejected() {
let mut w = AntiReplayWindow::new();
assert!(w.check_and_update(0).is_ok());
// Advance well past the window
assert!(w.check_and_update(2000).is_ok());
// seq 0 is now too old
assert!(w.check_and_update(0).is_err());
}
#[test]
fn wrapping_works() {
let mut w = AntiReplayWindow::new();
assert!(w.check_and_update(0xFFFF_FFF0).is_ok());
assert!(w.check_and_update(0xFFFF_FFFF).is_ok());
assert!(w.check_and_update(0).is_ok()); // wrapped
assert!(w.check_and_update(1).is_ok());
assert!(w.check_and_update(0xFFFF_FFFF).is_err()); // duplicate
}
#[test]
fn within_window_boundary() {
let mut w = AntiReplayWindow::new();
assert!(w.check_and_update(1023).is_ok());
// 1023 - 0 = 1023, exactly at window boundary
assert!(w.check_and_update(0).is_ok());
// But 1024 behind would be out
assert!(w.check_and_update(1024).is_ok());
// Now 0 is 1024 behind 1024, which is at the boundary limit
assert!(w.check_and_update(0).is_err()); // already seen or too old
}
#[test]
fn custom_window_size() {
let mut w = AntiReplayWindow::with_window(64);
for i in 0..64 {
assert!(w.check_and_update(i).is_ok());
}
// seq 0 is now exactly at the boundary (64 behind 64)
assert!(w.check_and_update(0).is_err());
}
#[test]
fn video_burst_200_with_one_reorder() {
let mut w = AntiReplayWindow::with_window(1024);
// Simulate a 200-packet burst
for i in 0..200 {
assert!(
w.check_and_update(i).is_ok(),
"seq {} should be accepted",
i
);
}
// One packet reordered (arrives late)
assert!(w.check_and_update(50).is_err(), "seq 50 is a duplicate");
// But a packet just behind the window should still be ok
assert!(w.check_and_update(199).is_err(), "seq 199 is a duplicate");
// Continue the burst
for i in 200..400 {
assert!(
w.check_and_update(i).is_ok(),
"seq {} should be accepted",
i
);
}
}
#[test]
fn u32_high_range_works() {
let mut w = AntiReplayWindow::with_window(64);
let base = 1000u32;
assert!(w.check_and_update(base).is_ok());
assert!(w.check_and_update(base + 1).is_ok());
// 65 behind highest (base+1) is outside the 64-packet window
assert!(w.check_and_update(base.wrapping_sub(64)).is_err());
// 63 behind is inside
assert!(w.check_and_update(base.wrapping_sub(62)).is_ok());
// base itself is now a duplicate
assert!(w.check_and_update(base).is_err());
}
}