eggress_protocol_shadowsocks/
nonce.rs

1use crate::ShadowsocksError;
2
3#[derive(Debug, Clone)]
4pub struct NonceCounter {
5    nonce_size: usize,
6    counter: u64,
7}
8
9impl NonceCounter {
10    pub fn new(nonce_size: usize) -> Self {
11        Self {
12            nonce_size,
13            counter: 0,
14        }
15    }
16
17    pub fn starting_at(nonce_size: usize, counter: u64) -> Self {
18        Self {
19            nonce_size,
20            counter,
21        }
22    }
23
24    pub fn current(&self) -> Vec<u8> {
25        let mut buf = vec![0u8; self.nonce_size];
26        let end = self.nonce_size.min(8);
27        buf[..end].copy_from_slice(&self.counter.to_le_bytes()[..end]);
28        buf
29    }
30
31    pub fn advance(&mut self) -> Result<(), ShadowsocksError> {
32        self.counter = self
33            .counter
34            .checked_add(1)
35            .ok_or_else(|| ShadowsocksError::Other("nonce counter overflow".into()))?;
36        Ok(())
37    }
38
39    pub fn nonce_size(&self) -> usize {
40        self.nonce_size
41    }
42}
43
44#[cfg(test)]
45mod tests {
46    use super::*;
47
48    #[test]
49    fn test_nonce_starts_at_zero() {
50        let nonce = NonceCounter::new(12);
51        assert_eq!(nonce.current(), vec![0u8; 12]);
52    }
53
54    #[test]
55    fn test_nonce_increments() {
56        let mut nonce = NonceCounter::new(12);
57        nonce.advance().unwrap();
58        let bytes = nonce.current();
59        assert_eq!(bytes.len(), 12);
60        // little-endian: counter in first 8 bytes, rest zero
61        assert_eq!(&bytes[..8], &1u64.to_le_bytes());
62        assert_eq!(&bytes[8..], &[0, 0, 0, 0]);
63    }
64
65    #[test]
66    fn test_nonce_advance_multiple() {
67        let mut nonce = NonceCounter::new(12);
68        for i in 1u64..=10 {
69            nonce.advance().unwrap();
70            let bytes = nonce.current();
71            assert_eq!(&bytes[..8], &i.to_le_bytes());
72        }
73    }
74
75    #[test]
76    fn test_nonce_overflow_returns_error() {
77        let mut nonce = NonceCounter::new(12);
78        nonce.counter = u64::MAX;
79        assert!(nonce.advance().is_err());
80    }
81}
eggress_protocol_shadowsocks/nonce.rs

eggress_protocol_shadowsocks/
nonce.rs
