superfluid_daemon/
completion_bucket.rs

1//! Keystroke-storm protection for FIM completions.
2
3use std::collections::HashMap;
4use std::hash::{Hash, Hasher};
5use std::sync::Mutex;
6use std::time::{Duration, Instant};
7
8pub const DEFAULT_RATE_PER_SEC: f64 = 10.0;
9pub const DEFAULT_BURST: u32 = 20;
10
11const MAX_CLIENTS: usize = 100_000;
12const IDLE_EVICT: Duration = Duration::from_secs(300);
13const SWEEP_EVERY: u32 = 1000;
14const MIN_RATE_PER_SEC: f64 = 0.001;
15const MAX_WAIT: Duration = Duration::from_secs(24 * 3600);
16
17#[derive(Debug, Clone, Copy, PartialEq)]
18pub struct BucketConfig {
19    pub rate_per_sec: f64,
20    pub burst: u32,
21}
22
23impl Default for BucketConfig {
24    fn default() -> BucketConfig {
25        BucketConfig {
26            rate_per_sec: DEFAULT_RATE_PER_SEC,
27            burst: DEFAULT_BURST,
28        }
29    }
30}
31
32impl BucketConfig {
33    pub const OFF: BucketConfig = BucketConfig {
34        rate_per_sec: 0.0,
35        burst: 0,
36    };
37
38    pub fn enabled(&self) -> bool {
39        self.rate_per_sec > 0.0
40    }
41
42    pub fn validate(&self) -> Result<(), &'static str> {
43        if !self.rate_per_sec.is_finite() || self.rate_per_sec < 0.0 {
44            return Err("completion rate must be finite and non-negative");
45        }
46        if self.enabled() && self.rate_per_sec < MIN_RATE_PER_SEC {
47            return Err("completion rate must be 0 (off) or at least 0.001 per second");
48        }
49        if self.enabled() && self.burst == 0 {
50            return Err("completion burst must be at least 1 when the rate is non-zero");
51        }
52        Ok(())
53    }
54}
55
56#[derive(Debug, Clone)]
57pub struct TokenBucket {
58    tokens: f64,
59    last: Instant,
60}
61
62impl TokenBucket {
63    pub fn new(cfg: BucketConfig, now: Instant) -> TokenBucket {
64        TokenBucket {
65            tokens: cfg.burst as f64,
66            last: now,
67        }
68    }
69
70    pub fn try_take(&mut self, cfg: BucketConfig, now: Instant) -> Result<(), Duration> {
71        if !cfg.enabled() {
72            return Ok(());
73        }
74        let elapsed = now.saturating_duration_since(self.last).as_secs_f64();
75        self.tokens = (self.tokens + elapsed * cfg.rate_per_sec).min(cfg.burst as f64);
76        self.last = now;
77        if self.tokens >= 1.0 - 1e-9 {
78            self.tokens = (self.tokens - 1.0).max(0.0);
79            return Ok(());
80        }
81        let wait = (1.0 - self.tokens) / cfg.rate_per_sec;
82        Err(Duration::try_from_secs_f64(wait).map_or(MAX_WAIT, |d| d.min(MAX_WAIT)))
83    }
84}
85
86pub struct ClientBuckets {
87    state: Mutex<State>,
88}
89
90struct State {
91    buckets: HashMap<u64, TokenBucket>,
92    sweep_counter: u32,
93}
94
95impl Default for ClientBuckets {
96    fn default() -> ClientBuckets {
97        ClientBuckets::new()
98    }
99}
100
101impl ClientBuckets {
102    pub fn new() -> ClientBuckets {
103        ClientBuckets {
104            state: Mutex::new(State {
105                buckets: HashMap::new(),
106                sweep_counter: 0,
107            }),
108        }
109    }
110
111    pub fn try_take(&self, client: u64, cfg: BucketConfig) -> Result<(), Duration> {
112        if !cfg.enabled() {
113            return Ok(());
114        }
115        let now = Instant::now();
116        let mut st = self.state.lock().expect("completion buckets");
117        st.sweep_counter += 1;
118        if st.sweep_counter >= SWEEP_EVERY {
119            st.sweep_counter = 0;
120            st.buckets.retain(|_, b| now.saturating_duration_since(b.last) <= IDLE_EVICT);
121        }
122        if st.buckets.len() >= MAX_CLIENTS && !st.buckets.contains_key(&client) {
123            st.buckets.retain(|_, b| now.saturating_duration_since(b.last) <= IDLE_EVICT);
124            if st.buckets.len() >= MAX_CLIENTS {
125                let mut stamps: Vec<Instant> = st.buckets.values().map(|b| b.last).collect();
126                stamps.sort_unstable();
127                let cutoff = stamps[(stamps.len() / 10).min(stamps.len() - 1)];
128                st.buckets.retain(|_, b| b.last > cutoff);
129            }
130        }
131        st.buckets
132            .entry(client)
133            .or_insert_with(|| TokenBucket::new(cfg, now))
134            .try_take(cfg, now)
135    }
136
137    pub fn len(&self) -> usize {
138        self.state.lock().expect("completion buckets").buckets.len()
139    }
140
141    pub fn is_empty(&self) -> bool {
142        self.len() == 0
143    }
144}
145
146pub fn http_client_key(peer: std::net::IpAddr, credential: Option<&str>) -> u64 {
147    let mut h = std::collections::hash_map::DefaultHasher::new();
148    peer.hash(&mut h);
149    credential.hash(&mut h);
150    h.finish()
151}
152
153pub fn retry_after_secs(wait: Duration) -> u64 {
154    wait.as_secs_f64().ceil().max(1.0) as u64
155}
156
157pub fn retry_after_ms(wait: Duration) -> u64 {
158    let ms = wait.as_nanos().div_ceil(1_000_000);
159    u64::try_from(ms).unwrap_or(u64::MAX).max(1)
160}
161
162#[cfg(test)]
163mod tests {
164    use super::*;
165
166    const CFG: BucketConfig = BucketConfig {
167        rate_per_sec: 2.0,
168        burst: 3,
169    };
170
171    #[test]
172    fn burst_then_refuse_then_refill() {
173        let t0 = Instant::now();
174        let mut b = TokenBucket::new(CFG, t0);
175        for _ in 0..3 {
176            assert!(b.try_take(CFG, t0).is_ok());
177        }
178        let wait = b.try_take(CFG, t0).unwrap_err();
179        assert!((wait.as_secs_f64() - 0.5).abs() < 1e-6, "{wait:?}");
180        let t1 = t0 + Duration::from_millis(500);
181        assert!(b.try_take(CFG, t1).is_ok());
182        assert!(b.try_take(CFG, t1).is_err());
183        let t2 = t1 + Duration::from_secs(60);
184        for _ in 0..3 {
185            assert!(b.try_take(CFG, t2).is_ok());
186        }
187        assert!(b.try_take(CFG, t2).is_err());
188    }
189
190    #[test]
191    fn the_advertised_wait_is_enough() {
192        for rate in [3.0, 7.0, 9.7, 13.0] {
193            let cfg = BucketConfig { rate_per_sec: rate, burst: 1 };
194            let t0 = Instant::now();
195            let mut b = TokenBucket::new(cfg, t0);
196            b.try_take(cfg, t0).unwrap();
197            let wait = b.try_take(cfg, t0).unwrap_err();
198            let t1 = t0 + Duration::from_millis(retry_after_ms(wait));
199            assert!(b.try_take(cfg, t1).is_ok(), "rate {rate}: {wait:?}");
200        }
201    }
202
203    #[test]
204    fn a_denormal_rate_saturates_instead_of_panicking() {
205        let cfg = BucketConfig { rate_per_sec: 1e-320, burst: 1 };
206        let t0 = Instant::now();
207        let mut b = TokenBucket::new(cfg, t0);
208        b.try_take(cfg, t0).unwrap();
209        let wait = b.try_take(cfg, t0).unwrap_err();
210        assert_eq!(wait, MAX_WAIT);
211        assert_eq!(retry_after_secs(wait), 24 * 3600);
212    }
213
214    #[test]
215    fn off_never_refuses() {
216        let t0 = Instant::now();
217        let mut b = TokenBucket::new(BucketConfig::OFF, t0);
218        for _ in 0..1000 {
219            assert!(b.try_take(BucketConfig::OFF, t0).is_ok());
220        }
221        let m = ClientBuckets::new();
222        assert!(m.try_take(1, BucketConfig::OFF).is_ok());
223        assert!(m.is_empty(), "a disabled bucket tracks nobody");
224    }
225
226    #[test]
227    fn clients_are_independent() {
228        let m = ClientBuckets::new();
229        let ip: std::net::IpAddr = "127.0.0.1".parse().unwrap();
230        let a = http_client_key(ip, Some("key-a"));
231        let b = http_client_key(ip, Some("key-b"));
232        assert_ne!(a, b, "two credentials on one host are two clients");
233        assert_eq!(a, http_client_key(ip, Some("key-a")), "stable across requests");
234        for _ in 0..3 {
235            assert!(m.try_take(a, CFG).is_ok());
236        }
237        assert!(m.try_take(a, CFG).is_err());
238        assert!(m.try_take(b, CFG).is_ok(), "b is not charged for a's storm");
239    }
240
241    #[test]
242    fn validation() {
243        assert!(BucketConfig::default().validate().is_ok());
244        assert!(BucketConfig::OFF.validate().is_ok());
245        assert!(BucketConfig { rate_per_sec: 1.0, burst: 0 }.validate().is_err());
246        assert!(BucketConfig { rate_per_sec: f64::NAN, burst: 1 }.validate().is_err());
247        assert!(BucketConfig { rate_per_sec: -1.0, burst: 1 }.validate().is_err());
248        assert!(BucketConfig { rate_per_sec: 1e-320, burst: 1 }.validate().is_err());
249        assert!(BucketConfig { rate_per_sec: 0.001, burst: 1 }.validate().is_ok());
250        assert_eq!(retry_after_secs(Duration::from_millis(1)), 1);
251        assert_eq!(retry_after_secs(Duration::from_millis(1500)), 2);
252        assert_eq!(retry_after_secs(Duration::from_nanos(1_000_000_001)), 2);
253        assert_eq!(retry_after_ms(Duration::from_nanos(99_000_001)), 100);
254        assert_eq!(retry_after_ms(Duration::from_millis(99)), 99);
255        assert_eq!(retry_after_ms(Duration::ZERO), 1);
256    }
257}
superfluid_daemon/completion_bucket.rs

superfluid_daemon/
completion_bucket.rs
