didbot_serve/proxy/
limits.rs

1//! What one account may spend of the proxy: a rate, a number of calls in
2//! flight, and a cache of where its targets live.
3//!
4//! The target of a proxied call is any https origin a DID document names,
5//! and the call leaves from this server's address. Without a per-account
6//! rate, one account can point this server at a third party as fast as it
7//! likes, and the blocklist that follows hurts every account. Without caps
8//! on calls in flight, accounts calling a service that never answers can
9//! hold every one of [`crate::MAX_CONCURRENT_REQUESTS`] for
10//! [`super::PROXY_TIMEOUT`] each. One creator may hold thousands of
11//! accounts, so the per-account cap alone does not bound that; the
12//! deployment-wide cap does, and the per-tree cap keeps one creator's
13//! accounts from taking all of it.
14
15use std::collections::HashMap;
16use std::num::{NonZeroU32, NonZeroUsize};
17use std::sync::{Arc, Mutex};
18use std::time::{Duration, Instant};
19
20use didbot_http::PrivateAddresses;
21use governor::clock::Clock;
22use governor::{DefaultKeyedRateLimiter, Quota};
23use url::Url;
24
25/// The longest target reference the cache keeps an entry for: the longest
26/// DID there is, with room for a service id. A longer one is resolved on
27/// every call, so a header's length does not decide how much the cache
28/// holds.
29const MAX_CACHED_TARGET: usize = didbot_identity::MAX_DID_LENGTH + 64;
30
31/// How often buckets that have filled back up are dropped. A full bucket
32/// says what an absent one says.
33const SWEEP_EVERY: Duration = Duration::from_secs(60);
34
35/// The bounds on proxied calls, and where an `app.bsky.*` call with no
36/// `atproto-proxy` goes; see [`ProxyLimits`].
37#[derive(Debug, Clone, PartialEq, Eq)]
38pub struct ProxySettings {
39    /// How many proxied calls one account's budget refills by each minute.
40    ///
41    /// **Reasonable range: 1 to 6000.** Sixty, the default, is one a second
42    /// sustained: a hundred times the few calls a minute an agent makes,
43    /// and the most one account can make a third party see from this
44    /// server's address.
45    pub calls_per_minute: u32,
46    /// How many proxied calls one account may make at once from a full
47    /// budget.
48    ///
49    /// **Reasonable range: 1 to 1000.** Sixty, the default, covers a client
50    /// that opens a view by fanning out a few dozen calls.
51    pub burst_calls: u32,
52    /// How many of one account's proxied calls may be open at once, from
53    /// the request to the last byte of the answer.
54    ///
55    /// **Reasonable range: 1 to 64.** Eight, the default, is more than an
56    /// agent opens at once, and keeps one account to a sixteenth of
57    /// `in_flight_total`.
58    pub in_flight_per_account: usize,
59    /// How many proxied calls may be open at once across one account tree:
60    /// an account this server's operator admitted, and every account
61    /// created beneath it.
62    ///
63    /// **Reasonable range: 1 to 512.** 64, the default, is half of
64    /// `in_flight_total`, so one creator's accounts leave half for every
65    /// other tree. A per-tree cap at or above `in_flight_total` has no
66    /// effect.
67    pub in_flight_per_tree: usize,
68    /// How many proxied calls may be open at once across every account.
69    ///
70    /// **Reasonable range: 1 to 512.** 128, the default, is a quarter of
71    /// [`crate::MAX_CONCURRENT_REQUESTS`], so however many accounts call a
72    /// stalled service, three quarters of the request permits stay free
73    /// for everything else. A per-account cap above it has no effect.
74    pub in_flight_total: usize,
75    /// How long the endpoint a target's DID document names is kept.
76    ///
77    /// **Reasonable range: 10 to 3600 seconds.** Five minutes, the default,
78    /// fetches a document once for every few hundred calls an agent makes
79    /// to it, and a service that moves is reached at its new address within
80    /// five minutes.
81    pub target_cache_ttl: Duration,
82    /// How many targets' endpoints are kept, the least recently used
83    /// dropped first.
84    ///
85    /// **Reasonable range: 16 to 100_000.** An entry is a few kilobytes at
86    /// most. 1024, the default, holds every service a deployment's agents
87    /// call in common use, in a few megabytes.
88    pub target_cache_entries: usize,
89    /// The service an `app.bsky.*` call with no `atproto-proxy` header is
90    /// sent to, as that header would name it: `did#service`, such as
91    /// `did:web:api.bsky.app#bsky_appview`. The reference PDS sends these
92    /// calls to its configured app view the same way.
93    ///
94    /// `None`, the default, sends nowhere: such a call is
95    /// `MethodNotImplemented`.
96    pub app_view: Option<String>,
97}
98
99impl Default for ProxySettings {
100    fn default() -> Self {
101        Self {
102            calls_per_minute: 60,
103            burst_calls: 60,
104            in_flight_per_account: 8,
105            in_flight_per_tree: 64,
106            in_flight_total: 128,
107            target_cache_ttl: Duration::from_secs(300),
108            target_cache_entries: 1024,
109            app_view: None,
110        }
111    }
112}
113
114impl ProxySettings {
115    /// Refuses a setting outside its documented range, or an `app_view`
116    /// that is not a DID and a service id, naming its `[proxy]` key.
117    pub fn check(&self) -> Result<(), String> {
118        if let Some(app_view) = &self.app_view {
119            let target = super::Target::parse(app_view)
120                .map_err(|error| format!("[proxy] app_view {app_view:?}: {}", error.message()))?;
121            didbot_identity::document_url(target.did)
122                .map_err(|error| format!("[proxy] app_view {app_view:?}: {error}"))?;
123        }
124        let ttl = self.target_cache_ttl.as_secs();
125        let out_of_range = [
126            (
127                "calls_per_minute",
128                u64::from(self.calls_per_minute),
129                1,
130                6000,
131            ),
132            ("burst_calls", u64::from(self.burst_calls), 1, 1000),
133            (
134                "in_flight_per_account",
135                self.in_flight_per_account as u64,
136                1,
137                64,
138            ),
139            ("in_flight_per_tree", self.in_flight_per_tree as u64, 1, 512),
140            ("in_flight_total", self.in_flight_total as u64, 1, 512),
141            ("target_cache_secs", ttl, 10, 3600),
142            (
143                "target_cache_entries",
144                self.target_cache_entries as u64,
145                16,
146                100_000,
147            ),
148        ]
149        .into_iter()
150        .find(|(_, value, low, high)| !(low..=high).contains(&value));
151        match out_of_range {
152            Some((key, value, low, high)) => Err(format!(
153                "[proxy] {key} must be {low} to {high}, not {value}"
154            )),
155            None => Ok(()),
156        }
157    }
158}
159
160/// Why a proxied call was refused before it left this server.
161#[derive(Debug, Clone, Copy, PartialEq, Eq)]
162pub(crate) enum Refusal {
163    /// The account has spent its budget; the next call is allowed after
164    /// this long.
165    Rate(Duration),
166    /// The account already has this many calls open.
167    InFlight(usize),
168    /// The account's tree already has this many calls open.
169    Tree(usize),
170    /// Every account together already has this many calls open.
171    Full(usize),
172}
173
174/// The proxied calls open now.
175#[derive(Default)]
176struct Open {
177    per_account: HashMap<String, usize>,
178    per_tree: HashMap<String, usize>,
179    total: usize,
180}
181
182/// One account's call held open, released when dropped.
183pub(crate) struct InFlight {
184    open: Arc<Mutex<Open>>,
185    did: String,
186    tree: String,
187}
188
189impl Drop for InFlight {
190    fn drop(&mut self) {
191        let mut open = lock(&self.open);
192        open.total -= 1;
193        let Open {
194            per_account,
195            per_tree,
196            ..
197        } = &mut *open;
198        release(per_account, &self.did);
199        release(per_tree, &self.tree);
200    }
201}
202
203/// Gives back one of `key`'s calls, dropping its entry at none.
204fn release(counts: &mut HashMap<String, usize>, key: &str) {
205    if let Some(count) = counts.get_mut(key) {
206        *count -= 1;
207        if *count == 0 {
208            counts.remove(key);
209        }
210    }
211}
212
213/// An endpoint a target's document named, and under what it was found.
214struct Endpoint {
215    url: Url,
216    /// The setting the document was fetched under. An entry found while
217    /// private addresses were allowed is not used once they are refused.
218    private: PrivateAddresses,
219    fetched: Instant,
220}
221
222/// The per-account bounds on proxied calls, and the cache of targets'
223/// endpoints, for one deployment.
224pub struct ProxyLimits {
225    settings: ProxySettings,
226    rate: DefaultKeyedRateLimiter<String>,
227    last_swept: Mutex<Instant>,
228    open: Arc<Mutex<Open>>,
229    targets: Mutex<lru::LruCache<String, Endpoint>>,
230}
231
232impl ProxyLimits {
233    /// Bounds set by `settings`. A zero anywhere counts as one.
234    pub fn new(settings: ProxySettings) -> Self {
235        let rate = NonZeroU32::new(settings.calls_per_minute).unwrap_or(NonZeroU32::MIN);
236        let burst = NonZeroU32::new(settings.burst_calls).unwrap_or(NonZeroU32::MIN);
237        let entries = NonZeroUsize::new(settings.target_cache_entries).unwrap_or(NonZeroUsize::MIN);
238        Self {
239            settings,
240            rate: DefaultKeyedRateLimiter::keyed(Quota::per_minute(rate).allow_burst(burst)),
241            last_swept: Mutex::new(Instant::now()),
242            open: Arc::new(Mutex::new(Open::default())),
243            targets: Mutex::new(lru::LruCache::new(entries)),
244        }
245    }
246
247    /// The settings these bounds were built from.
248    pub fn settings(&self) -> &ProxySettings {
249        &self.settings
250    }
251
252    /// Spends one of `did`'s calls and holds it open, or says why not.
253    /// `tree` names the account tree `did` is in; see
254    /// [`ProxySettings::in_flight_per_tree`].
255    ///
256    /// The rate is checked first, so a call refused for being in flight
257    /// has still spent its place in the budget: an account retrying
258    /// against its own stalled calls slows itself down.
259    pub(crate) fn admit(&self, did: &str, tree: &str) -> Result<InFlight, Refusal> {
260        self.sweep();
261        let key = did.to_owned();
262        if let Err(not_until) = self.rate.check_key(&key) {
263            return Err(Refusal::Rate(
264                not_until.wait_time_from(self.rate.clock().now()),
265            ));
266        }
267        let mut open = lock(&self.open);
268        let account = open.per_account.get(&key).copied().unwrap_or(0);
269        if account >= self.settings.in_flight_per_account {
270            return Err(Refusal::InFlight(account));
271        }
272        let held = open.per_tree.get(tree).copied().unwrap_or(0);
273        if held >= self.settings.in_flight_per_tree {
274            return Err(Refusal::Tree(held));
275        }
276        if open.total >= self.settings.in_flight_total {
277            return Err(Refusal::Full(open.total));
278        }
279        open.total += 1;
280        *open.per_account.entry(key.clone()).or_insert(0) += 1;
281        *open.per_tree.entry(tree.to_owned()).or_insert(0) += 1;
282        Ok(InFlight {
283            open: self.open.clone(),
284            did: key,
285            tree: tree.to_owned(),
286        })
287    }
288
289    /// The endpoint cached for `target`, if one is fresh and was found
290    /// under `private`.
291    pub(crate) fn cached_target(&self, target: &str, private: PrivateAddresses) -> Option<Url> {
292        let mut targets = lock(&self.targets);
293        let entry = targets.get(target)?;
294        if entry.fetched.elapsed() >= self.settings.target_cache_ttl || entry.private != private {
295            targets.pop(target);
296            return None;
297        }
298        Some(entry.url.clone())
299    }
300
301    /// Keeps the endpoint `target`'s document named, found under `private`.
302    pub(crate) fn cache_target(&self, target: &str, private: PrivateAddresses, url: &Url) {
303        if target.len() > MAX_CACHED_TARGET {
304            return;
305        }
306        lock(&self.targets).put(
307            target.to_owned(),
308            Endpoint {
309                url: url.clone(),
310                private,
311                fetched: Instant::now(),
312            },
313        );
314    }
315
316    /// Drops full buckets, at most once per [`SWEEP_EVERY`], so the rate
317    /// map holds the accounts calling now rather than every account that
318    /// ever has.
319    fn sweep(&self) {
320        let mut last = lock(&self.last_swept);
321        if last.elapsed() < SWEEP_EVERY {
322            return;
323        }
324        *last = Instant::now();
325        self.rate.retain_recent();
326        self.rate.shrink_to_fit();
327    }
328}
329
330impl Default for ProxyLimits {
331    fn default() -> Self {
332        Self::new(ProxySettings::default())
333    }
334}
335
336fn lock<T>(mutex: &Mutex<T>) -> std::sync::MutexGuard<'_, T> {
337    mutex
338        .lock()
339        .unwrap_or_else(|poisoned| poisoned.into_inner())
340}
341
342#[cfg(test)]
343mod tests {
344    use super::*;
345
346    fn endpoint() -> Url {
347        Url::parse("https://service.example").unwrap()
348    }
349
350    /// An entry past its TTL is fetched again, and one found while private
351    /// addresses were allowed is not used once they are refused.
352    #[test]
353    fn a_cached_target_expires_and_does_not_cross_the_private_setting() {
354        let limits = ProxyLimits::new(ProxySettings {
355            target_cache_ttl: Duration::from_millis(50),
356            ..ProxySettings::default()
357        });
358        let target = "did:web:service.example#svc";
359        limits.cache_target(target, PrivateAddresses::Allowed, &endpoint());
360        assert_eq!(
361            limits.cached_target(target, PrivateAddresses::Allowed),
362            Some(endpoint())
363        );
364        assert_eq!(
365            limits.cached_target(target, PrivateAddresses::Refused),
366            None
367        );
368
369        limits.cache_target(target, PrivateAddresses::Refused, &endpoint());
370        std::thread::sleep(Duration::from_millis(60));
371        assert_eq!(
372            limits.cached_target(target, PrivateAddresses::Refused),
373            None
374        );
375    }
376
377    /// Releasing a call frees its place, and an account with none open
378    /// holds no entry.
379    #[test]
380    fn a_released_call_frees_its_place() {
381        let limits = ProxyLimits::new(ProxySettings {
382            in_flight_per_account: 1,
383            ..ProxySettings::default()
384        });
385        let held = limits
386            .admit("did:web:one.example", "did:web:one.example")
387            .unwrap();
388        assert_eq!(
389            limits
390                .admit("did:web:one.example", "did:web:one.example")
391                .err(),
392            Some(Refusal::InFlight(1))
393        );
394        drop(held);
395        assert!(lock(&limits.open).per_account.is_empty());
396        assert!(lock(&limits.open).per_tree.is_empty());
397        assert!(limits
398            .admit("did:web:one.example", "did:web:one.example")
399            .is_ok());
400    }
401
402    /// One tree's accounts together hold no more than `in_flight_per_tree`,
403    /// and another tree's accounts still get theirs.
404    #[test]
405    fn one_trees_accounts_hold_no_more_than_the_trees_share() {
406        let limits = ProxyLimits::new(ProxySettings {
407            in_flight_per_account: 2,
408            in_flight_per_tree: 3,
409            ..ProxySettings::default()
410        });
411        let creator = "did:web:creator.example";
412        let held: Vec<_> = [
413            "did:web:a.example",
414            "did:web:b.example",
415            "did:web:b.example",
416        ]
417        .into_iter()
418        .map(|did| limits.admit(did, creator).unwrap())
419        .collect();
420        assert_eq!(
421            limits.admit("did:web:c.example", creator).err(),
422            Some(Refusal::Tree(3))
423        );
424        assert!(limits
425            .admit("did:web:other.example", "did:web:other.example")
426            .is_ok());
427        drop(held);
428        assert!(limits.admit("did:web:c.example", creator).is_ok());
429    }
430
431    /// Many accounts together hold no more than `in_flight_total`, which is
432    /// what stops many trees together taking every request permit.
433    #[test]
434    fn every_account_together_holds_no_more_than_the_total() {
435        let limits = ProxyLimits::new(ProxySettings {
436            in_flight_per_account: 2,
437            in_flight_total: 3,
438            ..ProxySettings::default()
439        });
440        let first: Vec<_> = (0..2)
441            .map(|_| {
442                limits
443                    .admit("did:web:one.example", "did:web:one.example")
444                    .unwrap()
445            })
446            .collect();
447        let second = limits
448            .admit("did:web:two.example", "did:web:two.example")
449            .unwrap();
450        assert_eq!(
451            limits
452                .admit("did:web:three.example", "did:web:three.example")
453                .err(),
454            Some(Refusal::Full(3))
455        );
456        drop(first);
457        drop(second);
458        assert_eq!(lock(&limits.open).total, 0);
459        assert!(limits
460            .admit("did:web:three.example", "did:web:three.example")
461            .is_ok());
462    }
463}
didbot_serve/proxy/limits.rs

didbot_serve/proxy/
limits.rs
