aztec_pxe/stores/
sender_store.rs1use std::sync::Arc;
4
5use aztec_core::error::Error;
6use aztec_core::types::{AztecAddress, Fr};
7
8use super::kv::KvStore;
9
10pub struct SenderStore {
12 kv: Arc<dyn KvStore>,
13}
14
15impl SenderStore {
16 pub fn new(kv: Arc<dyn KvStore>) -> Self {
17 Self { kv }
18 }
19
20 pub async fn add(&self, sender: &AztecAddress) -> Result<(), Error> {
22 let key = sender_key(sender);
23 self.kv.put(&key, &[1]).await
25 }
26
27 pub async fn remove(&self, sender: &AztecAddress) -> Result<(), Error> {
29 let key = sender_key(sender);
30 self.kv.delete(&key).await
31 }
32
33 pub async fn contains(&self, sender: &AztecAddress) -> Result<bool, Error> {
35 let key = sender_key(sender);
36 Ok(self.kv.get(&key).await?.is_some())
37 }
38
39 pub async fn get_all(&self) -> Result<Vec<AztecAddress>, Error> {
41 let entries = self.kv.list_prefix(b"sender:").await?;
42 entries
43 .into_iter()
44 .map(|(k, _)| {
45 let key_str = String::from_utf8_lossy(&k);
46 let hex_part = key_str
47 .strip_prefix("sender:")
48 .ok_or_else(|| Error::InvalidData("invalid sender key prefix".into()))?;
49 Ok(AztecAddress(Fr::from_hex(hex_part)?))
50 })
51 .collect()
52 }
53}
54
55fn sender_key(sender: &AztecAddress) -> Vec<u8> {
56 format!("sender:{sender}").into_bytes()
57}
58
59#[cfg(test)]
60mod tests {
61 use super::*;
62 use crate::stores::InMemoryKvStore;
63
64 #[tokio::test]
65 async fn add_remove_senders() {
66 let kv = Arc::new(InMemoryKvStore::new());
67 let store = SenderStore::new(kv);
68 let sender = AztecAddress::from(99u64);
69
70 store.add(&sender).await.unwrap();
71 assert!(store.contains(&sender).await.unwrap());
72
73 let all = store.get_all().await.unwrap();
74 assert_eq!(all.len(), 1);
75
76 store.remove(&sender).await.unwrap();
77 assert!(!store.contains(&sender).await.unwrap());
78 assert!(store.get_all().await.unwrap().is_empty());
79 }
80
81 #[tokio::test]
82 async fn add_is_idempotent() {
83 let kv = Arc::new(InMemoryKvStore::new());
84 let store = SenderStore::new(kv);
85 let sender = AztecAddress::from(42u64);
86
87 store.add(&sender).await.unwrap();
88 store.add(&sender).await.unwrap();
89
90 let all = store.get_all().await.unwrap();
91 assert_eq!(all.len(), 1);
92 }
93}