zeyn/
security.rs

1//! Rate limits and IP rules.
2
3use crate::ctx::BoxFuture;
4use std::collections::HashMap;
5use std::net::IpAddr;
6use std::sync::Mutex;
7use std::time::{Duration, Instant};
8
9/// What a rate limit counts by.
10#[derive(Clone, Copy, Debug, PartialEq, Eq)]
11pub enum RateKey {
12    /// The client's IP address.
13    Ip,
14    /// The signed in user; anonymous requests fall back to the IP address.
15    User,
16}
17
18/// At most `requests` per `window` for each key. Bursts up to `requests` are allowed, then the
19/// allowance refills evenly over the window (token bucket).
20///
21/// ```ignore
22/// App::new("Shop")
23///     .rate_limit(RateLimit::per_ip(300, Duration::from_secs(60)))
24///     .write_limit(RateLimit::per_user(30, Duration::from_secs(60)))
25/// ```
26#[derive(Clone, Copy, Debug, PartialEq)]
27pub struct RateLimit {
28    pub requests: u32,
29    pub window: Duration,
30    pub key: RateKey,
31}
32
33impl RateLimit {
34    pub fn per_ip(requests: u32, window: Duration) -> Self {
35        RateLimit {
36            requests: requests.max(1),
37            window,
38            key: RateKey::Ip,
39        }
40    }
41
42    pub fn per_user(requests: u32, window: Duration) -> Self {
43        RateLimit {
44            requests: requests.max(1),
45            window,
46            key: RateKey::User,
47        }
48    }
49}
50
51/// Counters behind a rate limit. The built in one keeps them in memory, which is right for a single
52/// server; several servers behind a load balancer need a shared store such as `RedisStore`
53/// (Cargo feature `redis`).
54pub trait RateStore: Send + Sync + 'static {
55    /// Takes one request from the key's allowance. `Err` carries how long to wait. A store that
56    /// cannot reach its backend lets the request through: a cache outage must not stop the app.
57    fn take(&self, key: &str, limit: &RateLimit) -> BoxFuture<Result<(), Duration>>;
58
59    /// Takes `amount` units at once, such as the tokens of a model answer. The allowance may go
60    /// below zero, so a large answer is paid in full and the next call waits. `Err` when less than
61    /// one unit is left (also with `amount` 0, which only checks). Stores that only count requests
62    /// may keep the default, which takes one unit.
63    fn charge(&self, key: &str, limit: &RateLimit, amount: u32) -> BoxFuture<Result<(), Duration>> {
64        let _ = amount;
65        self.take(key, limit)
66    }
67}
68
69/// In memory token buckets with a cap on the number of keys, so a flood of addresses cannot
70/// exhaust memory.
71pub struct MemoryRateStore {
72    buckets: Mutex<HashMap<String, (f64, Instant)>>,
73    max_keys: usize,
74}
75
76impl Default for MemoryRateStore {
77    fn default() -> Self {
78        MemoryRateStore {
79            buckets: Mutex::new(HashMap::new()),
80            max_keys: 100_000,
81        }
82    }
83}
84
85impl MemoryRateStore {
86    pub fn with_max_keys(max_keys: usize) -> Self {
87        MemoryRateStore {
88            max_keys: max_keys.max(1),
89            ..Self::default()
90        }
91    }
92
93    pub fn len(&self) -> usize {
94        self.buckets
95            .lock()
96            .map(|b| b.len())
97            .unwrap_or(0)
98    }
99
100    pub fn is_empty(&self) -> bool {
101        self.len() == 0
102    }
103}
104
105impl RateStore for MemoryRateStore {
106    fn take(&self, key: &str, limit: &RateLimit) -> BoxFuture<Result<(), Duration>> {
107        let taken = MemoryRateStore::take(self, key, limit);
108        Box::pin(async move { taken })
109    }
110
111    fn charge(&self, key: &str, limit: &RateLimit, amount: u32) -> BoxFuture<Result<(), Duration>> {
112        let charged = MemoryRateStore::charge(self, key, limit, amount);
113        Box::pin(async move { charged })
114    }
115}
116
117impl MemoryRateStore {
118    /// Takes one request from the key's allowance, right away.
119    pub fn take(&self, key: &str, limit: &RateLimit) -> Result<(), Duration> {
120        self.bucket(key, limit, |tokens| {
121            if *tokens >= 1.0 {
122                *tokens -= 1.0;
123                true
124            } else {
125                false
126            }
127        })
128    }
129
130    /// Takes `amount` units right away; see [`RateStore::charge`].
131    pub fn charge(&self, key: &str, limit: &RateLimit, amount: u32) -> Result<(), Duration> {
132        self.bucket(key, limit, |tokens| {
133            *tokens -= f64::from(amount);
134            *tokens >= 1.0
135        })
136    }
137
138    /// Refills the key's bucket and lets `use_tokens` take from it; `Err` with the wait until one
139    /// unit is back when it says no.
140    fn bucket(
141        &self,
142        key: &str,
143        limit: &RateLimit,
144        use_tokens: impl FnOnce(&mut f64) -> bool,
145    ) -> Result<(), Duration> {
146        let now = Instant::now();
147        let capacity = f64::from(limit.requests);
148        let per_sec = capacity / limit.window.as_secs_f64().max(0.001);
149        let Ok(mut buckets) = self.buckets.lock() else {
150            return Ok(());
151        };
152        if buckets.len() >= self.max_keys && !buckets.contains_key(key) {
153            buckets.retain(|_, (tokens, at)| {
154                let refilled = *tokens + now.duration_since(*at).as_secs_f64() * per_sec;
155                refilled < capacity
156            });
157            if buckets.len() >= self.max_keys {
158                let oldest = buckets
159                    .iter()
160                    .min_by_key(|(_, (_, at))| *at)
161                    .map(|(k, _)| k.clone());
162                if let Some(k) = oldest {
163                    buckets.remove(&k);
164                }
165            }
166        }
167        let (tokens, at) = buckets
168            .entry(key.to_string())
169            .or_insert((capacity, now));
170        *tokens = (*tokens + now.duration_since(*at).as_secs_f64() * per_sec).min(capacity);
171        *at = now;
172        if use_tokens(tokens) {
173            Ok(())
174        } else {
175            Err(Duration::from_secs_f64((1.0 - *tokens).max(0.0) / per_sec))
176        }
177    }
178}
179
180/// An address range such as `10.0.0.0/8`, `192.168.1.10` or `2001:db8::/32`.
181#[derive(Clone, Copy, Debug, PartialEq, Eq)]
182pub struct IpRange {
183    base: IpAddr,
184    prefix: u8,
185}
186
187impl IpRange {
188    pub fn parse(s: &str) -> Option<IpRange> {
189        let (ip, prefix) = match s.trim().split_once('/') {
190            Some((ip, p)) => (ip.parse::<IpAddr>().ok()?, p.parse::<u8>().ok()?),
191            None => {
192                let ip: IpAddr = s.trim().parse().ok()?;
193                (ip, if ip.is_ipv4() { 32 } else { 128 })
194            }
195        };
196        let max = if ip.is_ipv4() { 32 } else { 128 };
197        (prefix <= max).then_some(IpRange { base: ip, prefix })
198    }
199
200    pub fn contains(&self, ip: IpAddr) -> bool {
201        let bits = |ip: IpAddr| -> (u128, u8) {
202            match ip {
203                IpAddr::V4(v) => (u128::from(u32::from(v)), 32),
204                IpAddr::V6(v) => match v.to_ipv4_mapped() {
205                    Some(v4) => (u128::from(u32::from(v4)), 32),
206                    None => (u128::from(v), 128),
207                },
208            }
209        };
210        let (a, wa) = bits(self.base);
211        let (b, wb) = bits(ip);
212        if wa != wb {
213            return false;
214        }
215        if self.prefix == 0 {
216            return true;
217        }
218        let shift = u32::from(wa - self.prefix);
219        (a >> shift) == (b >> shift)
220    }
221}
222
223/// Parses a list of ranges; invalid entries are reported so a typo never opens a door silently.
224pub fn parse_ranges<I, S>(ranges: I) -> Result<Vec<IpRange>, String>
225where
226    I: IntoIterator<Item = S>,
227    S: AsRef<str>,
228{
229    ranges
230        .into_iter()
231        .map(|r| {
232            IpRange::parse(r.as_ref())
233                .ok_or_else(|| format!("`{}` is not an IP address or range", r.as_ref()))
234        })
235        .collect()
236}
237
238/// The client's address: the connection's peer, or, when the peer is a trusted proxy, the last
239/// address in `X-Forwarded-For` that is not a trusted proxy.
240pub fn client_ip(
241    peer: Option<IpAddr>,
242    forwarded_for: Option<&str>,
243    trusted: &[IpRange],
244) -> Option<IpAddr> {
245    let peer = peer?;
246    if !trusted.iter().any(|r| r.contains(peer)) {
247        return Some(peer);
248    }
249    let Some(chain) = forwarded_for else {
250        return Some(peer);
251    };
252    chain
253        .rsplit(',')
254        .filter_map(|p| p.trim().parse::<IpAddr>().ok())
255        .find(|ip| !trusted.iter().any(|r| r.contains(*ip)))
256        .or(Some(peer))
257}
zeyn/security.rs

zeyn/
security.rs
