diff --git a/integration_tests/tests/auth_transfer/private.rs b/integration_tests/tests/auth_transfer/private.rs index cf02d0ac..f66dbdbb 100644 --- a/integration_tests/tests/auth_transfer/private.rs +++ b/integration_tests/tests/auth_transfer/private.rs @@ -170,14 +170,15 @@ async fn private_transfer_to_owned_account_using_claiming_path() -> Result<()> { }; // Get the keys for the newly created account - let (to_keys, _) = ctx + let bundle = ctx .wallet() .storage() .user_data .get_private_account(to_account_id) - .cloned() .context("Failed to get private account")?; + let to_keys = bundle.key_chain; + // Send to this account using claiming path (using npk and vpk instead of account ID) let command = Command::AuthTransfer(AuthTransferSubcommand::Send { from: Some(format_private_account_id(from)), @@ -336,14 +337,15 @@ async fn private_transfer_to_owned_account_continuous_run_path() -> Result<()> { }; // Get the newly created account's keys - let (to_keys, _) = ctx + let bundle = ctx .wallet() .storage() .user_data .get_private_account(to_account_id) - .cloned() .context("Failed to get private account")?; + let to_keys = bundle.key_chain; + // Send transfer using nullifier and viewing public keys let command = Command::AuthTransfer(AuthTransferSubcommand::Send { from: Some(format_private_account_id(from)), diff --git a/integration_tests/tests/keys_restoration.rs b/integration_tests/tests/keys_restoration.rs index 8dca027c..54b9474a 100644 --- a/integration_tests/tests/keys_restoration.rs +++ b/integration_tests/tests/keys_restoration.rs @@ -58,15 +58,15 @@ async fn sync_private_account_with_non_zero_chain_index() -> Result<()> { anyhow::bail!("Expected RegisterAccount return value"); }; - // Get the keys for the newly created account - let (to_keys, _) = ctx + let bundle = ctx .wallet() .storage() .user_data .get_private_account(to_account_id) - .cloned() .context("Failed to get private account")?; + let to_keys = bundle.key_chain; + // Send to this account using claiming path (using npk and vpk instead of account ID) let command = Command::AuthTransfer(AuthTransferSubcommand::Send { from: Some(format_private_account_id(from)), diff --git a/integration_tests/tests/token.rs b/integration_tests/tests/token.rs index e40e27c8..f69bb197 100644 --- a/integration_tests/tests/token.rs +++ b/integration_tests/tests/token.rs @@ -1131,13 +1131,13 @@ async fn token_claiming_path_with_private_accounts() -> Result<()> { }; // Get keys for foreign mint (claiming path) - let (holder_keys, _) = ctx + let holder_keys = ctx .wallet() .storage() .user_data .get_private_account(recipient_account_id) - .cloned() - .context("Failed to get private account keys")?; + .unwrap() + .key_chain; // Mint using claiming path (foreign account) let mint_amount = 9; diff --git a/integration_tests/tests/wallet_ffi.rs b/integration_tests/tests/wallet_ffi.rs index ac548280..f1dd9627 100644 --- a/integration_tests/tests/wallet_ffi.rs +++ b/integration_tests/tests/wallet_ffi.rs @@ -585,7 +585,7 @@ fn test_wallet_ffi_get_private_account_keys() -> Result<()> { .user_data .get_private_account(account_id) .unwrap() - .0; + .key_chain; let expected_npk = &key_chain.nullifier_public_key; let expected_vpk = &key_chain.viewing_public_key; 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 42130b1f..53ff223e 100644 --- a/key_protocol/src/key_management/key_tree/keys_private.rs +++ b/key_protocol/src/key_management/key_tree/keys_private.rs @@ -2,10 +2,13 @@ use k256::{Scalar, elliptic_curve::PrimeField as _}; use nssa_core::{NullifierPublicKey, encryption::ViewingPublicKey}; use serde::{Deserialize, Serialize}; -use crate::key_management::{ - KeyChain, - key_tree::traits::KeyNode, - secret_holders::{PrivateKeyHolder, SecretSpendingKey}, +use crate::{ + key_management::{ + KeyChain, + key_tree::traits::KeyNode, + secret_holders::{PrivateKeyHolder, SecretSpendingKey}, + }, + key_protocol_core::PrivateBundle, }; #[derive(Debug, Serialize, Deserialize, Clone)] @@ -129,9 +132,12 @@ impl<'a> From<&'a ChildKeysPrivate> for &'a (KeyChain, nssa::Account) { clippy::single_char_lifetime_names, reason = "TODO add meaningful name" )] -impl<'a> From<&'a mut ChildKeysPrivate> for &'a mut (KeyChain, nssa::Account) { +impl<'a> From<&'a mut ChildKeysPrivate> for PrivateBundle { fn from(value: &'a mut ChildKeysPrivate) -> Self { - &mut value.value + Self { + key_chain: value.value.0.clone(), + account: value.value.1.clone(), + } } } diff --git a/key_protocol/src/key_protocol_core/mod.rs b/key_protocol/src/key_protocol_core/mod.rs index 8186865f..bb771331 100644 --- a/key_protocol/src/key_protocol_core/mod.rs +++ b/key_protocol/src/key_protocol_core/mod.rs @@ -15,24 +15,36 @@ pub type PublicKey = AffinePoint; #[derive(Clone, Debug, Serialize, Deserialize)] pub struct NSSAUserData { /// Default public accounts. - pub default_pub_account_signing_keys: BTreeMap, + pub default_pub_account_signing_keys: BTreeMap, /// Default private accounts. - pub default_user_private_accounts: - BTreeMap, + pub default_user_private_accounts: BTreeMap, /// Tree of public keys. pub public_key_tree: KeyTreePublic, /// Tree of private keys. pub private_key_tree: KeyTreePrivate, } +/// TODO: eventually, this should have `sign_key: Option` and `pub_key: PublicKey`. +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct PublicBundle { + pub sign_key: nssa::PrivateKey, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct PrivateBundle { + pub key_chain: KeyChain, + pub account: nssa_core::account::Account, +} + impl NSSAUserData { fn valid_public_key_transaction_pairing_check( - accounts_keys_map: &BTreeMap, + accounts_keys_map: &BTreeMap, ) -> bool { let mut check_res = true; - for (account_id, key) in accounts_keys_map { - let expected_account_id = - nssa::AccountId::from(&nssa::PublicKey::new_from_private_key(key)); + for (account_id, public_bundle) in accounts_keys_map { + let expected_account_id = nssa::AccountId::from( + &nssa::PublicKey::new_from_private_key(&public_bundle.sign_key), + ); if &expected_account_id != account_id { println!("{expected_account_id}, {account_id}"); check_res = false; @@ -42,11 +54,11 @@ impl NSSAUserData { } fn valid_private_key_transaction_pairing_check( - accounts_keys_map: &BTreeMap, + accounts_keys_map: &BTreeMap, ) -> bool { let mut check_res = true; - for (account_id, (key, _)) in accounts_keys_map { - let expected_account_id = nssa::AccountId::from(&key.nullifier_public_key); + for (account_id, bundle) in accounts_keys_map { + let expected_account_id = nssa::AccountId::from(&bundle.key_chain.nullifier_public_key); if expected_account_id != *account_id { println!("{expected_account_id}, {account_id}"); check_res = false; @@ -56,11 +68,8 @@ impl NSSAUserData { } pub fn new_with_accounts( - default_accounts_keys: BTreeMap, - default_accounts_key_chains: BTreeMap< - nssa::AccountId, - (KeyChain, nssa_core::account::Account), - >, + default_accounts_keys: BTreeMap, + default_accounts_key_chains: BTreeMap, public_key_tree: KeyTreePublic, private_key_tree: KeyTreePrivate, ) -> Result { @@ -111,6 +120,7 @@ impl NSSAUserData { ) -> Option<&nssa::PrivateKey> { self.default_pub_account_signing_keys .get(&account_id) + .map(|bundle| &bundle.sign_key) .or_else(|| self.public_key_tree.get_node(account_id).map(Into::into)) } @@ -135,23 +145,28 @@ impl NSSAUserData { /// Returns the signing key for public transaction signatures. #[must_use] - pub fn get_private_account( - &self, - account_id: nssa::AccountId, - ) -> Option<&(KeyChain, nssa_core::account::Account)> { + pub fn get_private_account(&self, account_id: nssa::AccountId) -> Option { self.default_user_private_accounts .get(&account_id) - .or_else(|| self.private_key_tree.get_node(account_id).map(Into::into)) + .cloned() + .or_else(|| { + self.private_key_tree + .get_node(account_id) + .map(|child_keys_private| PrivateBundle { + key_chain: child_keys_private.value.0.clone(), + account: child_keys_private.value.1.clone(), + }) + }) } /// Returns the signing key for public transaction signatures. pub fn get_private_account_mut( &mut self, account_id: &nssa::AccountId, - ) -> Option<&mut (KeyChain, nssa_core::account::Account)> { + ) -> Option { // First seek in defaults - if let Some(key) = self.default_user_private_accounts.get_mut(account_id) { - Some(key) + if let Some(bundle) = self.default_user_private_accounts.get(account_id) { + Some(bundle).cloned() // Then seek in tree } else { self.private_key_tree @@ -209,7 +224,10 @@ mod tests { let account_id_private_str = account_id_private.to_string(); println!("{account_id_private_str:#?}"); - let key_chain = &user_data.get_private_account(account_id_private).unwrap().0; + let key_chain = &user_data + .get_private_account(account_id_private) + .unwrap() + .key_chain; println!("{key_chain:#?}"); } } diff --git a/wallet-ffi/src/keys.rs b/wallet-ffi/src/keys.rs index 4eeadd8f..b8c95c55 100644 --- a/wallet-ffi/src/keys.rs +++ b/wallet-ffi/src/keys.rs @@ -116,12 +116,13 @@ pub unsafe extern "C" fn wallet_ffi_get_private_account_keys( let account_id = AccountId::new(unsafe { (*account_id).data }); - let Some((key_chain, _account)) = wallet.storage().user_data.get_private_account(account_id) - else { + let Some(bundle) = wallet.storage().user_data.get_private_account(account_id) else { print_error("Private account not found in wallet"); return WalletFfiError::AccountNotFound; }; + let key_chain = bundle.key_chain; + // NPK is a 32-byte array let npk_bytes = key_chain.nullifier_public_key.0; diff --git a/wallet/src/chain_storage.rs b/wallet/src/chain_storage.rs index 3699609b..c9d07d08 100644 --- a/wallet/src/chain_storage.rs +++ b/wallet/src/chain_storage.rs @@ -7,7 +7,7 @@ use key_protocol::{ key_tree::{KeyTreePrivate, KeyTreePublic, chain_index::ChainIndex}, secret_holders::SeedHolder, }, - key_protocol_core::NSSAUserData, + key_protocol_core::{NSSAUserData, PrivateBundle, PublicBundle}, }; use log::debug; use nssa::program::Program; @@ -74,11 +74,21 @@ impl WalletChainStore { } PersistentAccountData::Preconfigured(acc_data) => match acc_data { InitialAccountData::Public(data) => { - public_init_acc_map.insert(data.account_id, data.pub_sign_key); + public_init_acc_map.insert( + data.account_id, + PublicBundle { + sign_key: data.pub_sign_key, + }, + ); } InitialAccountData::Private(data) => { - private_init_acc_map - .insert(data.account_id, (data.key_chain, data.account)); + private_init_acc_map.insert( + data.account_id, + PrivateBundle { + key_chain: data.key_chain, + account: data.account, + }, + ); } }, } @@ -108,7 +118,12 @@ impl WalletChainStore { for init_acc_data in initial_accounts { match init_acc_data { InitialAccountData::Public(data) => { - public_init_acc_map.insert(data.account_id, data.pub_sign_key); + public_init_acc_map.insert( + data.account_id, + PublicBundle { + sign_key: data.pub_sign_key, + }, + ); } InitialAccountData::Private(data) => { let mut account = data.account; @@ -117,7 +132,13 @@ impl WalletChainStore { // startup. Fix this when program id can be fetched // from the node and queried from the wallet. account.program_owner = Program::authenticated_transfer_program().id(); - private_init_acc_map.insert(data.account_id, (data.key_chain, account)); + private_init_acc_map.insert( + data.account_id, + PrivateBundle { + key_chain: data.key_chain, + account, + }, + ); } } } @@ -178,7 +199,7 @@ impl WalletChainStore { .user_data .default_user_private_accounts .entry(account_id) - .and_modify(|data| data.1 = account.clone()); + .and_modify(|data| data.account = account.clone()); if matches!(entry, Entry::Vacant(_)) { self.user_data diff --git a/wallet/src/cli/account.rs b/wallet/src/cli/account.rs index 86ae7e35..4283561e 100644 --- a/wallet/src/cli/account.rs +++ b/wallet/src/cli/account.rs @@ -156,7 +156,7 @@ impl WalletSubcommand for NewSubcommand { .insert(account_id.to_string(), Label::new(label)); } - let (key, _) = wallet_core + let bundle = wallet_core .storage .user_data .get_private_account(account_id) @@ -165,10 +165,13 @@ impl WalletSubcommand for NewSubcommand { println!( "Generated new account with account_id Private/{account_id} at path {chain_index}", ); - println!("With npk {}", hex::encode(key.nullifier_public_key.0)); + println!( + "With npk {}", + hex::encode(bundle.key_chain.nullifier_public_key.0) + ); println!( "With vpk {}", - hex::encode(key.viewing_public_key.to_bytes()) + hex::encode(bundle.key_chain.viewing_public_key.to_bytes()) ); wallet_core.store_persistent_data().await?; @@ -229,14 +232,20 @@ impl WalletSubcommand for AccountSubcommand { println!("pk {}", hex::encode(public_key.value())); } AccountPrivacyKind::Private => { - let (key, _) = wallet_core + let bundle = wallet_core .storage .user_data .get_private_account(account_id) .context("Private account not found in storage")?; - println!("npk {}", hex::encode(key.nullifier_public_key.0)); - println!("vpk {}", hex::encode(key.viewing_public_key.to_bytes())); + println!( + "npk {}", + hex::encode(bundle.key_chain.nullifier_public_key.0) + ); + println!( + "vpk {}", + hex::encode(bundle.key_chain.viewing_public_key.to_bytes()) + ); } } Ok(()) diff --git a/wallet/src/helperfunctions.rs b/wallet/src/helperfunctions.rs index 37a27409..6cfe00c8 100644 --- a/wallet/src/helperfunctions.rs +++ b/wallet/src/helperfunctions.rs @@ -178,22 +178,22 @@ pub fn produce_data_for_storage( } } - for (account_id, key) in &user_data.default_pub_account_signing_keys { + for (account_id, bundle) in &user_data.default_pub_account_signing_keys { vec_for_storage.push( InitialAccountData::Public(PublicAccountPrivateInitialData { account_id: *account_id, - pub_sign_key: key.clone(), + pub_sign_key: bundle.sign_key.clone(), }) .into(), ); } - for (account_id, (key_chain, account)) in &user_data.default_user_private_accounts { + for (account_id, bundle) in &user_data.default_user_private_accounts { vec_for_storage.push( InitialAccountData::Private(Box::new(PrivateAccountPrivateInitialData { account_id: *account_id, - account: account.clone(), - key_chain: key_chain.clone(), + account: bundle.account.clone(), + key_chain: bundle.key_chain.clone(), })) .into(), ); diff --git a/wallet/src/lib.rs b/wallet/src/lib.rs index 63ea8611..f4fdf2da 100644 --- a/wallet/src/lib.rs +++ b/wallet/src/lib.rs @@ -15,7 +15,10 @@ use bip39::Mnemonic; use chain_storage::WalletChainStore; use common::{HashType, transaction::NSSATransaction}; use config::WalletConfig; -use key_protocol::key_management::key_tree::{chain_index::ChainIndex, traits::KeyNode as _}; +use key_protocol::{ + key_management::key_tree::{chain_index::ChainIndex, traits::KeyNode as _}, + key_protocol_core::PrivateBundle, +}; use log::info; use nssa::{ Account, AccountId, PrivacyPreservingTransaction, @@ -295,13 +298,16 @@ impl WalletCore { self.storage .user_data .get_private_account(account_id) - .map(|value| value.1.clone()) + .map(|bundle| bundle.account) } #[must_use] pub fn get_private_account_commitment(&self, account_id: AccountId) -> Option { - let (keys, account) = self.storage.user_data.get_private_account(account_id)?; - Some(Commitment::new(&keys.nullifier_public_key, account)) + let bundle = self.storage.user_data.get_private_account(account_id)?; + Some(Commitment::new( + &bundle.key_chain.nullifier_public_key, + &bundle.account, + )) } /// Poll transactions. @@ -484,7 +490,15 @@ impl WalletCore { .user_data .default_user_private_accounts .iter() - .map(|(acc_account_id, (key_chain, _))| (*acc_account_id, key_chain, None)) + .map( + |( + acc_account_id, + PrivateBundle { + key_chain, + account: _, + }, + )| (*acc_account_id, key_chain, None), + ) .chain(self.storage.user_data.private_key_tree.key_map.iter().map( |(chain_index, keys_node)| { ( diff --git a/wallet/src/privacy_preserving_tx.rs b/wallet/src/privacy_preserving_tx.rs index 04056111..e70c928e 100644 --- a/wallet/src/privacy_preserving_tx.rs +++ b/wallet/src/privacy_preserving_tx.rs @@ -202,15 +202,13 @@ async fn private_acc_preparation( wallet: &WalletCore, account_id: AccountId, ) -> Result { - let Some((from_keys, from_acc)) = wallet - .storage - .user_data - .get_private_account(account_id) - .cloned() - else { + let Some(from_bundle) = wallet.storage.user_data.get_private_account(account_id) else { return Err(ExecutionFailureKind::KeyNotFoundError); }; + let from_keys = from_bundle.key_chain; + let from_acc = from_bundle.account; + let nsk = from_keys.private_key_holder.nullifier_secret_key; let from_npk = from_keys.nullifier_public_key;