Skip to main content

aauth_person_server/
store.rs

1use std::collections::hash_map::Entry;
2use std::collections::HashMap;
3use std::sync::Mutex;
4
5use async_trait::async_trait;
6
7use crate::types::{Approval, ApprovalStatus, LogEntry, Mission, Person};
8
9/// Whether [`Store::spend`] may take a mission past its budget.
10#[derive(Clone, Copy, Debug, PartialEq, Eq)]
11pub enum Spend {
12    /// Only if it fits what is left: a request issued without asking anyone.
13    WithinBudget,
14    /// Whatever is left: a request the approver approved.
15    Approved,
16}
17
18/// A store that could not answer, such as a database that is down. The
19/// Person Server answers the request with a `503`.
20#[derive(Debug)]
21pub struct StoreError(Box<dyn std::error::Error + Send + Sync>);
22
23impl StoreError {
24    /// Wrap the underlying error.
25    pub fn new(error: impl Into<Box<dyn std::error::Error + Send + Sync>>) -> Self {
26        Self(error.into())
27    }
28}
29
30impl std::fmt::Display for StoreError {
31    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
32        write!(f, "store unavailable: {}", self.0)
33    }
34}
35
36impl std::error::Error for StoreError {
37    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
38        Some(self.0.as_ref())
39    }
40}
41
42/// What a [`Store`] answers.
43pub type StoreResult<T> = Result<T, StoreError>;
44
45/// Where a Person Server keeps its state.
46///
47/// Every method takes the request's tenant: the AAuth `tenant` the agent
48/// named, or the server's default tenant when it named none (see
49/// [`crate::Config::default_tenant`]). A store keeps each tenant's records
50/// apart: a record is found only under the tenant it was saved with. A
51/// tenant the store cannot scope, such as one in a format it does not
52/// serve, finds nothing.
53///
54/// Three methods must be atomic, because concurrent requests race on them:
55/// [`Store::spend`] checks and spends in one step, so two bookings cannot
56/// both fit in what is left; [`Store::transition_approval`] changes an
57/// approval only from the status it was read in, so a decision, a
58/// withdrawal or a collection happens once; and [`Store::insert_mission`]
59/// inserts only when no mission has that hash, so it never replaces one.
60#[async_trait]
61pub trait Store: Send + Sync {
62    /// The person an agent key is bound to.
63    async fn binding(&self, tenant: &str, jkt: &str) -> StoreResult<Option<Person>>;
64    /// Bind a key to a person unless it is bound already; returns whoever it is bound to.
65    async fn bind(&self, tenant: &str, jkt: &str, person: &Person) -> StoreResult<Person>;
66    /// The mission `s256` names.
67    async fn mission(&self, tenant: &str, s256: &str) -> StoreResult<Option<Mission>>;
68    /// Save a newly approved mission, unless one with its hash exists.
69    /// Returns whether it did: an existing mission is never replaced.
70    async fn insert_mission(&self, tenant: &str, mission: &Mission) -> StoreResult<bool>;
71    /// Append `entry` to a live mission's log. Returns whether it was live.
72    async fn log_mission(
73        &self,
74        tenant: &str,
75        s256: &str,
76        entry: LogEntry,
77        now: i64,
78    ) -> StoreResult<bool>;
79    /// Mark a mission completed at `at`, unless it already was. Returns the
80    /// mission, or `None` when there is no such mission.
81    async fn end_mission(&self, tenant: &str, s256: &str, at: i64) -> StoreResult<Option<Mission>>;
82    /// Add `amount` to a live mission's spend, under `rule`. Returns the
83    /// mission as spent, or `None` when it was not live or `amount` did not fit.
84    async fn spend(
85        &self,
86        tenant: &str,
87        s256: &str,
88        amount: u64,
89        rule: Spend,
90        now: i64,
91    ) -> StoreResult<Option<Mission>>;
92    /// The tenant's missions, for the management routes.
93    async fn missions(&self, tenant: &str) -> StoreResult<Vec<Mission>>;
94    /// The approval `id` names.
95    async fn approval(&self, tenant: &str, id: &str) -> StoreResult<Option<Approval>>;
96    /// Save a new approval.
97    async fn insert_approval(&self, tenant: &str, approval: &Approval) -> StoreResult<()>;
98    /// Replace approval `next.id` with `next`, only if its status is still
99    /// `expected`. Returns whether it did.
100    async fn transition_approval(
101        &self,
102        tenant: &str,
103        expected: ApprovalStatus,
104        next: &Approval,
105    ) -> StoreResult<bool>;
106    /// The tenant's approvals, for the management routes.
107    async fn approvals(&self, tenant: &str) -> StoreResult<Vec<Approval>>;
108}
109
110/// A record's key: its tenant and its own identifier.
111type Key = (String, String);
112
113fn key(tenant: &str, id: &str) -> Key {
114    (tenant.to_string(), id.to_string())
115}
116
117#[derive(Default)]
118struct State {
119    bindings: HashMap<Key, Person>,
120    missions: HashMap<Key, Mission>,
121    approvals: HashMap<Key, Approval>,
122}
123
124/// Approvals are kept this long past their expiry, so a late poll still
125/// reads `expired` rather than `not_found`.
126const APPROVAL_RETENTION_SECS: i64 = 24 * 60 * 60;
127
128/// Everything in memory: for development, tests and single-process demos.
129/// Expired approvals are dropped a day after they expire; nothing else is
130/// ever dropped, and nothing survives a restart.
131#[derive(Default)]
132pub struct MemoryStore(Mutex<State>);
133
134impl MemoryStore {
135    /// An empty store.
136    pub fn new() -> Self {
137        Self::default()
138    }
139
140    fn with<T>(&self, f: impl FnOnce(&mut State) -> T) -> StoreResult<T> {
141        let mut state = self
142            .0
143            .lock()
144            .unwrap_or_else(|poisoned| poisoned.into_inner());
145        Ok(f(&mut state))
146    }
147}
148
149#[async_trait]
150impl Store for MemoryStore {
151    async fn binding(&self, tenant: &str, jkt: &str) -> StoreResult<Option<Person>> {
152        self.with(|s| s.bindings.get(&key(tenant, jkt)).cloned())
153    }
154
155    async fn bind(&self, tenant: &str, jkt: &str, person: &Person) -> StoreResult<Person> {
156        self.with(|s| {
157            s.bindings
158                .entry(key(tenant, jkt))
159                .or_insert_with(|| person.clone())
160                .clone()
161        })
162    }
163
164    async fn mission(&self, tenant: &str, s256: &str) -> StoreResult<Option<Mission>> {
165        self.with(|s| s.missions.get(&key(tenant, s256)).cloned())
166    }
167
168    async fn insert_mission(&self, tenant: &str, mission: &Mission) -> StoreResult<bool> {
169        self.with(|s| match s.missions.entry(key(tenant, &mission.s256)) {
170            Entry::Occupied(_) => false,
171            Entry::Vacant(slot) => {
172                slot.insert(mission.clone());
173                true
174            }
175        })
176    }
177
178    async fn log_mission(
179        &self,
180        tenant: &str,
181        s256: &str,
182        entry: LogEntry,
183        now: i64,
184    ) -> StoreResult<bool> {
185        self.with(|s| match s.missions.get_mut(&key(tenant, s256)) {
186            Some(mission) if mission.is_live(now) => {
187                mission.log.push(entry);
188                true
189            }
190            _ => false,
191        })
192    }
193
194    async fn end_mission(&self, tenant: &str, s256: &str, at: i64) -> StoreResult<Option<Mission>> {
195        self.with(|s| {
196            let mission = s.missions.get_mut(&key(tenant, s256))?;
197            mission.completed_at.get_or_insert(at);
198            Some(mission.clone())
199        })
200    }
201
202    async fn spend(
203        &self,
204        tenant: &str,
205        s256: &str,
206        amount: u64,
207        rule: Spend,
208        now: i64,
209    ) -> StoreResult<Option<Mission>> {
210        self.with(|s| {
211            let mission = s.missions.get_mut(&key(tenant, s256))?;
212            if !mission.is_live(now) {
213                return None;
214            }
215            if rule == Spend::WithinBudget {
216                let left = mission.budget.as_ref()?.amount.checked_sub(mission.spent)?;
217                if left < amount {
218                    return None;
219                }
220            }
221            mission.spent = mission.spent.saturating_add(amount);
222            Some(mission.clone())
223        })
224    }
225
226    async fn missions(&self, tenant: &str) -> StoreResult<Vec<Mission>> {
227        self.with(|s| {
228            s.missions
229                .iter()
230                .filter(|((t, _), _)| t == tenant)
231                .map(|(_, m)| m.clone())
232                .collect()
233        })
234    }
235
236    async fn approval(&self, tenant: &str, id: &str) -> StoreResult<Option<Approval>> {
237        self.with(|s| s.approvals.get(&key(tenant, id)).cloned())
238    }
239
240    async fn insert_approval(&self, tenant: &str, approval: &Approval) -> StoreResult<()> {
241        self.with(|s| {
242            let now = approval.created_at;
243            s.approvals
244                .retain(|_, a| a.expires_at + APPROVAL_RETENTION_SECS > now);
245            s.approvals
246                .insert(key(tenant, &approval.id), approval.clone());
247        })
248    }
249
250    async fn transition_approval(
251        &self,
252        tenant: &str,
253        expected: ApprovalStatus,
254        next: &Approval,
255    ) -> StoreResult<bool> {
256        self.with(|s| match s.approvals.get_mut(&key(tenant, &next.id)) {
257            Some(current) if current.status == expected => {
258                *current = next.clone();
259                true
260            }
261            _ => false,
262        })
263    }
264
265    async fn approvals(&self, tenant: &str) -> StoreResult<Vec<Approval>> {
266        self.with(|s| {
267            s.approvals
268                .iter()
269                .filter(|((t, _), _)| t == tenant)
270                .map(|(_, a)| a.clone())
271                .collect()
272        })
273    }
274}
275
276#[cfg(test)]
277mod tests {
278    use super::*;
279
280    fn mission(person: &str) -> Mission {
281        Mission {
282            s256: "h".into(),
283            mission: "{}".into(),
284            agent_id: "aauth:travel@agents.example".into(),
285            person: Person::new(person),
286            resources: Vec::new(),
287            budget: None,
288            spent: 0,
289            log: Vec::new(),
290            expires_at: i64::MAX,
291            completed_at: None,
292            tenant: "default".into(),
293        }
294    }
295
296    #[tokio::test]
297    async fn a_mission_is_never_replaced() {
298        let store = MemoryStore::default();
299        assert!(store
300            .insert_mission("default", &mission("sam"))
301            .await
302            .unwrap());
303        store.end_mission("default", "h", 10).await.unwrap();
304        assert!(!store
305            .insert_mission("default", &mission("dana"))
306            .await
307            .unwrap());
308        let kept = store.mission("default", "h").await.unwrap().unwrap();
309        assert_eq!(kept.person.id, "sam");
310        assert_eq!(kept.completed_at, Some(10));
311    }
312
313    #[tokio::test]
314    async fn each_tenant_sees_only_its_own() {
315        let store = MemoryStore::default();
316        assert!(store.insert_mission("acme", &mission("sam")).await.unwrap());
317        assert!(store
318            .insert_mission("globex", &mission("hank"))
319            .await
320            .unwrap());
321        let at = |t| store.mission(t, "h");
322        assert_eq!(at("acme").await.unwrap().unwrap().person.id, "sam");
323        assert_eq!(at("globex").await.unwrap().unwrap().person.id, "hank");
324        assert!(at("default").await.unwrap().is_none());
325        assert_eq!(store.missions("acme").await.unwrap().len(), 1);
326        store.bind("acme", "k", &Person::new("sam")).await.unwrap();
327        assert!(store.binding("globex", "k").await.unwrap().is_none());
328    }
329}