aztec_pxe/stores/
sender_store.rs

1//! Sender address storage for private log discovery.
2
3use std::sync::Arc;
4
5use aztec_core::error::Error;
6use aztec_core::types::{AztecAddress, Fr};
7
8use super::kv::KvStore;
9
10/// Stores registered sender addresses for private log discovery.
11pub 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    /// Register a sender address.
21    pub async fn add(&self, sender: &AztecAddress) -> Result<(), Error> {
22        let key = sender_key(sender);
23        // Store a minimal value — presence of the key is what matters.
24        self.kv.put(&key, &[1]).await
25    }
26
27    /// Remove a sender address.
28    pub async fn remove(&self, sender: &AztecAddress) -> Result<(), Error> {
29        let key = sender_key(sender);
30        self.kv.delete(&key).await
31    }
32
33    /// Check if a sender is registered.
34    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    /// List all registered sender addresses.
40    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}