diff --git a/key_protocol/src/key_management/key_tree/keys_private.rs b/key_protocol/src/key_management/key_tree/keys_private.rs index 6ffc8119..e30d638f 100644 --- a/key_protocol/src/key_management/key_tree/keys_private.rs +++ b/key_protocol/src/key_management/key_tree/keys_private.rs @@ -1,5 +1,5 @@ use k256::{Scalar, elliptic_curve::PrimeField as _}; -use nssa_core::{Identifier, NullifierPublicKey, encryption::ViewingPublicKey}; +use nssa_core::{NullifierPublicKey, PrivateAccountKind, encryption::ViewingPublicKey}; use serde::{Deserialize, Serialize}; use crate::key_management::{ @@ -10,7 +10,7 @@ use crate::key_management::{ #[derive(Debug, Serialize, Deserialize, Clone)] pub struct ChildKeysPrivate { - pub value: (KeyChain, Vec<(Identifier, nssa::Account)>), + pub value: (KeyChain, Vec<(PrivateAccountKind, nssa::Account)>), pub ccc: [u8; 32], /// Can be [`None`] if root. pub cci: Option, @@ -115,9 +115,8 @@ impl KeyTreeNode for ChildKeysPrivate { } fn account_ids(&self) -> impl Iterator { - self.value.1.iter().map(|(identifier, _)| { - nssa::AccountId::from((&self.value.0.nullifier_public_key, *identifier)) - }) + let npk = self.value.0.nullifier_public_key; + self.value.1.iter().map(move |(kind, _)| nssa::AccountId::for_private_account(&npk, kind)) } } diff --git a/key_protocol/src/key_management/key_tree/mod.rs b/key_protocol/src/key_management/key_tree/mod.rs index 0ae0a52f..47999ec7 100644 --- a/key_protocol/src/key_management/key_tree/mod.rs +++ b/key_protocol/src/key_management/key_tree/mod.rs @@ -319,6 +319,7 @@ mod tests { use std::{collections::HashSet, str::FromStr as _}; use nssa::AccountId; + use nssa_core::PrivateAccountKind; use super::*; @@ -532,7 +533,7 @@ mod tests { .get_mut(&ChainIndex::from_str("/1").unwrap()) .unwrap(); acc.value.1.push(( - 0, + PrivateAccountKind::Regular(0), nssa::Account { balance: 2, ..nssa::Account::default() @@ -544,7 +545,7 @@ mod tests { .get_mut(&ChainIndex::from_str("/2").unwrap()) .unwrap(); acc.value.1.push(( - 0, + PrivateAccountKind::Regular(0), nssa::Account { balance: 3, ..nssa::Account::default() @@ -556,7 +557,7 @@ mod tests { .get_mut(&ChainIndex::from_str("/0/1").unwrap()) .unwrap(); acc.value.1.push(( - 0, + PrivateAccountKind::Regular(0), nssa::Account { balance: 5, ..nssa::Account::default() @@ -568,7 +569,7 @@ mod tests { .get_mut(&ChainIndex::from_str("/1/0").unwrap()) .unwrap(); acc.value.1.push(( - 0, + PrivateAccountKind::Regular(0), nssa::Account { balance: 6, ..nssa::Account::default() diff --git a/key_protocol/src/key_protocol_core/mod.rs b/key_protocol/src/key_protocol_core/mod.rs index 4df6df82..e6c08d90 100644 --- a/key_protocol/src/key_protocol_core/mod.rs +++ b/key_protocol/src/key_protocol_core/mod.rs @@ -3,7 +3,7 @@ use std::collections::BTreeMap; use anyhow::Result; use k256::AffinePoint; use nssa::{Account, AccountId}; -use nssa_core::Identifier; +use nssa_core::{Identifier, PrivateAccountKind}; use serde::{Deserialize, Serialize}; use crate::key_management::{ @@ -17,7 +17,7 @@ pub type PublicKey = AffinePoint; #[derive(Clone, Debug, Serialize, Deserialize)] pub struct UserPrivateAccountData { pub key_chain: KeyChain, - pub accounts: Vec<(Identifier, Account)>, + pub accounts: Vec<(PrivateAccountKind, Account)>, } #[derive(Clone, Debug, Serialize, Deserialize)] @@ -53,9 +53,9 @@ impl NSSAUserData { ) -> bool { let mut check_res = true; for (account_id, entry) in accounts_keys_map { - let any_match = entry.accounts.iter().any(|(identifier, _)| { - nssa::AccountId::from((&entry.key_chain.nullifier_public_key, *identifier)) - == *account_id + let npk = &entry.key_chain.nullifier_public_key; + let any_match = entry.accounts.iter().any(|(kind, _)| { + nssa::AccountId::for_private_account(npk, kind) == *account_id }); if !any_match { println!("No matching entry found for account_id {account_id}"); @@ -155,24 +155,22 @@ impl NSSAUserData { ) -> Option<(KeyChain, nssa_core::account::Account, Identifier)> { // Check default accounts if let Some(entry) = self.default_user_private_accounts.get(&account_id) { - for (identifier, account) in &entry.accounts { - let expected_id = - nssa::AccountId::from((&entry.key_chain.nullifier_public_key, *identifier)); - if expected_id == account_id { - return Some((entry.key_chain.clone(), account.clone(), *identifier)); - } + let npk = &entry.key_chain.nullifier_public_key; + if let Some((kind, account)) = + entry.accounts.iter().find(|(kind, _)| nssa::AccountId::for_private_account(npk, kind) == account_id) + { + return Some((entry.key_chain.clone(), account.clone(), kind.identifier())); } return None; } // Check tree if let Some(node) = self.private_key_tree.get_node(account_id) { let key_chain = &node.value.0; - for (identifier, account) in &node.value.1 { - let expected_id = - nssa::AccountId::from((&key_chain.nullifier_public_key, *identifier)); - if expected_id == account_id { - return Some((key_chain.clone(), account.clone(), *identifier)); - } + let npk = &key_chain.nullifier_public_key; + if let Some((kind, account)) = + node.value.1.iter().find(|(kind, _)| nssa::AccountId::for_private_account(npk, kind) == account_id) + { + return Some((key_chain.clone(), account.clone(), kind.identifier())); } } None diff --git a/nssa/core/src/encryption/mod.rs b/nssa/core/src/encryption/mod.rs index 5ac5f020..6cf3c06a 100644 --- a/nssa/core/src/encryption/mod.rs +++ b/nssa/core/src/encryption/mod.rs @@ -22,9 +22,9 @@ pub type Scalar = [u8; 32]; /// to reconstruct the account's [`AccountId`] on the receiver side. /// /// [`AccountId`]: crate::account::AccountId -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub enum PrivateAccountKind { - Account(Identifier), + Regular(Identifier), Pda { program_id: ProgramId, seed: PdaSeed, @@ -33,14 +33,14 @@ pub enum PrivateAccountKind { } impl PrivateAccountKind { - /// Account(ident): 0x00 || ident (16 LE) || [0u8; 64] + /// Regular(ident): 0x00 || ident (16 LE) || [0u8; 64] /// Pda { program_id, seed, ident }: 0x01 || program_id (32 LE) || seed (32) || ident (16 LE) pub const HEADER_LEN: usize = 81; #[must_use] pub fn identifier(&self) -> Identifier { match self { - Self::Account(identifier) => *identifier, + Self::Regular(identifier) => *identifier, Self::Pda { identifier, .. } => *identifier, } } @@ -49,7 +49,7 @@ impl PrivateAccountKind { pub fn to_header_bytes(&self) -> [u8; Self::HEADER_LEN] { let mut bytes = [0u8; Self::HEADER_LEN]; match self { - Self::Account(identifier) => { + Self::Regular(identifier) => { bytes[0] = 0x00; bytes[1..17].copy_from_slice(&identifier.to_le_bytes()); // bytes[17..81] are zero padding @@ -72,7 +72,7 @@ impl PrivateAccountKind { match bytes[0] { 0x00 => { let identifier = Identifier::from_le_bytes(bytes[1..17].try_into().unwrap()); - Some(Self::Account(identifier)) + Some(Self::Regular(identifier)) } 0x01 => { let mut program_id = [0u32; 8]; @@ -208,7 +208,7 @@ mod tests { let account_ct = EncryptionScheme::encrypt( &account, - &PrivateAccountKind::Account(42), + &PrivateAccountKind::Regular(42), &secret, &commitment, 0, diff --git a/nssa/core/src/program.rs b/nssa/core/src/program.rs index 8a4a3a25..55c80b32 100644 --- a/nssa/core/src/program.rs +++ b/nssa/core/src/program.rs @@ -8,6 +8,7 @@ use serde::{Deserialize, Serialize}; use crate::{ BlockId, Identifier, NullifierPublicKey, Timestamp, account::{Account, AccountId, AccountWithMetadata}, + encryption::PrivateAccountKind, }; pub const DEFAULT_PROGRAM_ID: ProgramId = [0; 8]; @@ -96,6 +97,17 @@ impl AccountId { .expect("Hash output must be exactly 32 bytes long"), ) } + + /// Derives the [`AccountId`] for a private account from the nullifier public key and kind. + #[must_use] + pub fn for_private_account(npk: &NullifierPublicKey, kind: &PrivateAccountKind) -> Self { + match kind { + PrivateAccountKind::Regular(identifier) => Self::from((npk, *identifier)), + PrivateAccountKind::Pda { program_id, seed, identifier } => { + Self::for_private_pda(program_id, seed, npk, *identifier) + } + } + } } #[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq)] diff --git a/nssa/src/privacy_preserving_transaction/message.rs b/nssa/src/privacy_preserving_transaction/message.rs index 8256ab7c..6024ada7 100644 --- a/nssa/src/privacy_preserving_transaction/message.rs +++ b/nssa/src/privacy_preserving_transaction/message.rs @@ -253,7 +253,7 @@ pub mod tests { let esk = [3; 32]; let shared_secret = SharedSecretKey::new(&esk, &vpk); let epk = EphemeralPublicKey::from_scalar(esk); - let ciphertext = EncryptionScheme::encrypt(&account, &PrivateAccountKind::Account(0), &shared_secret, &commitment, 2); + let ciphertext = EncryptionScheme::encrypt(&account, &PrivateAccountKind::Regular(0), &shared_secret, &commitment, 2); let encrypted_account_data = EncryptedAccountData::new(ciphertext.clone(), &npk, &vpk, epk.clone()); diff --git a/program_methods/guest/src/bin/privacy_preserving_circuit.rs b/program_methods/guest/src/bin/privacy_preserving_circuit.rs index 3d16fdbb..0313f424 100644 --- a/program_methods/guest/src/bin/privacy_preserving_circuit.rs +++ b/program_methods/guest/src/bin/privacy_preserving_circuit.rs @@ -45,11 +45,11 @@ struct ExecutionState { /// `AccountId` entry or as an equality check against the existing one, making the rule: one /// `(program, seed)` → one account per tx. pda_family_binding: HashMap<(ProgramId, PdaSeed), AccountId>, - /// Map from a mask-3 `pre_state`'s position in `visibility_mask` to the npk supplied for - /// that position in `private_account_keys`. Built once in `derive_from_outputs` by walking - /// `visibility_mask` in lock-step with `private_account_keys`, used later by the claim and - /// caller-seeds authorization paths. - private_pda_npk_by_position: HashMap, + /// Map from a mask-3 `pre_state`'s position in `visibility_mask` to the (npk, identifier) + /// supplied for that position in `private_account_keys`. Built once in `derive_from_outputs` + /// by walking `visibility_mask` in lock-step with `private_account_keys`, used later by the + /// claim and caller-seeds authorization paths. + private_pda_npk_by_position: HashMap, } impl ExecutionState { @@ -64,18 +64,18 @@ impl ExecutionState { // pre_state order across all masks 1/2/3, so walk `visibility_mask` in lock-step. The // downstream `compute_circuit_output` also consumes the same iterator and its trailing // assertions catch an over-supply of keys; under-supply surfaces here. - let mut private_pda_npk_by_position: HashMap = HashMap::new(); + let mut private_pda_npk_by_position: HashMap = HashMap::new(); { let mut keys_iter = private_account_keys.iter(); for (pos, &mask) in visibility_mask.iter().enumerate() { if matches!(mask, 1..=3) { - let (npk, _, _) = keys_iter.next().unwrap_or_else(|| { + let (npk, identifier, _) = keys_iter.next().unwrap_or_else(|| { panic!( "private_account_keys shorter than visibility_mask demands: no key for masked position {pos} (mask {mask})" ) }); if mask == 3 { - private_pda_npk_by_position.insert(pos, *npk); + private_pda_npk_by_position.insert(pos, (*npk, *identifier)); } } } @@ -363,11 +363,11 @@ impl ExecutionState { ); } Claim::Pda(seed) => { - let npk = self + let (npk, identifier) = self .private_pda_npk_by_position .get(&pre_state_position) .expect("private PDA pre_state must have an npk in the position map"); - let pda = AccountId::for_private_pda(&program_id, &seed, npk, u128::MAX); + let pda = AccountId::for_private_pda(&program_id, &seed, npk, *identifier); assert_eq!( pre_account_id, pda, "Invalid private PDA claim for account {pre_account_id}" @@ -454,7 +454,7 @@ fn assert_family_binding( fn resolve_authorization_and_record_bindings( pda_family_binding: &mut HashMap<(ProgramId, PdaSeed), AccountId>, private_pda_bound_positions: &mut HashMap, - private_pda_npk_by_position: &HashMap, + private_pda_npk_by_position: &HashMap, pre_account_id: AccountId, pre_state_position: usize, caller_program_id: Option, @@ -467,8 +467,8 @@ fn resolve_authorization_and_record_bindings( if AccountId::for_public_pda(&caller, seed) == pre_account_id { return Some((*seed, false, caller)); } - if let Some(npk) = private_pda_npk_by_position.get(&pre_state_position) - && AccountId::for_private_pda(&caller, seed, npk, u128::MAX) == pre_account_id + if let Some((npk, identifier)) = private_pda_npk_by_position.get(&pre_state_position) + && AccountId::for_private_pda(&caller, seed, npk, *identifier) == pre_account_id { return Some((*seed, true, caller)); } @@ -608,7 +608,7 @@ fn compute_circuit_output( // Encrypt and push post state let encrypted_account = EncryptionScheme::encrypt( &post_with_updated_nonce, - &PrivateAccountKind::Account(*identifier), + &PrivateAccountKind::Regular(*identifier), shared_secret, &commitment_post, output_index, diff --git a/wallet/src/chain_storage.rs b/wallet/src/chain_storage.rs index 4b160b66..81d83d22 100644 --- a/wallet/src/chain_storage.rs +++ b/wallet/src/chain_storage.rs @@ -73,8 +73,8 @@ impl WalletChainStore { PersistentAccountData::Private(data) => { let npk = data.data.value.0.nullifier_public_key; let chain_index = data.chain_index; - for identifier in &data.identifiers { - let account_id = nssa::AccountId::from((&npk, *identifier)); + for kind in &data.kinds { + let account_id = nssa::AccountId::for_private_account(&npk, kind); private_tree .account_id_map .insert(account_id, chain_index.clone()); @@ -90,7 +90,7 @@ impl WalletChainStore { data.account_id(), UserPrivateAccountData { key_chain: data.key_chain, - accounts: vec![(data.identifier, data.account)], + accounts: vec![(PrivateAccountKind::Regular(data.identifier), data.account)], }, ); } @@ -136,7 +136,7 @@ impl WalletChainStore { account_id, UserPrivateAccountData { key_chain: data.key_chain, - accounts: vec![(data.identifier, account)], + accounts: vec![(PrivateAccountKind::Regular(data.identifier), account)], }, ); } @@ -195,7 +195,6 @@ impl WalletChainStore { account: nssa_core::account::Account, ) { debug!("inserting at address {account_id}, this account {account:?}"); - let identifier = kind.identifier(); // Update default accounts if present if let Entry::Occupied(mut entry) = self @@ -204,10 +203,10 @@ impl WalletChainStore { .entry(account_id) { let entry = entry.get_mut(); - if let Some((_, acc)) = entry.accounts.iter_mut().find(|(id, _)| *id == identifier) { + if let Some((_, acc)) = entry.accounts.iter_mut().find(|(k, _)| k == kind) { *acc = account; } else { - entry.accounts.push((identifier, account)); + entry.accounts.push((kind.clone(), account)); } return; } @@ -230,29 +229,21 @@ impl WalletChainStore { .key_map .get_mut(&chain_index) { - if let Some((_, acc)) = node.value.1.iter_mut().find(|(id, _)| *id == identifier) { + if let Some((_, acc)) = node.value.1.iter_mut().find(|(k, _)| k == kind) { *acc = account; } else { - node.value.1.push((identifier, account)); + node.value.1.push((kind.clone(), account)); } } } else { // Node not yet in account_id_map — find it by checking all nodes for (ci, node) in &mut self.user_data.private_key_tree.key_map { let npk = &node.value.0.nullifier_public_key; - let expected_id = match kind { - PrivateAccountKind::Account(id) => nssa::AccountId::from((npk, *id)), - PrivateAccountKind::Pda { program_id, seed, identifier: id } => { - nssa::AccountId::for_private_pda(program_id, seed, npk, *id) - } - }; - if expected_id == account_id { - if let Some((_, acc)) = - node.value.1.iter_mut().find(|(id, _)| *id == identifier) - { + if nssa::AccountId::for_private_account(npk, kind) == account_id { + if let Some((_, acc)) = node.value.1.iter_mut().find(|(k, _)| k == kind) { *acc = account; } else { - node.value.1.push((identifier, account)); + node.value.1.push((kind.clone(), account)); } // Register in account_id_map self.user_data @@ -298,7 +289,7 @@ mod tests { data: public_data, }), PersistentAccountData::Private(Box::new(PersistentAccountDataPrivate { - identifiers: vec![], + kinds: vec![], chain_index: ChainIndex::root(), data: private_data, })), diff --git a/wallet/src/config.rs b/wallet/src/config.rs index bbd98ac7..e2905103 100644 --- a/wallet/src/config.rs +++ b/wallet/src/config.rs @@ -28,7 +28,7 @@ pub struct PersistentAccountDataPublic { #[derive(Debug, Clone, Serialize, Deserialize)] pub struct PersistentAccountDataPrivate { - pub identifiers: Vec, + pub kinds: Vec, pub chain_index: ChainIndex, pub data: ChildKeysPrivate, } diff --git a/wallet/src/helperfunctions.rs b/wallet/src/helperfunctions.rs index 94755f6e..04f780dd 100644 --- a/wallet/src/helperfunctions.rs +++ b/wallet/src/helperfunctions.rs @@ -166,10 +166,10 @@ pub fn produce_data_for_storage( } for (chain_index, node) in &user_data.private_key_tree.key_map { - let identifiers = node.value.1.iter().map(|(id, _)| *id).collect(); + let kinds = node.value.1.iter().map(|(kind, _)| kind.clone()).collect(); vec_for_storage.push( PersistentAccountDataPrivate { - identifiers, + kinds, chain_index: chain_index.clone(), data: node.clone(), } @@ -188,12 +188,12 @@ pub fn produce_data_for_storage( } for entry in user_data.default_user_private_accounts.values() { - for (identifier, account) in &entry.accounts { + for (kind, account) in &entry.accounts { vec_for_storage.push( InitialAccountData::Private(Box::new(PrivateAccountPrivateInitialData { account: account.clone(), key_chain: entry.key_chain.clone(), - identifier: *identifier, + identifier: kind.identifier(), })) .into(), ); diff --git a/wallet/src/lib.rs b/wallet/src/lib.rs index ee9f3dd1..97c4674b 100644 --- a/wallet/src/lib.rs +++ b/wallet/src/lib.rs @@ -284,7 +284,7 @@ impl WalletCore { .nullifier_public_key; let account_id = AccountId::from((&npk, identifier)); self.storage - .insert_private_account_data(account_id, &PrivateAccountKind::Account(identifier), Account::default()); + .insert_private_account_data(account_id, &PrivateAccountKind::Regular(identifier), Account::default()); (account_id, cci) } @@ -548,7 +548,7 @@ impl WalletCore { .map(|(kind, res_acc)| { let npk = &key_chain.nullifier_public_key; let account_id = match &kind { - PrivateAccountKind::Account(identifier) => { + PrivateAccountKind::Regular(identifier) => { nssa::AccountId::from((npk, *identifier)) } PrivateAccountKind::Pda { program_id, seed, identifier } => { diff --git a/wallet/src/privacy_preserving_tx.rs b/wallet/src/privacy_preserving_tx.rs index 3df2ecc1..d890bcc8 100644 --- a/wallet/src/privacy_preserving_tx.rs +++ b/wallet/src/privacy_preserving_tx.rs @@ -18,6 +18,17 @@ pub enum PrivacyPreservingAccount { vpk: ViewingPublicKey, identifier: Identifier, }, + /// An owned private PDA: wallet holds the nsk/npk; account_id was derived via + /// `AccountId::for_private_pda`. Produces visibility mask 3. + PrivatePdaOwned(AccountId), + /// A foreign private PDA: wallet knows the recipient's npk/vpk but not their nsk. + /// Produces visibility mask 3 with a default (uninitialised) account. + PrivatePdaForeign { + account_id: AccountId, + npk: NullifierPublicKey, + vpk: ViewingPublicKey, + identifier: Identifier, + }, } impl PrivacyPreservingAccount { @@ -31,11 +42,9 @@ impl PrivacyPreservingAccount { matches!( &self, Self::PrivateOwned(_) - | Self::PrivateForeign { - npk: _, - vpk: _, - identifier: _ - } + | Self::PrivateForeign { .. } + | Self::PrivatePdaOwned(_) + | Self::PrivatePdaForeign { .. } ) } } @@ -106,6 +115,28 @@ impl AccountManager { (State::Private(pre), 2) } + PrivacyPreservingAccount::PrivatePdaOwned(account_id) => { + let pre = private_acc_preparation(wallet, account_id).await?; + (State::Private(pre), 3) + } + PrivacyPreservingAccount::PrivatePdaForeign { + account_id, + npk, + vpk, + identifier, + } => { + let acc = nssa_core::account::Account::default(); + let auth_acc = AccountWithMetadata::new(acc, false, account_id); + let pre = AccountPreparedData { + nsk: None, + npk, + identifier, + vpk, + pre_state: auth_acc, + proof: None, + }; + (State::Private(pre), 3) + } }; pre_states.push(state); @@ -235,7 +266,7 @@ async fn private_acc_preparation( // TODO: Technically we could allow unauthorized owned accounts, but currently we don't have // support from that in the wallet. - let sender_pre = AccountWithMetadata::new(from_acc.clone(), true, (&from_npk, from_identifier)); + let sender_pre = AccountWithMetadata::new(from_acc.clone(), true, account_id); Ok(AccountPreparedData { nsk: Some(nsk),