diff --git a/key_protocol/src/key_management/key_tree/mod.rs b/key_protocol/src/key_management/key_tree/mod.rs index 3f847ebd..01a4922e 100644 --- a/key_protocol/src/key_management/key_tree/mod.rs +++ b/key_protocol/src/key_management/key_tree/mod.rs @@ -255,6 +255,28 @@ impl KeyTree { } impl KeyTree { + /// Generate a new private key node, registering the identifier=0 account_id immediately. + pub fn generate_new_private_node( + &mut self, + parent_cci: &ChainIndex, + ) -> Option<(nssa::AccountId, ChainIndex)> { + let cci = self.generate_new_node(parent_cci)?; + let node = self.key_map.get(&cci)?; + let account_id = nssa::AccountId::from((&node.value.0.nullifier_public_key, 0_u128)); + self.account_id_map.insert(account_id, cci.clone()); + Some((account_id, cci)) + } + + /// Generate a new private key node using layered placement, registering the identifier=0 + /// account_id immediately. + pub fn generate_new_private_node_layered(&mut self) -> Option<(nssa::AccountId, ChainIndex)> { + let cci = self.generate_new_node_layered()?; + let node = self.key_map.get(&cci)?; + let account_id = nssa::AccountId::from((&node.value.0.nullifier_public_key, 0_u128)); + self.account_id_map.insert(account_id, cci.clone()); + Some((account_id, cci)) + } + /// Cleanup of non-initialized accounts in a private tree. /// /// If account has no synced entries, removes it, stops at first initialized account. diff --git a/key_protocol/src/key_protocol_core/mod.rs b/key_protocol/src/key_protocol_core/mod.rs index 449196ea..b756c8f9 100644 --- a/key_protocol/src/key_protocol_core/mod.rs +++ b/key_protocol/src/key_protocol_core/mod.rs @@ -123,19 +123,19 @@ impl NSSAUserData { /// Generated new private key for privacy preserving transactions. /// - /// Returns the `ChainIndex` of the new node. + /// Returns the `AccountId` (for identifier=0) and `ChainIndex` of the new node. pub fn generate_new_privacy_preserving_transaction_key_chain( &mut self, parent_cci: Option, - ) -> ChainIndex { + ) -> (nssa::AccountId, ChainIndex) { match parent_cci { Some(parent_cci) => self .private_key_tree - .generate_new_node(&parent_cci) + .generate_new_private_node(&parent_cci) .expect("Parent must be present in a tree"), None => self .private_key_tree - .generate_new_node_layered() + .generate_new_private_node_layered() .expect("Search for new node slot failed"), } } @@ -211,13 +211,16 @@ mod tests { fn new_account() { let mut user_data = NSSAUserData::default(); - let chain_index = user_data + let (account_id, chain_index) = user_data .generate_new_privacy_preserving_transaction_key_chain(Some(ChainIndex::root())); let is_key_chain_generated = user_data.private_key_tree.key_map.contains_key(&chain_index); - assert!(is_key_chain_generated); + let is_account_id_registered = + user_data.private_key_tree.account_id_map.contains_key(&account_id); + assert!(is_account_id_registered); + let key_chain = &user_data.private_key_tree.key_map[&chain_index].value.0; println!("{key_chain:#?}"); } diff --git a/wallet-ffi/src/account.rs b/wallet-ffi/src/account.rs index 5bbf8c53..49f6a8de 100644 --- a/wallet-ffi/src/account.rs +++ b/wallet-ffi/src/account.rs @@ -98,16 +98,7 @@ pub unsafe extern "C" fn wallet_ffi_create_account_private( } }; - let chain_index = wallet.create_new_account_private(None); - - let node = wallet - .storage() - .user_data - .private_key_tree - .key_map - .get(&chain_index) - .expect("Node was just inserted"); - let account_id = AccountId::from((&node.value.0.nullifier_public_key, 0_u128)); + let (account_id, _chain_index) = wallet.create_new_account_private(None); unsafe { (*out_account_id).data = *account_id.value(); diff --git a/wallet/src/cli/account.rs b/wallet/src/cli/account.rs index 577f8042..c3b9183f 100644 --- a/wallet/src/cli/account.rs +++ b/wallet/src/cli/account.rs @@ -147,7 +147,7 @@ impl WalletSubcommand for NewSubcommand { anyhow::bail!("Label '{label}' is already in use by another account"); } - let chain_index = wallet_core.create_new_account_private(cci); + let (account_id, chain_index) = wallet_core.create_new_account_private(cci); let node = wallet_core .storage @@ -157,7 +157,6 @@ impl WalletSubcommand for NewSubcommand { .get(&chain_index) .expect("Node was just inserted"); let key = &node.value.0; - let account_id = AccountId::from((&key.nullifier_public_key, 0_u128)); if let Some(label) = label { wallet_core diff --git a/wallet/src/lib.rs b/wallet/src/lib.rs index b2458390..3301b513 100644 --- a/wallet/src/lib.rs +++ b/wallet/src/lib.rs @@ -259,7 +259,7 @@ impl WalletCore { pub fn create_new_account_private( &mut self, chain_index: Option, - ) -> ChainIndex { + ) -> (AccountId, ChainIndex) { self.storage .user_data .generate_new_privacy_preserving_transaction_key_chain(chain_index)