didbot_policy_count/
lib.rs

1//! The counting engine, behind the same [`Evaluator`] contract as
2//! `didbot-policy-regex` and `didbot-policy-cedar`.
3//!
4//! A document is `bot.did.policy#count`: statements, each naming the
5//! writes it counts, whether it counts each account on its own or the
6//! deployment as a whole, the window, the limit, and what a refusal says.
7//! The engine holds the counts; the numbers are the operator's, or the
8//! floor's, and they live in the document.
9//!
10//! ```json
11//! {
12//!   "collections": ["app.bsky.graph.follow"],
13//!   "naming": "anotherAccount",
14//!   "scope": "account",
15//!   "window": "hour",
16//!   "limit": 50,
17//!   "reason": "this account has reached this server's limit of 50 follows in an hour."
18//! }
19//! ```
20//!
21//! # The count judged is the count before this write
22//!
23//! The write path calls `judge` and then `observe` for every attempt,
24//! admitted or refused, so the limit is the last write admitted: with a
25//! limit of fifty, the fiftieth follow inside the window lands and the
26//! fifty-first is refused. A write another policy refused still counts,
27//! which is the point of counting attempts rather than commits.
28//!
29//! A statement counting the deployment as a whole counts only the writes
30//! the policy set allowed. An account retrying refused writes then spends
31//! its own count and nobody else's: once its own limit refuses it, its
32//! attempts stop reaching the deployment's count.
33//!
34//! # Windows are whole hours
35//!
36//! Counts are kept in a ring of hourly buckets. The hour window is the
37//! current bucket and the day window is the ring, so the hour advances in
38//! steps rather than sliding: a burst may spend the hour's budget in a
39//! second.
40//!
41//! # A count outlives the tree that judged it
42//!
43//! The tree is rebuilt on every poll tick, and this engine is constructed
44//! fresh with it. The counts are not: they live in one [`Counts`] per
45//! process, held under the statement as written, so a tick is not a way
46//! around a limit and an edit to a statement's limit keeps the count the
47//! old one was keeping. Two statements written differently count for
48//! themselves, whoever wrote them, so an operator's own counting policy
49//! can only ever refuse more.
50//!
51//! # Time is an input
52//!
53//! Nothing here calls a clock. The hour comes off
54//! [`didbot_policy::Universal::now`], which the write path supplies and
55//! the evaluation log records.
56
57#![forbid(unsafe_code)]
58
59mod counts;
60mod interaction;
61mod statement;
62
63use std::sync::{Arc, OnceLock, RwLock};
64
65use didbot_policy::{
66    CompileOutcome, CompiledId, EvalError, Evaluator, Lifecycle, Observed, Outcome, Staleness,
67    Subject,
68};
69use serde_json::Value;
70
71pub use counts::{hour_of, Counts, BUCKETS, HELD};
72pub use interaction::names_another_account;
73pub use statement::{Scope, Statement, Window};
74
75/// The counts this process holds, shared by every document it compiles.
76/// See the crate doc comment for why they are the process's and not the
77/// tree's.
78#[must_use]
79pub fn counts() -> &'static Arc<Counts> {
80    static COUNTS: OnceLock<Arc<Counts>> = OnceLock::new();
81    COUNTS.get_or_init(|| Arc::new(Counts::new()))
82}
83
84/// The engine: the counts it judges against, and every statement compiled
85/// into it.
86pub struct CountEvaluator {
87    counts: Arc<Counts>,
88    statements: RwLock<Vec<Statement>>,
89}
90
91impl Default for CountEvaluator {
92    fn default() -> Self {
93        Self::new()
94    }
95}
96
97impl CountEvaluator {
98    /// An engine counting into this process's [`counts()`].
99    #[must_use]
100    pub fn new() -> Self {
101        Self::over(Arc::clone(counts()))
102    }
103
104    /// An engine counting into `counts` instead, for a caller that holds
105    /// its own.
106    #[must_use]
107    pub fn over(counts: Arc<Counts>) -> Self {
108        Self {
109            counts,
110            statements: RwLock::new(Vec::new()),
111        }
112    }
113
114    fn statement(&self, compiled: CompiledId) -> Result<Statement, EvalError> {
115        let index = usize::try_from(compiled.0)
116            .map_err(|_| EvalError::Failed(format!("no compiled statement {}", compiled.0)))?;
117        self.statements
118            .read()
119            .unwrap_or_else(std::sync::PoisonError::into_inner)
120            .get(index)
121            .cloned()
122            .ok_or_else(|| EvalError::Failed(format!("no compiled statement {}", compiled.0)))
123    }
124
125    /// The account this write counts against for `statement`, or `None`
126    /// when the statement does not count it: another collection, or a
127    /// record that names nobody else when the statement asks for one that
128    /// does.
129    fn counted<'a>(statement: &Statement, subject: &'a Subject<'a>) -> Option<Option<&'a str>> {
130        let Subject::Write {
131            account,
132            attributes,
133            collection,
134            after,
135            ..
136        } = subject
137        else {
138            return None;
139        };
140        if !statement.collections.is_empty()
141            && !statement
142                .collections
143                .iter()
144                .any(|named| named == *collection)
145        {
146            return None;
147        }
148        if statement.names_another
149            && !after.is_some_and(|record| {
150                names_another_account(collection, record, account, attributes.handle)
151            })
152        {
153            return None;
154        }
155        Some(match statement.scope {
156            Scope::Account => Some(*account),
157            Scope::Server => None,
158        })
159    }
160}
161
162impl Evaluator for CountEvaluator {
163    fn compile(&self, predicate: &Value) -> CompileOutcome {
164        match Statement::parse(predicate) {
165            Ok(statement) => {
166                let mut statements = self
167                    .statements
168                    .write()
169                    .unwrap_or_else(std::sync::PoisonError::into_inner);
170                let id = CompiledId(statements.len() as u64);
171                statements.push(statement);
172                CompileOutcome::Compiled(id)
173            }
174            Err(reason) => CompileOutcome::Rejected(reason),
175        }
176    }
177
178    fn evaluate(&self, subject: &Subject<'_>, compiled: CompiledId) -> Result<Outcome, EvalError> {
179        let statement = self.statement(compiled)?;
180        let Some(account) = Self::counted(&statement, subject) else {
181            return Ok(Outcome::Allow);
182        };
183        let at = hour_of(subject.universal().now);
184        let held = self
185            .counts
186            .read(&statement.counter, account, statement.window, at);
187        if held >= u64::from(statement.limit) {
188            return Ok(Outcome::Reject {
189                reason: statement.reason.clone(),
190            });
191        }
192        Ok(Outcome::Allow)
193    }
194
195    /// Every attempt on the declared surface, admitted or refused, for a
196    /// statement counting each account. A statement counting the
197    /// deployment counts only the attempts the policy set allowed.
198    fn observe(&self, event: &Observed<'_>) {
199        let Ok(statement) = self.statement(event.compiled) else {
200            return;
201        };
202        let Some(account) = Self::counted(&statement, event.subject) else {
203            return;
204        };
205        if account.is_none() && !matches!(event.outcome, Outcome::Allow) {
206            return;
207        }
208        let at = hour_of(event.subject.universal().now);
209        self.counts.add(&statement.counter, account, at);
210    }
211
212    /// A deleted account's counts go with it. A freeze does not reset a
213    /// count and an unfreeze does not hand back a spent hour: this engine
214    /// refuses and never freezes, so there is no freeze here to re-trip.
215    fn lifecycle(&self, event: &Lifecycle<'_>) {
216        if let Lifecycle::Deleted { account } = event {
217            self.counts.forget(account);
218        }
219    }
220
221    /// The count is exact and in memory, and every write on the server is
222    /// routed to a statement counting every collection, so this engine
223    /// blocks rather than answer from a stale view.
224    fn staleness(&self) -> Staleness {
225        Staleness::Linearized
226    }
227}
228
229#[cfg(test)]
230mod tests {
231    use super::*;
232    use didbot_policy::{
233        AccountKind, Applicability, DeploymentAttributes, Diff, KindSet, PolicyDeclaration,
234        PolicyId, PolicyTree, Predicate, Revision, Scope as AnyScope, SubjectAttributes,
235        SubjectKind, Universal, WriteAction,
236    };
237    use serde_json::json;
238    use time::OffsetDateTime;
239
240    const AGENT: &str = "did:web:agent.pds.example";
241    const POST: &str = "app.bsky.feed.post";
242
243    fn statement(scope: &str, window: &str, limit: u32) -> Value {
244        json!({
245            "scope": scope,
246            "window": window,
247            "limit": limit,
248            "reason": format!("this account has reached its {window}ly limit of {limit}."),
249        })
250    }
251
252    fn at(hours: i64) -> OffsetDateTime {
253        OffsetDateTime::from_unix_timestamp(hours * 3_600).expect("an instant")
254    }
255
256    fn write<'a>(account: &'a str, record: &'a Value, now: OffsetDateTime) -> Subject<'a> {
257        Subject::Write {
258            account,
259            attributes: SubjectAttributes {
260                kind: SubjectKind::Known(AccountKind::Agent),
261                handle: None,
262                depth: None,
263            },
264            universal: Universal {
265                deployment: DeploymentAttributes {
266                    pds_did: "did:web:pds.example",
267                    pds_hostname: "pds.example",
268                    operator_did: None,
269                },
270                now,
271            },
272            client_id: None,
273            collection: POST,
274            action: WriteAction::Create,
275            diff: Diff { changes: &[] },
276            before: None,
277            after: Some(record),
278            lookups: &[],
279        }
280    }
281
282    /// A tree position to label an observation with. This engine never
283    /// reads it -- it tells its statements apart by `CompiledId` -- but
284    /// `Observed` always carries one.
285    fn a_policy() -> PolicyId {
286        let applicability = Applicability::Write {
287            kinds: KindSet::ANY,
288            collections: AnyScope::Any,
289            actions: didbot_policy::ActionSet::ALL,
290            path_prefixes: AnyScope::Any,
291        };
292        PolicyTree::builder().add_policy(PolicyDeclaration {
293            observes: applicability.clone(),
294            applicability,
295            staleness: Staleness::Linearized,
296            predicate: Predicate::Unusable {
297                reason: "a label, never asked to decide anything".to_owned(),
298            },
299            revision: Revision(0),
300        })
301    }
302
303    /// Judges the subject and then observes it, the way the write path
304    /// does.
305    fn attempt(engine: &CountEvaluator, compiled: CompiledId, subject: &Subject<'_>) -> Outcome {
306        let outcome = engine.evaluate(subject, compiled).expect("judged");
307        engine.observe(&Observed {
308            subject,
309            policy: a_policy(),
310            compiled,
311            outcome: &outcome,
312        });
313        outcome
314    }
315
316    fn compiled(engine: &CountEvaluator, statement: &Value) -> CompiledId {
317        match engine.compile(statement) {
318            CompileOutcome::Compiled(id) => id,
319            other => panic!("expected a compiled statement, got {other:?}"),
320        }
321    }
322
323    /// The write past the limit is the one refused, and a refused attempt
324    /// still counts.
325    #[test]
326    fn the_write_past_the_limit_is_refused() {
327        let engine = CountEvaluator::over(Arc::new(Counts::new()));
328        let id = compiled(&engine, &statement("account", "hour", 3));
329        let record = json!({ "text": "hello" });
330        let now = at(10);
331        for n in 0..3 {
332            assert_eq!(
333                attempt(&engine, id, &write(AGENT, &record, now)),
334                Outcome::Allow,
335                "write {n}"
336            );
337        }
338        assert!(matches!(
339            attempt(&engine, id, &write(AGENT, &record, now)),
340            Outcome::Reject { reason } if reason.contains("hourly limit of 3")
341        ));
342        // The next hour starts from what this one spent, not from nothing.
343        assert_eq!(
344            attempt(&engine, id, &write(AGENT, &record, at(11))),
345            Outcome::Allow
346        );
347    }
348
349    /// A statement over the deployment counts every account together.
350    #[test]
351    fn a_server_statement_counts_every_account_together() {
352        let engine = CountEvaluator::over(Arc::new(Counts::new()));
353        let id = compiled(&engine, &statement("server", "hour", 2));
354        let record = json!({ "text": "hello" });
355        let now = at(10);
356        assert_eq!(
357            attempt(&engine, id, &write("did:web:one", &record, now)),
358            Outcome::Allow
359        );
360        assert_eq!(
361            attempt(&engine, id, &write("did:web:two", &record, now)),
362            Outcome::Allow
363        );
364        assert!(matches!(
365            attempt(&engine, id, &write("did:web:three", &record, now)),
366            Outcome::Reject { .. }
367        ));
368    }
369
370    /// An account retrying past its own limit spends its own count and not
371    /// the deployment's: the deployment's count holds only the writes the
372    /// policy set allowed, so another account can still write.
373    #[test]
374    fn a_refused_attempt_does_not_spend_the_deployments_count() {
375        let engine = CountEvaluator::over(Arc::new(Counts::new()));
376        let own = compiled(&engine, &statement("account", "hour", 3));
377        let shared = compiled(&engine, &statement("server", "hour", 5));
378        let record = json!({ "text": "hello" });
379        let now = at(10);
380        // As the write path does: every statement judges, and each observes
381        // the verdict of the whole set.
382        let attempt_both = |account: &str| {
383            let subject = write(account, &record, now);
384            let verdict = didbot_policy::worst(
385                [own, shared].map(|id| engine.evaluate(&subject, id).expect("judged")),
386            );
387            for compiled in [own, shared] {
388                engine.observe(&Observed {
389                    subject: &subject,
390                    policy: a_policy(),
391                    compiled,
392                    outcome: &verdict,
393                });
394            }
395            verdict
396        };
397        for n in 0..3 {
398            assert_eq!(attempt_both("did:web:busy"), Outcome::Allow, "write {n}");
399        }
400        for _ in 0..50 {
401            assert!(matches!(
402                attempt_both("did:web:busy"),
403                Outcome::Reject { .. }
404            ));
405        }
406        assert_eq!(attempt_both("did:web:quiet"), Outcome::Allow);
407        assert_eq!(attempt_both("did:web:quiet"), Outcome::Allow);
408        assert!(matches!(
409            attempt_both("did:web:quiet"),
410            Outcome::Reject { reason } if reason.contains("hourly limit of 5")
411        ));
412    }
413
414    /// The counts outlive the engine, because a poll tick builds a new
415    /// one and must not hand back a spent hour.
416    #[test]
417    fn a_rebuilt_engine_finds_the_count_the_last_one_left() {
418        let counts = Arc::new(Counts::new());
419        let record = json!({ "text": "hello" });
420        let now = at(10);
421        let first = CountEvaluator::over(Arc::clone(&counts));
422        let id = compiled(&first, &statement("account", "hour", 2));
423        assert_eq!(
424            attempt(&first, id, &write(AGENT, &record, now)),
425            Outcome::Allow
426        );
427        assert_eq!(
428            attempt(&first, id, &write(AGENT, &record, now)),
429            Outcome::Allow
430        );
431        let rebuilt = CountEvaluator::over(Arc::clone(&counts));
432        let id = compiled(&rebuilt, &statement("account", "hour", 2));
433        assert!(matches!(
434            attempt(&rebuilt, id, &write(AGENT, &record, now)),
435            Outcome::Reject { .. }
436        ));
437        // And a restart is a fresh window.
438        counts.forget_everything();
439        assert_eq!(
440            attempt(&rebuilt, id, &write(AGENT, &record, now)),
441            Outcome::Allow
442        );
443    }
444
445    /// `naming` counts only what reaches somebody else, and a collection
446    /// the statement does not name is not counted at all.
447    #[test]
448    fn a_statement_counts_only_what_it_names() {
449        let engine = CountEvaluator::over(Arc::new(Counts::new()));
450        let mut narrowed = statement("account", "hour", 1);
451        narrowed["collections"] = json!([POST]);
452        narrowed["naming"] = json!("anotherAccount");
453        let id = compiled(&engine, &narrowed);
454        let now = at(10);
455        let alone = json!({ "text": "hello" });
456        let reply = json!({
457            "text": "hello",
458            "reply": { "parent": { "uri": "at://did:plc:stranger/app.bsky.feed.post/3l" } },
459        });
460        for _ in 0..3 {
461            assert_eq!(
462                attempt(&engine, id, &write(AGENT, &alone, now)),
463                Outcome::Allow
464            );
465        }
466        assert_eq!(
467            attempt(&engine, id, &write(AGENT, &reply, now)),
468            Outcome::Allow
469        );
470        assert!(matches!(
471            attempt(&engine, id, &write(AGENT, &reply, now)),
472            Outcome::Reject { .. }
473        ));
474        assert_eq!(
475            attempt(&engine, id, &write(AGENT, &alone, now)),
476            Outcome::Allow,
477            "a record that names nobody else is not held to this statement"
478        );
479    }
480
481    /// An account's counts go when the account does.
482    #[test]
483    fn a_deleted_account_takes_its_counts_with_it() {
484        let counts = Arc::new(Counts::new());
485        let engine = CountEvaluator::over(Arc::clone(&counts));
486        let id = compiled(&engine, &statement("account", "hour", 1));
487        let record = json!({ "text": "hello" });
488        let now = at(10);
489        assert_eq!(
490            attempt(&engine, id, &write(AGENT, &record, now)),
491            Outcome::Allow
492        );
493        assert!(matches!(
494            attempt(&engine, id, &write(AGENT, &record, now)),
495            Outcome::Reject { .. }
496        ));
497        engine.lifecycle(&Lifecycle::Deleted { account: AGENT });
498        assert_eq!(
499            attempt(&engine, id, &write(AGENT, &record, now)),
500            Outcome::Allow
501        );
502    }
503
504    /// A statement this engine cannot read is a rejection at compile
505    /// time, which the loader turns into a scoped denial.
506    #[test]
507    fn a_statement_that_does_not_read_is_rejected_at_compile() {
508        let engine = CountEvaluator::over(Arc::new(Counts::new()));
509        assert!(matches!(
510            engine.compile(&json!({ "scope": "account" })),
511            CompileOutcome::Rejected(_)
512        ));
513    }
514}
didbot_policy_count/lib.rs

didbot_policy_count/
lib.rs
