furnace/
access.rs

1//! Who may call the server and how often: API keys (`--api-key`) and a per-client rate limit
2//! (`--rate-limit`), shared by the HTTP and gRPC transports.
3//!
4//! The rate limit is a GCRA (generic cell rate algorithm), i.e. a token bucket kept as a single
5//! "theoretical arrival time" per client: `per_second` requests a second on average, and bursts
6//! of up to `burst` at once. Clients are the API key they present or, without keys, their IP
7//! address. Only a hash of either is kept.
8
9use std::collections::HashMap;
10use std::hash::{BuildHasher, RandomState};
11use std::net::IpAddr;
12use std::sync::Mutex;
13use std::time::{Duration, Instant};
14
15/// `--rate-limit` and `--rate-limit-burst`
16#[derive(Debug, Clone, Copy, PartialEq)]
17pub struct ClientRateLimit {
18    /// Requests a second per client, on average
19    pub per_second: f64,
20    /// Requests a client may send at once
21    pub burst: u32,
22}
23
24/// Why a request is turned away
25#[derive(Debug, PartialEq)]
26pub(crate) enum Rejection {
27    /// API keys are required and it presented none or an unknown one
28    Unauthorized,
29    /// Over the rate limit; may retry after this long
30    RateLimited(Duration),
31}
32
33/// API keys and rate limiting for one server
34pub(crate) struct Access {
35    api_keys: Vec<Vec<u8>>,
36    limiter: Option<RateLimiter>,
37}
38
39impl Access {
40    /// Admit requests presenting one of `api_keys` (any request when empty), each client
41    /// within `rate_limit` if one is set
42    pub fn new(api_keys: &[String], rate_limit: Option<ClientRateLimit>) -> Self {
43        Self {
44            api_keys: api_keys.iter().map(|k| k.as_bytes().to_vec()).collect(),
45            limiter: rate_limit.map(RateLimiter::new),
46        }
47    }
48
49    /// Whether requests need an API key
50    pub fn requires_key(&self) -> bool {
51        !self.api_keys.is_empty()
52    }
53
54    /// Whether every request is admitted: no API keys and no rate limit
55    pub fn is_open(&self) -> bool {
56        self.api_keys.is_empty() && self.limiter.is_none()
57    }
58
59    /// Admit a request presenting `key` (from `Authorization: Bearer` or `X-API-Key`) from
60    /// `ip`, or say why not. `check_key` is false for routes with their own authentication.
61    pub fn admit(
62        &self,
63        key: Option<&str>,
64        ip: Option<IpAddr>,
65        check_key: bool,
66    ) -> Result<(), Rejection> {
67        let known_key = key.filter(|key| self.is_known(key));
68        if check_key && self.requires_key() && known_key.is_none() {
69            return Err(Rejection::Unauthorized);
70        }
71        if let Some(limiter) = &self.limiter {
72            let client = match (known_key, ip) {
73                (Some(key), _) => Client::Key(key),
74                (None, Some(ip)) => Client::Ip(ip),
75                (None, None) => Client::Unknown,
76            };
77            limiter.check(client, Instant::now())?;
78        }
79        Ok(())
80    }
81
82    /// Compares with every key, each in constant time, so timing reveals neither the key nor
83    /// which one matched
84    fn is_known(&self, key: &str) -> bool {
85        self.api_keys.iter().fold(false, |found, known| {
86            found | constant_time_eq(known, key.as_bytes())
87        })
88    }
89}
90
91/// `a == b` in time depending only on the lengths
92pub(crate) fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
93    a.len() == b.len() && a.iter().zip(b).fold(0u8, |acc, (x, y)| acc | (x ^ y)) == 0
94}
95
96/// The API key in `Authorization: Bearer <key>` or `X-API-Key: <key>`
97pub(crate) fn presented_key<'a>(
98    authorization: Option<&'a str>,
99    x_api_key: Option<&'a str>,
100) -> Option<&'a str> {
101    authorization
102        .and_then(|value| value.strip_prefix("Bearer "))
103        .or(x_api_key)
104        .map(str::trim)
105}
106
107#[derive(Hash)]
108enum Client<'a> {
109    Key(&'a str),
110    Ip(IpAddr),
111    Unknown,
112}
113
114const SHARDS: usize = 16;
115/// A shard with more clients than this drops the ones whose bucket has refilled
116const SHARD_PRUNE_AT: usize = 4096;
117
118struct RateLimiter {
119    /// Time one request uses up
120    interval: u64,
121    /// How far ahead of now a client's arrival time may run: (burst - 1) intervals
122    tolerance: u64,
123    hasher: RandomState,
124    epoch: Instant,
125    /// Per client hash, its theoretical arrival time in nanoseconds since `epoch`
126    shards: [Mutex<HashMap<u64, u64>>; SHARDS],
127}
128
129impl RateLimiter {
130    fn new(limit: ClientRateLimit) -> Self {
131        let interval = (1e9 / limit.per_second).round().max(1.0) as u64;
132        Self {
133            interval,
134            tolerance: interval.saturating_mul(u64::from(limit.burst.max(1) - 1)),
135            hasher: RandomState::new(),
136            epoch: Instant::now(),
137            shards: std::array::from_fn(|_| Mutex::new(HashMap::new())),
138        }
139    }
140
141    fn check(&self, client: Client<'_>, now: Instant) -> Result<(), Rejection> {
142        let id = self.hasher.hash_one(&client);
143        let now =
144            u64::try_from(now.saturating_duration_since(self.epoch).as_nanos()).unwrap_or(u64::MAX);
145        let mut shard = self.shards[id as usize % SHARDS]
146            .lock()
147            .unwrap_or_else(|e| e.into_inner());
148        let arrival = shard.get(&id).copied().unwrap_or(now).max(now);
149        if arrival - now > self.tolerance {
150            return Err(Rejection::RateLimited(Duration::from_nanos(
151                arrival - now - self.tolerance,
152            )));
153        }
154        if shard.len() >= SHARD_PRUNE_AT && !shard.contains_key(&id) {
155            // A client whose arrival time has passed is the same as one never seen
156            shard.retain(|_, &mut arrival| arrival > now);
157        }
158        shard.insert(id, arrival + self.interval);
159        Ok(())
160    }
161}
162
163#[cfg(test)]
164mod tests {
165    use super::*;
166
167    fn ip(last: u8) -> Option<IpAddr> {
168        Some(IpAddr::from([10, 0, 0, last]))
169    }
170
171    #[test]
172    fn keys_are_required_when_configured() {
173        let open = Access::new(&[], None);
174        assert_eq!(open.admit(None, ip(1), true), Ok(()));
175        let keys = ["first-key".to_string(), "second-key".to_string()];
176        let access = Access::new(&keys, None);
177        assert_eq!(access.admit(Some("second-key"), ip(1), true), Ok(()));
178        assert_eq!(
179            access.admit(None, ip(1), true),
180            Err(Rejection::Unauthorized)
181        );
182        assert_eq!(
183            access.admit(Some("second-ke"), ip(1), true),
184            Err(Rejection::Unauthorized)
185        );
186        // Routes with their own authentication
187        assert_eq!(access.admit(None, ip(1), false), Ok(()));
188    }
189
190    #[test]
191    fn presented_key_from_either_header() {
192        assert_eq!(presented_key(Some("Bearer abc"), None), Some("abc"));
193        assert_eq!(presented_key(None, Some("abc")), Some("abc"));
194        assert_eq!(presented_key(Some("Basic abc"), None), None);
195        assert_eq!(presented_key(Some("Bearer a"), Some("b")), Some("a"));
196    }
197
198    #[test]
199    fn bursts_then_the_average_rate() {
200        let limiter = RateLimiter::new(ClientRateLimit {
201            per_second: 10.0,
202            burst: 3,
203        });
204        let t0 = Instant::now();
205        let check = |at_ms: u64| {
206            limiter.check(
207                Client::Ip(IpAddr::from([10, 0, 0, 1])),
208                t0 + Duration::from_millis(at_ms),
209            )
210        };
211        // Three at once, then one every 100 ms
212        for _ in 0..3 {
213            assert_eq!(check(0), Ok(()));
214        }
215        let Err(Rejection::RateLimited(retry)) = check(0) else {
216            panic!("the fourth request should be limited");
217        };
218        assert!(retry > Duration::from_millis(99) && retry <= Duration::from_millis(100));
219        assert_eq!(check(100), Ok(()));
220        assert!(check(150).is_err());
221        assert_eq!(check(200), Ok(()));
222        // An idle client gets its whole burst back
223        for _ in 0..3 {
224            assert_eq!(check(10_000), Ok(()));
225        }
226        assert!(check(10_000).is_err());
227    }
228
229    #[test]
230    fn clients_are_limited_separately() {
231        let keys = ["key-a".to_string(), "key-b".to_string()];
232        let access = Access::new(
233            &keys,
234            Some(ClientRateLimit {
235                per_second: 1.0,
236                burst: 1,
237            }),
238        );
239        assert_eq!(access.admit(Some("key-a"), ip(1), true), Ok(()));
240        assert!(access.admit(Some("key-a"), ip(2), true).is_err());
241        // Another key from the same address has its own budget
242        assert_eq!(access.admit(Some("key-b"), ip(1), true), Ok(()));
243
244        let open = Access::new(
245            &[],
246            Some(ClientRateLimit {
247                per_second: 1.0,
248                burst: 1,
249            }),
250        );
251        assert_eq!(open.admit(None, ip(1), true), Ok(()));
252        assert!(open.admit(None, ip(1), true).is_err());
253        assert_eq!(open.admit(None, ip(2), true), Ok(()));
254        // An unknown key does not open a new budget
255        assert!(open.admit(Some("made-up"), ip(1), true).is_err());
256    }
257
258    #[test]
259    fn many_clients_stay_bounded() {
260        let limiter = RateLimiter::new(ClientRateLimit {
261            per_second: 1000.0,
262            burst: 1,
263        });
264        let t0 = Instant::now();
265        for i in 0..20_000u32 {
266            let ip = IpAddr::from(i.to_be_bytes());
267            // 10 ms apart: each earlier client's bucket (1 ms) has refilled
268            let _ = limiter.check(
269                Client::Ip(ip),
270                t0 + Duration::from_millis(u64::from(i) * 10),
271            );
272        }
273        let clients: usize = limiter.shards.iter().map(|s| s.lock().unwrap().len()).sum();
274        assert!(clients <= SHARDS * SHARD_PRUNE_AT, "{clients}");
275    }
276}
furnace/access.rs

furnace/
access.rs
