Skip to main content

lexe_connect/
mailbox.rs

1//! The mailbox address and an in-memory mailbox store.
2
3use std::{
4    collections::HashMap,
5    sync::Mutex,
6    time::{Duration, Instant},
7};
8
9use lexe_byte_array::ByteArray;
10
11#[derive(Copy, Clone, Debug, Eq, PartialEq)]
12pub enum PutError {
13    /// The address holds a different blob.
14    Occupied,
15    /// The blob exceeds the store's max blob length.
16    TooLarge,
17    /// The store is at capacity with unexpired blobs.
18    Full,
19}
20
21/// Where a response blob is deposited in a mailbox:
22/// `SHA256(DOMAIN_SEPARATOR || one_time_secret)`, i.e. the hash of the
23/// request's HPKE `info`, rendered as 64 lowercase hex characters.
24/// Only the holders of the connection string can compute it.
25//
26// Spec: `mailbox_url` delivery.
27#[derive(Copy, Clone, Eq, Hash, PartialEq)]
28pub struct Address([u8; 32]);
29
30lexe_byte_array::impl_byte_array!(Address, 32);
31lexe_byte_array::impl_fromstr_fromhex!(Address, 32);
32lexe_byte_array::impl_debug_display_as_hex!(Address);
33lexe_serde::impl_serde_hexstr_or_bytes!(Address);
34
35/// An in-memory mailbox: the first blob written to an address wins, and
36/// blobs expire after a TTL. Callers pass `now` so expiration is testable.
37//
38// Spec: `mailbox_url` delivery.
39pub struct MailboxStore {
40    blobs: Mutex<HashMap<Address, StoredBlob>>,
41    max_blob_len: usize,
42    max_blobs: usize,
43    ttl: Duration,
44}
45
46struct StoredBlob {
47    blob: Vec<u8>,
48    expires_at: Instant,
49}
50
51impl MailboxStore {
52    /// Lexe's public mailbox TTL.
53    //
54    // Spec: `mailbox_url` delivery.
55    pub const DEFAULT_TTL: Duration = Duration::from_secs(5 * 60);
56    /// A sealed response echoing max-length `account` and `metadata` around
57    /// a real Lexe credential is ~4.3 KiB; see `largest_response_fits`.
58    pub const DEFAULT_MAX_BLOB_LEN: usize = 8192;
59    /// Bounds memory at roughly `DEFAULT_MAX_BLOBS * DEFAULT_MAX_BLOB_LEN`,
60    /// ~4 MiB. Each blob lives for `DEFAULT_TTL`, so this also caps
61    /// sustained throughput at `DEFAULT_MAX_BLOBS / DEFAULT_TTL`, ~1.7
62    /// responses per second.
63    pub const DEFAULT_MAX_BLOBS: usize = 512;
64
65    /// A store with the `DEFAULT_*` limits.
66    pub fn new() -> Self {
67        Self {
68            blobs: Mutex::new(HashMap::new()),
69            max_blob_len: Self::DEFAULT_MAX_BLOB_LEN,
70            max_blobs: Self::DEFAULT_MAX_BLOBS,
71            ttl: Self::DEFAULT_TTL,
72        }
73    }
74
75    /// The largest blob [`put`](Self::put) accepts.
76    pub fn with_max_blob_len(mut self, max_blob_len: usize) -> Self {
77        self.max_blob_len = max_blob_len;
78        self
79    }
80
81    /// The most blobs held at once; bounds memory at roughly
82    /// `max_blobs * max_blob_len`.
83    pub fn with_max_blobs(mut self, max_blobs: usize) -> Self {
84        self.max_blobs = max_blobs;
85        self
86    }
87
88    /// How long a blob is held after it is [`put`](Self::put).
89    pub fn with_ttl(mut self, ttl: Duration) -> Self {
90        self.ttl = ttl;
91        self
92    }
93
94    /// The blob at `address`, if one is present and unexpired. Blobs
95    /// survive reads until they expire, so REQUESTERs can retry.
96    pub fn get(&self, address: &Address, now: Instant) -> Option<Vec<u8>> {
97        let mut locked_blobs = self.blobs.lock().unwrap();
98        let stored = locked_blobs.get(address)?;
99
100        if stored.expires_at <= now {
101            locked_blobs.remove(address);
102            return None;
103        }
104
105        Some(stored.blob.clone())
106    }
107
108    pub fn put(
109        &self,
110        address: Address,
111        blob: Vec<u8>,
112        now: Instant,
113    ) -> Result<(), PutError> {
114        if blob.len() > self.max_blob_len {
115            return Err(PutError::TooLarge);
116        }
117
118        let mut locked_blobs = self.blobs.lock().unwrap();
119        if locked_blobs.len() >= self.max_blobs {
120            locked_blobs.retain(|_, stored| stored.expires_at > now);
121            if locked_blobs.len() >= self.max_blobs {
122                return Err(PutError::Full);
123            }
124        }
125
126        // A retry of the same delivery succeeds without extending the TTL.
127        if let Some(stored) = locked_blobs.get(&address)
128            && stored.expires_at > now
129        {
130            return if stored.blob == blob {
131                Ok(())
132            } else {
133                Err(PutError::Occupied)
134            };
135        }
136
137        locked_blobs.insert(
138            address,
139            StoredBlob {
140                blob,
141                expires_at: now + self.ttl,
142            },
143        );
144
145        Ok(())
146    }
147}
148
149#[cfg(test)]
150mod test {
151    use lexe_byte_array::ByteArray;
152
153    use super::*;
154
155    #[test]
156    fn first_write_wins_retries_then_expires() {
157        let store = MailboxStore::new().with_ttl(Duration::from_secs(10));
158        let a = Address::from_array([1; 32]);
159        let t0 = Instant::now();
160
161        assert_eq!(store.get(&a, t0), None);
162        assert_eq!(store.put(a, vec![1], t0), Ok(()));
163        assert_eq!(store.put(a, vec![2], t0), Err(PutError::Occupied));
164        assert_eq!(store.get(&a, t0), Some(vec![1]));
165        let t_retry = t0 + Duration::from_secs(9);
166        assert_eq!(store.put(a, vec![1], t_retry), Ok(()));
167        assert_eq!(store.get(&a, t_retry), Some(vec![1]));
168
169        // The retry did not extend the TTL.
170        let t1 = t0 + Duration::from_secs(10);
171        assert_eq!(store.get(&a, t1), None);
172        assert_eq!(store.put(a, vec![2], t1), Ok(()));
173        assert_eq!(store.get(&a, t1), Some(vec![2]));
174
175        let too_large = vec![0; MailboxStore::DEFAULT_MAX_BLOB_LEN + 1];
176        let b = Address::from_array([2; 32]);
177        assert_eq!(store.put(b, too_large, t1), Err(PutError::TooLarge));
178    }
179
180    #[test]
181    fn full_store_sweeps_expired() {
182        const MAX_BLOBS: usize = 4;
183        let store = MailboxStore::new()
184            .with_ttl(Duration::from_secs(1))
185            .with_max_blobs(MAX_BLOBS);
186        let t0 = Instant::now();
187        for i in 0..MAX_BLOBS {
188            let mut bytes = [0u8; 32];
189            bytes[..8].copy_from_slice(&(i as u64).to_le_bytes());
190            store.put(Address::from_array(bytes), vec![], t0).unwrap();
191        }
192        let a = Address::from_array([0xff; 32]);
193        assert_eq!(store.put(a, vec![], t0), Err(PutError::Full));
194        assert_eq!(store.put(a, vec![], t0 + Duration::from_secs(1)), Ok(()));
195    }
196}