1 files changed, 148 insertions, 0 deletions
diff --git a/src/protocol/anti_replay.rs b/src/protocol/anti_replay.rs
new file mode 100644
index 0000000..4fb2e11
--- /dev/null
+++ b/src/protocol/anti_replay.rs
@@ -0,0 +1,148 @@
+// Copyright 2017 Sopium
+
+// This file is part of WireGuard.rs.
+
+// WireGuard.rs is free software: you can redistribute it and/or
+// modify it under the terms of the GNU General Public License as
+// published by the Free Software Foundation, either version 3 of the
+// License, or (at your option) any later version.
+
+// WireGuard.rs is distributed in the hope that it will be useful, but
+// WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+// General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with WireGuard.rs.  If not, see <https://www.gnu.org/licenses/>.
+
+
+// This is RFC 6479.
+
+use std::cmp::min;
+
+// Power of 2.
+const BITMAP_BITLEN: u64 = 2048;
+
+const SIZE_OF_INTEGER: u64 = 32;
+const BITMAP_LEN: usize = (BITMAP_BITLEN / SIZE_OF_INTEGER) as usize;
+const BITMAP_INDEX_MASK: u64 = BITMAP_LEN as u64 - 1;
+// REDUNDANT_BIT_SHIFTS = log2(SIZE_OF_INTEGER).
+const REDUNDANT_BIT_SHIFTS: u64 = 5;
+const BITMAP_LOC_MASK: u64 = SIZE_OF_INTEGER - 1;
+/// Size of anti-replay window.
+pub const WINDOW_SIZE: u64 = BITMAP_BITLEN - SIZE_OF_INTEGER;
+
+pub struct AntiReplay {
+    last: u64,
+    bitmap: [u32; BITMAP_LEN],
+}
+
+impl Default for AntiReplay {
+    fn default() -> Self {
+        AntiReplay::new()
+    }
+}
+
+impl AntiReplay {
+    pub fn new() -> Self {
+        AntiReplay {
+            last: 0,
+            bitmap: [0; BITMAP_LEN],
+        }
+    }
+
+    /// Returns true if check is passed, i.e., not a replay or too old.
+    ///
+    /// Unlike RFC 6479, zero is allowed.
+    pub fn check(&self, seq: u64) -> bool {
+        // Larger is always good.
+        if seq > self.last {
+            return true;
+        }
+
+        if self.last - seq > WINDOW_SIZE {
+            return false;
+        }
+
+        let bit_location = seq & BITMAP_LOC_MASK;
+        let index = (seq >> REDUNDANT_BIT_SHIFTS) & BITMAP_INDEX_MASK;
+
+        self.bitmap[index as usize] & (1 << bit_location) == 0
+    }
+
+    /// Should only be called if check returns true.
+    pub fn update(&mut self, seq: u64) {
+        debug_assert!(self.check(seq));
+
+        let index = seq >> REDUNDANT_BIT_SHIFTS;
+
+        if seq > self.last {
+            let index_cur = self.last >> REDUNDANT_BIT_SHIFTS;
+            let diff = min(index - index_cur, BITMAP_LEN as u64);
+
+            for i in 0..diff {
+                let real_index = (index_cur + i + 1) & BITMAP_INDEX_MASK;
+                self.bitmap[real_index as usize] = 0;
+            }
+            self.last = seq;
+        }
+
+        let index = index & BITMAP_INDEX_MASK;
+        let bit_location = seq & BITMAP_LOC_MASK;
+        self.bitmap[index as usize] |= 1 << bit_location;
+    }
+
+    pub fn check_and_update(&mut self, seq: u64) -> bool {
+        if self.check(seq) {
+            self.update(seq);
+            true
+        } else {
+            false
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn anti_replay() {
+        let mut ar = AntiReplay::new();
+
+        for i in 0..20000 {
+            assert!(ar.check_and_update(i));
+        }
+
+        for i in (0..20000).rev() {
+            assert!(!ar.check(i));
+        }
+
+        assert!(ar.check_and_update(65536));
+        for i in (65536 - WINDOW_SIZE)..65535 {
+            assert!(ar.check_and_update(i));
+        }
+        for i in (65536 - 10 * WINDOW_SIZE)..65535 {
+            assert!(!ar.check(i));
+        }
+
+        ar.check_and_update(66000);
+        for i in 65537..66000 {
+            assert!(ar.check_and_update(i));
+        }
+        for i in 65537..66000 {
+            assert!(!ar.check_and_update(i));
+        }
+
+        // Test max u64.
+        let next = u64::max_value();
+        assert!(ar.check_and_update(next));
+        assert!(!ar.check(next));
+        for i in (next - WINDOW_SIZE)..next {
+            assert!(ar.check_and_update(i));
+        }
+        for i in (next - 20 * WINDOW_SIZE)..next {
+            assert!(!ar.check(i));
+        }
+    }
+}