ojak/deliverer/
keys.rs

1//! Keeping the keys deliveries are signed with, once loaded.
2
3use super::SenderKeys;
4use crate::queue::QueueError;
5use crate::sig::SenderKey;
6use std::collections::HashMap;
7use std::sync::Mutex;
8use std::time::{Duration, Instant};
9
10/// [`SenderKeys`] that keeps each key `inner` gives for `ttl` after loading
11/// it, so that it is not loaded again for every delivery it signs.
12///
13/// A post fans out to thousands of inboxes signed with one key; loading it
14/// for each — from a database, decrypting it, parsing its PEM — costs as
15/// many queries and as much work, and with many deliveries in flight can
16/// hold every connection an application has. A key replaced at its source is
17/// signed with once its `ttl` has run out. Only keys are kept: a sender with
18/// none, and an error, are asked about again next time. Whether a sender is
19/// [gone](SenderKeys::gone) is always asked of `inner`.
20pub struct CachingSenderKeys<K> {
21    inner: K,
22    ttl: Duration,
23    loaded: Mutex<HashMap<String, (Instant, SenderKey)>>,
24}
25
26impl<K> CachingSenderKeys<K> {
27    /// `inner`'s keys, each kept for `ttl`.
28    pub fn new(inner: K, ttl: Duration) -> Self {
29        Self {
30            inner,
31            ttl,
32            loaded: Mutex::new(HashMap::new()),
33        }
34    }
35
36    /// Where the keys come from.
37    pub fn inner(&self) -> &K {
38        &self.inner
39    }
40
41    fn cached(&self, sender: &str) -> Option<SenderKey> {
42        let loaded = self.loaded.lock().expect("sender keys");
43        let (at, key) = loaded.get(sender)?;
44        (at.elapsed() < self.ttl).then(|| key.clone())
45    }
46
47    fn keep(&self, sender: &str, key: &SenderKey) {
48        let mut loaded = self.loaded.lock().expect("sender keys");
49        // Only the keys signing now are worth keeping.
50        loaded.retain(|_, (at, _)| at.elapsed() < self.ttl);
51        loaded.insert(sender.to_owned(), (Instant::now(), key.clone()));
52    }
53}
54
55impl<K: SenderKeys> SenderKeys for CachingSenderKeys<K> {
56    async fn key(&self, sender: &str) -> Result<Option<SenderKey>, QueueError> {
57        if let Some(key) = self.cached(sender) {
58            return Ok(Some(key));
59        }
60        let key = self.inner.key(sender).await?;
61        if let Some(key) = &key {
62            self.keep(sender, key);
63        }
64        Ok(key)
65    }
66
67    async fn gone(&self, sender: &str) -> bool {
68        self.inner.gone(sender).await
69    }
70}
71
72#[cfg(test)]
73mod tests {
74    use super::*;
75    use std::sync::Arc;
76    use std::sync::atomic::{AtomicUsize, Ordering};
77
78    struct Counting {
79        key: SenderKey,
80        loads: Arc<AtomicUsize>,
81    }
82
83    impl SenderKeys for Counting {
84        async fn key(&self, sender: &str) -> Result<Option<SenderKey>, QueueError> {
85            self.loads.fetch_add(1, Ordering::SeqCst);
86            Ok((sender == "alice").then(|| self.key.clone()))
87        }
88
89        async fn gone(&self, sender: &str) -> bool {
90            sender == "gone"
91        }
92    }
93
94    fn key() -> SenderKey {
95        let pem = crate::sig::signature::generate_rsa_keypair(&mut rand_core::UnwrapErr(
96            getrandom::SysRng,
97        ))
98        .unwrap()
99        .0;
100        SenderKey {
101            key_id: "https://a.example/users/alice#main-key".into(),
102            private_key: Arc::new(crate::sig::PrivateKey::from_pem(&pem).unwrap()),
103        }
104    }
105
106    #[tokio::test]
107    async fn a_key_is_loaded_once_until_it_is_stale() {
108        let loads = Arc::new(AtomicUsize::new(0));
109        let keys = CachingSenderKeys::new(
110            Counting {
111                key: key(),
112                loads: loads.clone(),
113            },
114            Duration::from_secs(300),
115        );
116        assert!(keys.key("alice").await.unwrap().is_some());
117        assert!(keys.key("alice").await.unwrap().is_some());
118        assert_eq!(loads.load(Ordering::SeqCst), 1);
119        // A sender with no key is asked about each time.
120        assert!(keys.key("bob").await.unwrap().is_none());
121        assert!(keys.key("bob").await.unwrap().is_none());
122        assert_eq!(loads.load(Ordering::SeqCst), 3);
123        assert!(keys.gone("gone").await);
124
125        let stale = CachingSenderKeys::new(
126            Counting {
127                key: key(),
128                loads: loads.clone(),
129            },
130            Duration::ZERO,
131        );
132        stale.key("alice").await.unwrap();
133        stale.key("alice").await.unwrap();
134        assert_eq!(loads.load(Ordering::SeqCst), 5);
135    }
136}
ojak/deliverer/keys.rs

ojak/deliverer/
keys.rs
