aztec_pxe/stores/
key_store.rs

1//! Key storage with master key derivation.
2
3use std::sync::Arc;
4
5use aztec_core::error::Error;
6use aztec_core::types::{AztecAddress, Fr, GrumpkinScalar, PublicKeys};
7use aztec_crypto::keys::{
8    compute_app_nullifier_hiding_key, compute_app_secret_key, compute_ovsk_app, derive_keys,
9    DerivedKeys, KeyType,
10};
11
12use super::kv::KvStore;
13
14/// Stores master secret keys and derives app-scoped keys on demand.
15pub struct KeyStore {
16    kv: Arc<dyn KvStore>,
17}
18
19impl KeyStore {
20    pub fn new(kv: Arc<dyn KvStore>) -> Self {
21        Self { kv }
22    }
23
24    /// Add an account by storing its secret key, indexed by the public keys hash.
25    pub async fn add_account(&self, secret_key: &Fr) -> Result<DerivedKeys, Error> {
26        let derived = derive_keys(secret_key);
27        let pk_hash = derived.public_keys.hash();
28        let key = account_key(&pk_hash);
29        let value = secret_key.to_be_bytes();
30        self.kv.put(&key, &value).await?;
31        Ok(derived)
32    }
33
34    /// Get the secret key for an account identified by its public keys hash.
35    pub async fn get_secret_key(&self, pk_hash: &Fr) -> Result<Option<Fr>, Error> {
36        let key = account_key(pk_hash);
37        match self.kv.get(&key).await? {
38            Some(bytes) => {
39                let mut arr = [0u8; 32];
40                if bytes.len() == 32 {
41                    arr.copy_from_slice(&bytes);
42                    Ok(Some(Fr::from(arr)))
43                } else {
44                    Err(Error::InvalidData("invalid secret key length".into()))
45                }
46            }
47            None => Ok(None),
48        }
49    }
50
51    /// Get the master nullifier hiding key for an account.
52    pub async fn get_master_nullifier_hiding_key(
53        &self,
54        pk_hash: &Fr,
55    ) -> Result<Option<GrumpkinScalar>, Error> {
56        match self.get_secret_key(pk_hash).await? {
57            Some(sk) => Ok(Some(derive_keys(&sk).master_nullifier_hiding_key)),
58            None => Ok(None),
59        }
60    }
61
62    /// Get the master incoming viewing secret key for an account.
63    pub async fn get_master_incoming_viewing_secret_key(
64        &self,
65        pk_hash: &Fr,
66    ) -> Result<Option<GrumpkinScalar>, Error> {
67        match self.get_secret_key(pk_hash).await? {
68            Some(sk) => Ok(Some(derive_keys(&sk).master_incoming_viewing_secret_key)),
69            None => Ok(None),
70        }
71    }
72
73    /// Get the master outgoing viewing secret key for an account.
74    pub async fn get_master_outgoing_viewing_secret_key(
75        &self,
76        pk_hash: &Fr,
77    ) -> Result<Option<GrumpkinScalar>, Error> {
78        match self.get_secret_key(pk_hash).await? {
79            Some(sk) => Ok(Some(derive_keys(&sk).master_outgoing_viewing_secret_key)),
80            None => Ok(None),
81        }
82    }
83
84    /// Get the master tagging secret key for an account.
85    pub async fn get_master_tagging_secret_key(
86        &self,
87        pk_hash: &Fr,
88    ) -> Result<Option<GrumpkinScalar>, Error> {
89        match self.get_secret_key(pk_hash).await? {
90            Some(sk) => Ok(Some(derive_keys(&sk).master_tagging_secret_key)),
91            None => Ok(None),
92        }
93    }
94
95    /// Get the public keys for an account.
96    pub async fn get_public_keys(&self, pk_hash: &Fr) -> Result<Option<PublicKeys>, Error> {
97        match self.get_secret_key(pk_hash).await? {
98            Some(sk) => Ok(Some(derive_keys(&sk).public_keys)),
99            None => Ok(None),
100        }
101    }
102
103    /// Compute the app-scoped nullifier hiding key.
104    pub async fn get_app_nullifier_hiding_key(
105        &self,
106        pk_hash: &Fr,
107        app: &AztecAddress,
108    ) -> Result<Option<Fr>, Error> {
109        match self.get_master_nullifier_hiding_key(pk_hash).await? {
110            Some(nhk_m) => Ok(Some(compute_app_nullifier_hiding_key(&nhk_m, app))),
111            None => Ok(None),
112        }
113    }
114
115    /// Find which account owns a given master public key hash, then return
116    /// the corresponding public key point and app-siloed secret key.
117    ///
118    /// Mirrors the TS `getKeyValidationRequest(pkMHash, contractAddress)`.
119    pub async fn get_key_validation_request(
120        &self,
121        pk_m_hash: &Fr,
122        app: &AztecAddress,
123    ) -> Result<Option<(aztec_core::types::Point, Fr)>, Error> {
124        use aztec_core::hash::poseidon2_hash;
125
126        let accounts = self.get_accounts().await?;
127        for account_pk_hash in &accounts {
128            let Some(sk) = self.get_secret_key(account_pk_hash).await? else {
129                continue;
130            };
131            let derived = derive_keys(&sk);
132            let pk = &derived.public_keys;
133            let keys_and_types = [
134                (
135                    &pk.master_nullifier_public_key,
136                    &derived.master_nullifier_hiding_key,
137                    KeyType::Nullifier,
138                ),
139                (
140                    &pk.master_incoming_viewing_public_key,
141                    &derived.master_incoming_viewing_secret_key,
142                    KeyType::IncomingViewing,
143                ),
144                (
145                    &pk.master_outgoing_viewing_public_key,
146                    &derived.master_outgoing_viewing_secret_key,
147                    KeyType::OutgoingViewing,
148                ),
149                (
150                    &pk.master_tagging_public_key,
151                    &derived.master_tagging_secret_key,
152                    KeyType::Tagging,
153                ),
154            ];
155
156            for (pk_m, sk_m, key_type) in &keys_and_types {
157                let hash = poseidon2_hash(&[pk_m.x, pk_m.y, Fr::from(pk_m.is_infinite)]);
158                if hash == *pk_m_hash {
159                    let sk_app = compute_app_secret_key(sk_m, app, *key_type);
160                    return Ok(Some(((*pk_m).clone(), sk_app)));
161                }
162            }
163        }
164        Ok(None)
165    }
166
167    /// Compute an app-scoped secret key for a given key type.
168    pub async fn get_app_secret_key(
169        &self,
170        pk_hash: &Fr,
171        app: &AztecAddress,
172        key_type: KeyType,
173    ) -> Result<Option<Fr>, Error> {
174        let sk = match self.get_secret_key(pk_hash).await? {
175            Some(sk) => sk,
176            None => return Ok(None),
177        };
178        let derived = derive_keys(&sk);
179        let master_key = match key_type {
180            KeyType::Nullifier => &derived.master_nullifier_hiding_key,
181            KeyType::IncomingViewing => &derived.master_incoming_viewing_secret_key,
182            KeyType::OutgoingViewing => &derived.master_outgoing_viewing_secret_key,
183            KeyType::Tagging => &derived.master_tagging_secret_key,
184        };
185        Ok(Some(compute_app_secret_key(master_key, app, key_type)))
186    }
187
188    /// Compute the app-scoped outgoing viewing secret key (returns GrumpkinScalar).
189    pub async fn get_app_ovsk(
190        &self,
191        pk_hash: &Fr,
192        app: &AztecAddress,
193    ) -> Result<Option<GrumpkinScalar>, Error> {
194        match self.get_master_outgoing_viewing_secret_key(pk_hash).await? {
195            Some(ovsk_m) => Ok(Some(compute_ovsk_app(&ovsk_m, app))),
196            None => Ok(None),
197        }
198    }
199
200    /// List all stored public keys hashes.
201    pub async fn get_accounts(&self) -> Result<Vec<Fr>, Error> {
202        let entries = self.kv.list_prefix(b"key:account:").await?;
203        entries
204            .into_iter()
205            .map(|(k, _)| {
206                let key_str = String::from_utf8_lossy(&k);
207                let hex_part = key_str
208                    .strip_prefix("key:account:")
209                    .ok_or_else(|| Error::InvalidData("invalid key prefix".into()))?;
210                Fr::from_hex(hex_part)
211            })
212            .collect()
213    }
214}
215
216fn account_key(pk_hash: &Fr) -> Vec<u8> {
217    format!("key:account:{pk_hash}").into_bytes()
218}
219
220#[cfg(test)]
221mod tests {
222    use super::*;
223    use crate::stores::InMemoryKvStore;
224
225    #[tokio::test]
226    async fn add_and_retrieve_account() {
227        let kv = Arc::new(InMemoryKvStore::new());
228        let store = KeyStore::new(kv);
229        let sk = Fr::from(8923u64);
230
231        let derived = store.add_account(&sk).await.unwrap();
232        let pk_hash = derived.public_keys.hash();
233
234        let retrieved_sk = store.get_secret_key(&pk_hash).await.unwrap().unwrap();
235        assert_eq!(retrieved_sk, sk);
236
237        let public_keys = store.get_public_keys(&pk_hash).await.unwrap().unwrap();
238        assert_eq!(public_keys.hash(), pk_hash);
239    }
240}
aztec_pxe/stores/key_store.rs

aztec_pxe/stores/
key_store.rs
