aztec_pxe/stores/
address_store.rs

1//! Complete address storage.
2
3use std::sync::Arc;
4
5use aztec_core::error::Error;
6use aztec_core::types::{AztecAddress, CompleteAddress};
7
8use super::kv::KvStore;
9
10/// Stores `CompleteAddress` records for registered accounts.
11pub struct AddressStore {
12    kv: Arc<dyn KvStore>,
13}
14
15impl AddressStore {
16    pub fn new(kv: Arc<dyn KvStore>) -> Self {
17        Self { kv }
18    }
19
20    /// Store a complete address.
21    pub async fn add(&self, complete: &CompleteAddress) -> Result<(), Error> {
22        let key = address_key(&complete.address);
23        let value = serde_json::to_vec(complete)?;
24        self.kv.put(&key, &value).await
25    }
26
27    /// Get a complete address by its Aztec address.
28    pub async fn get(&self, address: &AztecAddress) -> Result<Option<CompleteAddress>, Error> {
29        let key = address_key(address);
30        match self.kv.get(&key).await? {
31            Some(bytes) => Ok(Some(serde_json::from_slice(&bytes)?)),
32            None => Ok(None),
33        }
34    }
35
36    /// List all registered complete addresses.
37    pub async fn get_all(&self) -> Result<Vec<CompleteAddress>, Error> {
38        let entries = self.kv.list_prefix(b"address:").await?;
39        entries
40            .into_iter()
41            .map(|(_, v)| Ok(serde_json::from_slice(&v)?))
42            .collect()
43    }
44}
45
46fn address_key(address: &AztecAddress) -> Vec<u8> {
47    format!("address:{address}").into_bytes()
48}
49
50#[cfg(test)]
51mod tests {
52    use super::*;
53    use crate::stores::InMemoryKvStore;
54    use aztec_core::types::{Fr, PublicKeys};
55
56    #[tokio::test]
57    async fn store_and_retrieve_address() {
58        let kv = Arc::new(InMemoryKvStore::new());
59        let store = AddressStore::new(kv);
60        let complete = CompleteAddress {
61            address: AztecAddress::from(42u64),
62            public_keys: PublicKeys::default(),
63            partial_address: Fr::from(1u64),
64        };
65
66        store.add(&complete).await.unwrap();
67        let retrieved = store.get(&complete.address).await.unwrap().unwrap();
68        assert_eq!(retrieved.address, complete.address);
69    }
70}