shoes 0.3.2

A multi-protocol proxy server.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124

use std::collections::{HashSet, VecDeque};
use std::sync::Arc;
use std::time::Duration;

use parking_lot::Mutex;
use tokio::time::Instant;

use super::salt_checker::SaltChecker;

#[derive(Debug)]
struct TimeEntry {
    instant: Instant,
    salt: Arc<[u8]>,
}

#[derive(Debug)]
pub struct TimedSaltChecker {
    state: Arc<Mutex<SaltState>>,
    cleanup: Option<tokio::task::AbortHandle>,
}

#[derive(Debug)]
struct SaltState {
    last_salts: VecDeque<TimeEntry>,
    known_salts: HashSet<Arc<[u8]>>,
    timeout_secs: u64,
    capacity: usize,
}

impl TimedSaltChecker {
    pub fn new(timeout_secs: u64) -> Self {
        Self::with_capacity(timeout_secs, 65536)
    }

    fn with_capacity(timeout_secs: u64, capacity: usize) -> Self {
        let state = Arc::new(Mutex::new(SaltState {
            last_salts: VecDeque::new(),
            known_salts: HashSet::new(),
            timeout_secs,
            capacity,
        }));
        let weak = Arc::downgrade(&state);
        let cleanup = tokio::runtime::Handle::try_current().ok().map(|runtime| {
            runtime
                .spawn(async move {
                    let mut interval = tokio::time::interval(Duration::from_secs(1));
                    loop {
                        interval.tick().await;
                        let Some(state) = weak.upgrade() else { break };
                        state.lock().expire();
                    }
                })
                .abort_handle()
        });
        Self { state, cleanup }
    }
}

impl Drop for TimedSaltChecker {
    fn drop(&mut self) {
        if let Some(task) = &self.cleanup {
            task.abort();
        }
    }
}

impl SaltState {
    fn expire(&mut self) {
        while let Some(time_entry) = self.last_salts.front() {
            if time_entry.instant.elapsed().as_secs() < self.timeout_secs {
                break;
            }
            self.known_salts.remove(&time_entry.salt);
            self.last_salts.pop_front();
        }
        if self.known_salts.len() < self.known_salts.capacity() / 4 {
            self.known_salts.shrink_to(self.known_salts.len().max(128));
            self.last_salts.shrink_to(self.last_salts.len().max(128));
        }
    }
}

impl SaltChecker for TimedSaltChecker {
    fn insert_and_check(&mut self, salt: &[u8]) -> bool {
        let mut state = self.state.lock();
        state.expire();
        // Reject new work at capacity. Evicting live entries would admit replays.
        if state.known_salts.contains(salt) || state.known_salts.len() >= state.capacity {
            return false;
        }

        let salt: Arc<[u8]> = salt.into();
        state.known_salts.insert(salt.clone());
        state.last_salts.push_back(TimeEntry {
            instant: Instant::now(),
            salt,
        });

        true
    }
}

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

    #[tokio::test(start_paused = true)]
    async fn capacity_is_replay_safe_and_idle_entries_expire() {
        let mut checker = TimedSaltChecker::with_capacity(60, 2);
        assert!(checker.insert_and_check(b"first"));
        assert!(checker.insert_and_check(b"second"));
        assert!(!checker.insert_and_check(b"third"));
        assert!(!checker.insert_and_check(b"first"));
        tokio::task::yield_now().await;
        tokio::time::advance(Duration::from_secs(61)).await;
        tokio::task::yield_now().await;
        assert!(checker.state.lock().known_salts.is_empty());
        assert!(checker.insert_and_check(b"third"));
        let task = checker.cleanup.clone().unwrap();
        drop(checker);
        tokio::task::yield_now().await;
        assert!(task.is_finished());
    }
}