diff --git a/integration_tests/tests/auth_transfer/private.rs b/integration_tests/tests/auth_transfer/private.rs index 59b4719a..22396ebc 100644 --- a/integration_tests/tests/auth_transfer/private.rs +++ b/integration_tests/tests/auth_transfer/private.rs @@ -169,7 +169,8 @@ async fn private_transfer_to_owned_account_using_claiming_path() -> Result<()> { .storage() .user_data .get_private_account(to_account_id) - .cloned() + .unwrap() + .key_chain .context("Failed to get private account")?; // Send to this account using claiming path (using npk and vpk instead of account ID) @@ -329,7 +330,8 @@ async fn private_transfer_to_owned_account_continuous_run_path() -> Result<()> { .storage() .user_data .get_private_account(to_account_id) - .cloned() + .unwrap() + .key_chain .context("Failed to get private account")?; // Send transfer using nullifier and viewing public keys diff --git a/integration_tests/tests/keys_restoration.rs b/integration_tests/tests/keys_restoration.rs index cdbe2e6b..7706e108 100644 --- a/integration_tests/tests/keys_restoration.rs +++ b/integration_tests/tests/keys_restoration.rs @@ -64,7 +64,8 @@ async fn sync_private_account_with_non_zero_chain_index() -> Result<()> { .storage() .user_data .get_private_account(to_account_id) - .cloned() + .unwrap() + .key_chain .context("Failed to get private account")?; // Send to this account using claiming path (using npk and vpk instead of account ID) diff --git a/integration_tests/tests/token.rs b/integration_tests/tests/token.rs index b638b6c9..683f7ebe 100644 --- a/integration_tests/tests/token.rs +++ b/integration_tests/tests/token.rs @@ -1102,7 +1102,8 @@ async fn token_claiming_path_with_private_accounts() -> Result<()> { .storage() .user_data .get_private_account(recipient_account_id) - .cloned() + .unwrap() + .key_chain .context("Failed to get private account keys")?; // Mint using claiming path (foreign account) 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..c08f9e70 100644 --- a/key_protocol/src/key_management/key_tree/keys_private.rs +++ b/key_protocol/src/key_management/key_tree/keys_private.rs @@ -2,11 +2,11 @@ use k256::{Scalar, elliptic_curve::PrimeField as _}; use nssa_core::{NullifierPublicKey, encryption::ViewingPublicKey}; use serde::{Deserialize, Serialize}; -use crate::key_management::{ +use crate::{key_management::{ KeyChain, key_tree::traits::KeyNode, secret_holders::{PrivateKeyHolder, SecretSpendingKey}, -}; +}, key_protocol_core::PrivateBundle}; #[derive(Debug, Serialize, Deserialize, Clone)] pub struct ChildKeysPrivate { @@ -129,9 +129,10 @@ 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 + PrivateBundle { 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..6403ec9a 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,10 +68,10 @@ impl NSSAUserData { } pub fn new_with_accounts( - default_accounts_keys: BTreeMap, + default_accounts_keys: BTreeMap, default_accounts_key_chains: BTreeMap< nssa::AccountId, - (KeyChain, nssa_core::account::Account), + PrivateBundle, >, public_key_tree: KeyTreePublic, private_key_tree: KeyTreePrivate, @@ -111,6 +123,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)) } @@ -134,24 +147,33 @@ 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)> { + #[must_use] //Marvin: double check TODO + 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)) 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. + /// TODO: fix this comment (Marvin) 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 +231,7 @@ 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..16e152b5 100644 --- a/wallet-ffi/src/keys.rs +++ b/wallet-ffi/src/keys.rs @@ -116,12 +116,14 @@ 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) + 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..c5fc7d80 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,16 @@ 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)); + .insert(data.account_id, PrivateBundle {key_chain: data.key_chain, account: data.account}); } }, } @@ -108,7 +113,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 +127,7 @@ 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: account}); } } } @@ -178,7 +188,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 2a8ed2c7..1bb6ec74 100644 --- a/wallet/src/cli/account.rs +++ b/wallet/src/cli/account.rs @@ -137,7 +137,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) @@ -146,10 +146,10 @@ 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?; @@ -203,14 +203,14 @@ 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 d82dedaf..a790b77a 100644 --- a/wallet/src/helperfunctions.rs +++ b/wallet/src/helperfunctions.rs @@ -118,22 +118,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..0a01f185 100644 --- a/wallet/src/lib.rs +++ b/wallet/src/lib.rs @@ -15,7 +15,7 @@ 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 +295,13 @@ impl WalletCore { self.storage .user_data .get_private_account(account_id) - .map(|value| value.1.clone()) + .map(|bundle| bundle.account.clone()) } #[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 +484,7 @@ 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..5b7ce373 100644 --- a/wallet/src/privacy_preserving_tx.rs +++ b/wallet/src/privacy_preserving_tx.rs @@ -202,15 +202,18 @@ async fn private_acc_preparation( wallet: &WalletCore, account_id: AccountId, ) -> Result { - let Some((from_keys, from_acc)) = wallet + let Some(from_bundle) = wallet .storage .user_data .get_private_account(account_id) - .cloned() + .clone() 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;