aztec_pxe/stores/
capsule_store.rs

1//! Ephemeral capsule storage for private execution.
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 ephemeral capsule data that is consumed during execution.
11///
12/// Capsules are private data blobs passed to contract functions via oracle calls.
13/// They are consumed (deleted) after being read.
14pub struct CapsuleStore {
15    kv: Arc<dyn KvStore>,
16}
17
18impl CapsuleStore {
19    pub fn new(kv: Arc<dyn KvStore>) -> Self {
20        Self { kv }
21    }
22
23    /// Store a capsule (a list of field element arrays).
24    pub async fn add(&self, contract: &Fr, capsule: &[Vec<Fr>]) -> Result<(), Error> {
25        let key = capsule_key(contract);
26        let value = serde_json::to_vec(capsule)?;
27        self.kv.put(&key, &value).await
28    }
29
30    /// Pop a capsule for the given contract (consumes it).
31    pub async fn pop(&self, contract: &Fr) -> Result<Option<Vec<Vec<Fr>>>, Error> {
32        let key = capsule_key(contract);
33        match self.kv.get(&key).await? {
34            Some(bytes) => {
35                self.kv.delete(&key).await?;
36                Ok(Some(serde_json::from_slice(&bytes)?))
37            }
38            None => Ok(None),
39        }
40    }
41
42    /// Store capsule data in a contract-scoped slot.
43    pub async fn store_capsule(
44        &self,
45        contract_address: &AztecAddress,
46        slot: &Fr,
47        capsule: &[Fr],
48    ) -> Result<(), Error> {
49        let key = db_slot_key(contract_address, slot);
50        let value = serde_json::to_vec(capsule)?;
51        self.kv.put(&key, &value).await
52    }
53
54    /// Load capsule data from a contract-scoped slot.
55    pub async fn load_capsule(
56        &self,
57        contract_address: &AztecAddress,
58        slot: &Fr,
59    ) -> Result<Option<Vec<Fr>>, Error> {
60        let key = db_slot_key(contract_address, slot);
61        match self.kv.get(&key).await? {
62            Some(bytes) => Ok(Some(serde_json::from_slice(&bytes)?)),
63            None => Ok(None),
64        }
65    }
66
67    /// Delete capsule data from a contract-scoped slot.
68    pub async fn delete_capsule(
69        &self,
70        contract_address: &AztecAddress,
71        slot: &Fr,
72    ) -> Result<(), Error> {
73        let key = db_slot_key(contract_address, slot);
74        self.kv.delete(&key).await
75    }
76
77    /// Copy a contiguous region of contract-scoped slots.
78    pub async fn copy_capsule(
79        &self,
80        contract_address: &AztecAddress,
81        src_slot: &Fr,
82        dst_slot: &Fr,
83        num_entries: usize,
84    ) -> Result<(), Error> {
85        if num_entries == 0 {
86            return Ok(());
87        }
88
89        let mut copied = Vec::with_capacity(num_entries);
90        for i in 0..num_entries {
91            let current_src = Fr::from((src_slot.to_usize() + i) as u64);
92            let data = self
93                .load_capsule(contract_address, &current_src)
94                .await?
95                .ok_or_else(|| {
96                    Error::InvalidData(format!(
97                        "attempted to copy empty capsule slot {} for contract {}",
98                        current_src, contract_address
99                    ))
100                })?;
101            copied.push(data);
102        }
103
104        for (i, data) in copied.into_iter().enumerate() {
105            let current_dst = Fr::from((dst_slot.to_usize() + i) as u64);
106            self.store_capsule(contract_address, &current_dst, &data)
107                .await?;
108        }
109
110        Ok(())
111    }
112
113    /// Append entries to a capsule array stored at `base_slot`.
114    pub async fn append_to_capsule_array(
115        &self,
116        contract_address: &AztecAddress,
117        base_slot: &Fr,
118        content: &[Vec<Fr>],
119    ) -> Result<(), Error> {
120        let current_length = self
121            .load_capsule(contract_address, base_slot)
122            .await?
123            .and_then(|capsule| capsule.first().copied())
124            .unwrap_or_else(Fr::zero)
125            .to_usize();
126
127        for (i, capsule) in content.iter().enumerate() {
128            let next_slot = array_slot(base_slot, current_length + i);
129            self.store_capsule(contract_address, &next_slot, capsule)
130                .await?;
131        }
132
133        self.store_capsule(
134            contract_address,
135            base_slot,
136            &[Fr::from((current_length + content.len()) as u64)],
137        )
138        .await
139    }
140
141    /// Read all entries from a capsule array stored at `base_slot`.
142    pub async fn read_capsule_array(
143        &self,
144        contract_address: &AztecAddress,
145        base_slot: &Fr,
146    ) -> Result<Vec<Vec<Fr>>, Error> {
147        let length = self
148            .load_capsule(contract_address, base_slot)
149            .await?
150            .and_then(|capsule| capsule.first().copied())
151            .unwrap_or_else(Fr::zero)
152            .to_usize();
153
154        let mut values = Vec::with_capacity(length);
155        for i in 0..length {
156            let slot = array_slot(base_slot, i);
157            let value = self
158                .load_capsule(contract_address, &slot)
159                .await?
160                .ok_or_else(|| {
161                    Error::InvalidData(format!(
162                        "expected non-empty capsule array value at slot {} for contract {}",
163                        slot, contract_address
164                    ))
165                })?;
166            values.push(value);
167        }
168
169        Ok(values)
170    }
171
172    /// Replace the entire capsule array stored at `base_slot`.
173    pub async fn set_capsule_array(
174        &self,
175        contract_address: &AztecAddress,
176        base_slot: &Fr,
177        content: &[Vec<Fr>],
178    ) -> Result<(), Error> {
179        let original_length = self
180            .load_capsule(contract_address, base_slot)
181            .await?
182            .and_then(|capsule| capsule.first().copied())
183            .unwrap_or_else(Fr::zero)
184            .to_usize();
185
186        self.store_capsule(
187            contract_address,
188            base_slot,
189            &[Fr::from(content.len() as u64)],
190        )
191        .await?;
192
193        for (i, capsule) in content.iter().enumerate() {
194            let slot = array_slot(base_slot, i);
195            self.store_capsule(contract_address, &slot, capsule).await?;
196        }
197
198        for i in content.len()..original_length {
199            let slot = array_slot(base_slot, i);
200            self.delete_capsule(contract_address, &slot).await?;
201        }
202
203        Ok(())
204    }
205}
206
207fn capsule_key(contract: &Fr) -> Vec<u8> {
208    format!("capsule:{contract}").into_bytes()
209}
210
211fn db_slot_key(contract_address: &AztecAddress, slot: &Fr) -> Vec<u8> {
212    format!("capsule_db:{contract_address}:{slot}").into_bytes()
213}
214
215/// Compute the capsule array entry slot using field addition, matching
216/// the Noir formula: `base_slot + 1 + index`.
217fn array_slot(base_slot: &Fr, index: usize) -> Fr {
218    // Use field arithmetic (not integer arithmetic) to match Noir.
219    Fr(base_slot.0 + Fr::from(1u64 + index as u64).0)
220}
221
222#[cfg(test)]
223mod tests {
224    use super::*;
225    use crate::stores::InMemoryKvStore;
226
227    #[tokio::test]
228    async fn capsule_is_consumed_on_pop() {
229        let kv = Arc::new(InMemoryKvStore::new());
230        let store = CapsuleStore::new(kv);
231        let contract = Fr::from(1u64);
232        let capsule = vec![vec![Fr::from(10u64), Fr::from(20u64)]];
233
234        store.add(&contract, &capsule).await.unwrap();
235
236        let first = store.pop(&contract).await.unwrap();
237        assert!(first.is_some());
238
239        let second = store.pop(&contract).await.unwrap();
240        assert!(second.is_none());
241    }
242
243    #[tokio::test]
244    async fn contract_scoped_capsule_array_roundtrip() {
245        let kv = Arc::new(InMemoryKvStore::new());
246        let store = CapsuleStore::new(kv);
247        let contract = AztecAddress::from(1u64);
248        let base_slot = Fr::from(10u64);
249
250        store
251            .append_to_capsule_array(
252                &contract,
253                &base_slot,
254                &[
255                    vec![Fr::from(11u64)],
256                    vec![Fr::from(12u64), Fr::from(13u64)],
257                ],
258            )
259            .await
260            .unwrap();
261
262        let values = store
263            .read_capsule_array(&contract, &base_slot)
264            .await
265            .unwrap();
266        assert_eq!(values.len(), 2);
267        assert_eq!(values[0], vec![Fr::from(11u64)]);
268        assert_eq!(values[1], vec![Fr::from(12u64), Fr::from(13u64)]);
269
270        store
271            .set_capsule_array(&contract, &base_slot, &[vec![Fr::from(99u64)]])
272            .await
273            .unwrap();
274
275        let values = store
276            .read_capsule_array(&contract, &base_slot)
277            .await
278            .unwrap();
279        assert_eq!(values, vec![vec![Fr::from(99u64)]]);
280    }
281}
aztec_pxe/stores/capsule_store.rs

aztec_pxe/stores/
capsule_store.rs
