1use 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 Occupied,
15 TooLarge,
17 Full,
19}
20
21#[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
35pub 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 pub const DEFAULT_TTL: Duration = Duration::from_secs(5 * 60);
56 pub const DEFAULT_MAX_BLOB_LEN: usize = 8192;
59 pub const DEFAULT_MAX_BLOBS: usize = 512;
64
65 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 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 pub fn with_max_blobs(mut self, max_blobs: usize) -> Self {
84 self.max_blobs = max_blobs;
85 self
86 }
87
88 pub fn with_ttl(mut self, ttl: Duration) -> Self {
90 self.ttl = ttl;
91 self
92 }
93
94 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 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 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}