From 6cbc5028cfa786a041e27177e65d7d4596778be6 Mon Sep 17 00:00:00 2001 From: Pravdyvy Date: Tue, 11 Nov 2025 17:25:08 +0200 Subject: [PATCH 01/10] feat: tree construction --- .../src/key_management/key_tree/mod.rs | 17 +++++++++++++++++ wallet/src/lib.rs | 19 +++++++++++++++++++ wallet/src/main.rs | 5 ++++- 3 files changed, 40 insertions(+), 1 deletion(-) diff --git a/key_protocol/src/key_management/key_tree/mod.rs b/key_protocol/src/key_management/key_tree/mod.rs index dcc027b..7a25c81 100644 --- a/key_protocol/src/key_management/key_tree/mod.rs +++ b/key_protocol/src/key_management/key_tree/mod.rs @@ -127,6 +127,23 @@ impl KeyTree { self.addr_map.insert(addr, chain_index.clone()); self.key_map.insert(chain_index, node); } + + pub fn generate_tree_for_depth(&mut self, depth: u32) { + let mut id_stack = vec![ChainIndex::root()]; + + while !id_stack.is_empty() { + let curr_id = id_stack.pop().unwrap(); + + self.generate_new_node(curr_id.clone()); + + let mut next_id = curr_id.n_th_child(0); + + while (next_id.chain().iter().sum::()) < depth - 1 { + id_stack.push(next_id.clone()); + next_id = next_id.next_in_line(); + } + } + } } #[cfg(test)] diff --git a/wallet/src/lib.rs b/wallet/src/lib.rs index 3b494a9..c39669f 100644 --- a/wallet/src/lib.rs +++ b/wallet/src/lib.rs @@ -259,6 +259,13 @@ pub enum OverCommand { #[arg(short, long)] password: String, }, + /// !!!WARNING!!! will rewrite current storage + RestoreKeys { + #[arg(short, long)] + password: String, + #[arg(short, long)] + depth: u32, + } } ///To execute commands, env var NSSA_WALLET_HOME_DIR must be set into directory with config @@ -511,3 +518,15 @@ pub async fn execute_setup(password: String) -> Result<()> { Ok(()) } + +pub async fn execute_keys_restoration(password: String, depth: u32) -> Result<()> { + let config = fetch_config().await?; + let mut wallet_core = WalletCore::start_from_config_new_storage(config.clone(), password.clone()).await?; + + wallet_core.storage.user_data.public_key_tree.generate_tree_for_depth(depth); + wallet_core.storage.user_data.private_key_tree.generate_tree_for_depth(depth); + + wallet_core.store_persistent_data().await?; + + Ok(()) +} diff --git a/wallet/src/main.rs b/wallet/src/main.rs index 1fe52b3..986b27e 100644 --- a/wallet/src/main.rs +++ b/wallet/src/main.rs @@ -1,7 +1,7 @@ use anyhow::Result; use clap::{CommandFactory, Parser}; use tokio::runtime::Builder; -use wallet::{Args, OverCommand, execute_continious_run, execute_setup, execute_subcommand}; +use wallet::{Args, OverCommand, execute_continious_run, execute_keys_restoration, execute_setup, execute_subcommand}; pub const NUM_THREADS: usize = 2; @@ -25,6 +25,9 @@ fn main() -> Result<()> { OverCommand::Setup { password } => { execute_setup(password).await.unwrap(); } + OverCommand::RestoreKeys { password, depth } => { + execute_keys_restoration(password, depth).await.unwrap(); + } } } else if args.continious_run { execute_continious_run().await.unwrap(); From e92ad2132f769f60cd4ee7b885df0bf67fc2874e Mon Sep 17 00:00:00 2001 From: Pravdyvy Date: Wed, 12 Nov 2025 08:26:25 +0200 Subject: [PATCH 02/10] feat: keys restoration from mnemonic --- integration_tests/src/test_suite_map.rs | 131 ++++++++++++++++++ .../key_management/key_tree/chain_index.rs | 19 ++- .../src/key_management/key_tree/mod.rs | 73 +++++++++- wallet/src/lib.rs | 57 +++++++- wallet/src/main.rs | 5 +- 5 files changed, 271 insertions(+), 14 deletions(-) diff --git a/integration_tests/src/test_suite_map.rs b/integration_tests/src/test_suite_map.rs index ff4ce84..4399407 100644 --- a/integration_tests/src/test_suite_map.rs +++ b/integration_tests/src/test_suite_map.rs @@ -3,6 +3,7 @@ use std::{ collections::HashMap, path::PathBuf, pin::Pin, + str::FromStr, time::{Duration, Instant}, }; @@ -1647,6 +1648,136 @@ pub fn prepare_function_map() -> HashMap { info!("Success!"); } + #[nssa_integration_test] + pub async fn test_keys_restoration() { + info!("########## test_keys_restoration ##########"); + let from: Address = ACC_SENDER_PRIVATE.parse().unwrap(); + + let command = Command::Account(AccountSubcommand::New(NewSubcommand::Private { + cci: ChainIndex::root(), + })); + + let sub_ret = wallet::execute_subcommand(command).await.unwrap(); + let SubcommandReturnValue::RegisterAccount { addr: to_addr1 } = sub_ret else { + panic!("FAILED TO REGISTER ACCOUNT"); + }; + + let command = Command::Account(AccountSubcommand::New(NewSubcommand::Private { + cci: ChainIndex::from_str("/0").unwrap(), + })); + + let sub_ret = wallet::execute_subcommand(command).await.unwrap(); + let SubcommandReturnValue::RegisterAccount { addr: to_addr2 } = sub_ret else { + panic!("FAILED TO REGISTER ACCOUNT"); + }; + + let command = Command::AuthTransfer(AuthTransferSubcommand::Send { + from: make_private_account_input_from_str(&from.to_string()), + to: Some(make_private_account_input_from_str(&to_addr1.to_string())), + to_npk: None, + to_ipk: None, + amount: 100, + }); + + wallet::execute_subcommand(command).await.unwrap(); + + let command = Command::AuthTransfer(AuthTransferSubcommand::Send { + from: make_private_account_input_from_str(&from.to_string()), + to: Some(make_private_account_input_from_str(&to_addr2.to_string())), + to_npk: None, + to_ipk: None, + amount: 100, + }); + + wallet::execute_subcommand(command).await.unwrap(); + + let from: Address = ACC_SENDER.parse().unwrap(); + + let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public { + cci: ChainIndex::root(), + })); + + let sub_ret = wallet::execute_subcommand(command).await.unwrap(); + let SubcommandReturnValue::RegisterAccount { addr: to_addr3 } = sub_ret else { + panic!("FAILED TO REGISTER ACCOUNT"); + }; + + let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public { + cci: ChainIndex::from_str("/0").unwrap(), + })); + + let sub_ret = wallet::execute_subcommand(command).await.unwrap(); + let SubcommandReturnValue::RegisterAccount { addr: to_addr4 } = sub_ret else { + panic!("FAILED TO REGISTER ACCOUNT"); + }; + + let command = Command::AuthTransfer(AuthTransferSubcommand::Send { + from: make_public_account_input_from_str(&from.to_string()), + to: Some(make_public_account_input_from_str(&to_addr3.to_string())), + to_npk: None, + to_ipk: None, + amount: 100, + }); + + wallet::execute_subcommand(command).await.unwrap(); + + let command = Command::AuthTransfer(AuthTransferSubcommand::Send { + from: make_public_account_input_from_str(&from.to_string()), + to: Some(make_public_account_input_from_str(&to_addr4.to_string())), + to_npk: None, + to_ipk: None, + amount: 100, + }); + + wallet::execute_subcommand(command).await.unwrap(); + + info!("########## PREPARATION END ##########"); + + wallet::execute_keys_restoration("test_pass".to_string(), 10) + .await + .unwrap(); + + let wallet_config = fetch_config().await.unwrap(); + let wallet_storage = WalletCore::start_from_config_update_chain(wallet_config.clone()) + .await + .unwrap(); + + assert!( + wallet_storage + .storage + .user_data + .private_key_tree + .get_node(to_addr1) + .is_some() + ); + assert!( + wallet_storage + .storage + .user_data + .private_key_tree + .get_node(to_addr2) + .is_some() + ); + assert!( + wallet_storage + .storage + .user_data + .public_key_tree + .get_node(to_addr3) + .is_some() + ); + assert!( + wallet_storage + .storage + .user_data + .public_key_tree + .get_node(to_addr4) + .is_some() + ); + + info!("Success!"); + } + println!("{function_map:#?}"); function_map diff --git a/key_protocol/src/key_management/key_tree/chain_index.rs b/key_protocol/src/key_management/key_tree/chain_index.rs index e22abf0..b7190d1 100644 --- a/key_protocol/src/key_management/key_tree/chain_index.rs +++ b/key_protocol/src/key_management/key_tree/chain_index.rs @@ -42,10 +42,13 @@ impl FromStr for ChainIndex { impl Display for ChainIndex { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "/")?; - for cci in &self.0[..(self.0.len() - 1)] { - write!(f, "{cci}/")?; + if *self != Self::root() { + for cci in &self.0[..(self.0.len() - 1)] { + write!(f, "{cci}/")?; + } + write!(f, "{}", self.0.last().unwrap())?; } - write!(f, "{}", self.0.last().unwrap()) + Ok(()) } } @@ -74,6 +77,16 @@ impl ChainIndex { ChainIndex(chain) } + + pub fn depth(&self) -> u32 { + let mut res = 0; + + for cci in &self.0 { + res += cci + 1; + } + + res + } } #[cfg(test)] diff --git a/key_protocol/src/key_management/key_tree/mod.rs b/key_protocol/src/key_management/key_tree/mod.rs index 7a25c81..d6600fe 100644 --- a/key_protocol/src/key_management/key_tree/mod.rs +++ b/key_protocol/src/key_management/key_tree/mod.rs @@ -1,5 +1,10 @@ -use std::collections::{BTreeMap, HashMap}; +use std::{ + collections::{BTreeMap, HashMap}, + sync::Arc, +}; +use anyhow::Result; +use common::sequencer_client::SequencerClient; use serde::{Deserialize, Serialize}; use crate::key_management::{ @@ -128,24 +133,80 @@ impl KeyTree { self.key_map.insert(chain_index, node); } + pub fn remove(&mut self, addr: nssa::Address) -> Option { + let chain_index = self.addr_map.remove(&addr).unwrap(); + self.key_map.remove(&chain_index) + } + pub fn generate_tree_for_depth(&mut self, depth: u32) { let mut id_stack = vec![ChainIndex::root()]; - while !id_stack.is_empty() { - let curr_id = id_stack.pop().unwrap(); - + while let Some(curr_id) = id_stack.pop() { self.generate_new_node(curr_id.clone()); let mut next_id = curr_id.n_th_child(0); - while (next_id.chain().iter().sum::()) < depth - 1 { + while (next_id.depth()) < depth - 1 { id_stack.push(next_id.clone()); next_id = next_id.next_in_line(); - } + } } } } +impl KeyTree { + pub fn cleanup_tree_for_depth(&mut self, depth: u32) { + let mut id_stack = vec![ChainIndex::root()]; + + while let Some(curr_id) = id_stack.pop() { + if let Some(node) = self.key_map.get(&curr_id) + && node.value.1 == nssa::Account::default() + && curr_id != ChainIndex::root() + { + let addr = node.address(); + self.remove(addr); + } + + let mut next_id = curr_id.n_th_child(0); + + while (next_id.depth()) < depth - 1 { + id_stack.push(next_id.clone()); + next_id = next_id.next_in_line(); + } + } + } +} + +impl KeyTree { + pub async fn cleanup_tree_for_depth( + &mut self, + depth: u32, + client: Arc, + ) -> Result<()> { + let mut id_stack = vec![ChainIndex::root()]; + + while let Some(curr_id) = id_stack.pop() { + if let Some(node) = self.key_map.get(&curr_id) { + let address = node.address(); + let node_acc = client.get_account(address.to_string()).await?.account; + + if node_acc == nssa::Account::default() && curr_id != ChainIndex::root() { + self.remove(address); + } + } + + let mut next_id = curr_id.n_th_child(0); + + while (next_id.depth()) < depth - 1 { + id_stack.push(next_id.clone()); + next_id = next_id.next_in_line(); + } + } + + Ok(()) + } +} + #[cfg(test)] mod tests { use std::str::FromStr; diff --git a/wallet/src/lib.rs b/wallet/src/lib.rs index c39669f..8949e7b 100644 --- a/wallet/src/lib.rs +++ b/wallet/src/lib.rs @@ -265,7 +265,7 @@ pub enum OverCommand { password: String, #[arg(short, long)] depth: u32, - } + }, } ///To execute commands, env var NSSA_WALLET_HOME_DIR must be set into directory with config @@ -521,10 +521,59 @@ pub async fn execute_setup(password: String) -> Result<()> { pub async fn execute_keys_restoration(password: String, depth: u32) -> Result<()> { let config = fetch_config().await?; - let mut wallet_core = WalletCore::start_from_config_new_storage(config.clone(), password.clone()).await?; + let mut wallet_core = + WalletCore::start_from_config_new_storage(config.clone(), password.clone()).await?; - wallet_core.storage.user_data.public_key_tree.generate_tree_for_depth(depth); - wallet_core.storage.user_data.private_key_tree.generate_tree_for_depth(depth); + wallet_core + .storage + .user_data + .public_key_tree + .generate_tree_for_depth(depth); + + println!("Public tree generated"); + + wallet_core + .storage + .user_data + .private_key_tree + .generate_tree_for_depth(depth); + + println!("Private tree generated"); + + wallet_core + .storage + .user_data + .public_key_tree + .cleanup_tree_for_depth(depth, wallet_core.sequencer_client.clone()) + .await?; + + println!("Public tree cleaned up"); + + let last_block = wallet_core + .sequencer_client + .get_last_block() + .await? + .last_block; + + println!("Last block is {last_block}"); + + parse_block_range( + 1, + last_block, + wallet_core.sequencer_client.clone(), + &mut wallet_core, + ) + .await?; + + println!("Private tree clean up start"); + + wallet_core + .storage + .user_data + .private_key_tree + .cleanup_tree_for_depth(depth); + + println!("Private tree cleaned up"); wallet_core.store_persistent_data().await?; diff --git a/wallet/src/main.rs b/wallet/src/main.rs index 986b27e..1a8b35a 100644 --- a/wallet/src/main.rs +++ b/wallet/src/main.rs @@ -1,7 +1,10 @@ use anyhow::Result; use clap::{CommandFactory, Parser}; use tokio::runtime::Builder; -use wallet::{Args, OverCommand, execute_continious_run, execute_keys_restoration, execute_setup, execute_subcommand}; +use wallet::{ + Args, OverCommand, execute_continious_run, execute_keys_restoration, execute_setup, + execute_subcommand, +}; pub const NUM_THREADS: usize = 2; From 5affa4f9fdd03c9dacef2e734342b91f76867aa3 Mon Sep 17 00:00:00 2001 From: Pravdyvy Date: Wed, 3 Dec 2025 07:05:23 +0200 Subject: [PATCH 03/10] fix: suggestions fix --- integration_tests/src/test_suite_map.rs | 127 +++++++++++++----- .../src/key_management/key_tree/mod.rs | 20 +++ 2 files changed, 114 insertions(+), 33 deletions(-) diff --git a/integration_tests/src/test_suite_map.rs b/integration_tests/src/test_suite_map.rs index 9f61810..6af300a 100644 --- a/integration_tests/src/test_suite_map.rs +++ b/integration_tests/src/test_suite_map.rs @@ -1722,7 +1722,7 @@ pub fn prepare_function_map() -> HashMap { )), to_npk: None, to_ipk: None, - amount: 100, + amount: 101, }); wallet::execute_subcommand(command).await.unwrap(); @@ -1760,7 +1760,7 @@ pub fn prepare_function_map() -> HashMap { )), to_npk: None, to_ipk: None, - amount: 100, + amount: 102, }); wallet::execute_subcommand(command).await.unwrap(); @@ -1772,7 +1772,7 @@ pub fn prepare_function_map() -> HashMap { )), to_npk: None, to_ipk: None, - amount: 100, + amount: 103, }); wallet::execute_subcommand(command).await.unwrap(); @@ -1788,39 +1788,100 @@ pub fn prepare_function_map() -> HashMap { .await .unwrap(); - assert!( - wallet_storage - .storage - .user_data - .private_key_tree - .get_node(to_account_id1) - .is_some() + let acc1 = wallet_storage + .storage + .user_data + .private_key_tree + .get_node(to_account_id1) + .expect("Acc 1 should be restored"); + + let acc2 = wallet_storage + .storage + .user_data + .private_key_tree + .get_node(to_account_id2) + .expect("Acc 2 should be restored"); + + let _ = wallet_storage + .storage + .user_data + .public_key_tree + .get_node(to_account_id3) + .expect("Acc 3 should be restored"); + + let _ = wallet_storage + .storage + .user_data + .public_key_tree + .get_node(to_account_id4) + .expect("Acc 4 should be restored"); + + assert_eq!( + acc1.value.1.program_owner, + Program::authenticated_transfer_program().id() ); - assert!( - wallet_storage - .storage - .user_data - .private_key_tree - .get_node(to_account_id2) - .is_some() - ); - assert!( - wallet_storage - .storage - .user_data - .public_key_tree - .get_node(to_account_id3) - .is_some() - ); - assert!( - wallet_storage - .storage - .user_data - .public_key_tree - .get_node(to_account_id4) - .is_some() + assert_eq!( + acc2.value.1.program_owner, + Program::authenticated_transfer_program().id() ); + assert_eq!(acc1.value.1.balance, 100); + assert_eq!(acc2.value.1.balance, 101); + + info!("########## TREE CHECKS END ##########"); + + let command = Command::AuthTransfer(AuthTransferSubcommand::Send { + from: make_private_account_input_from_str(&to_account_id1.to_string()), + to: Some(make_private_account_input_from_str( + &to_account_id2.to_string(), + )), + to_npk: None, + to_ipk: None, + amount: 10, + }); + + wallet::execute_subcommand(command).await.unwrap(); + + let command = Command::AuthTransfer(AuthTransferSubcommand::Send { + from: make_public_account_input_from_str(&to_account_id3.to_string()), + to: Some(make_public_account_input_from_str( + &to_account_id4.to_string(), + )), + to_npk: None, + to_ipk: None, + amount: 11, + }); + + wallet::execute_subcommand(command).await.unwrap(); + + let wallet_config = fetch_config().await.unwrap(); + let seq_client = SequencerClient::new(wallet_config.sequencer_addr.clone()).unwrap(); + let wallet_storage = WalletCore::start_from_config_update_chain(wallet_config.clone()) + .await + .unwrap(); + + let comm1 = wallet_storage + .get_private_account_commitment(&to_account_id1) + .expect("Acc 1 commitment should exist"); + let comm2 = wallet_storage + .get_private_account_commitment(&to_account_id2) + .expect("Acc 2 commitment should exist"); + + assert!(verify_commitment_is_in_state(comm1, &seq_client).await); + assert!(verify_commitment_is_in_state(comm2, &seq_client).await); + + let acc3 = seq_client + .get_account_balance(to_account_id3.to_string()) + .await + .expect("Acc 3 must be present in public state"); + let acc4 = seq_client + .get_account_balance(to_account_id4.to_string()) + .await + .expect("Acc 4 must be present in public state"); + + assert_eq!(acc3.balance, 91); + assert_eq!(acc4.balance, 114); + info!("Success!"); } diff --git a/key_protocol/src/key_management/key_tree/mod.rs b/key_protocol/src/key_management/key_tree/mod.rs index 7f91143..3d55516 100644 --- a/key_protocol/src/key_management/key_tree/mod.rs +++ b/key_protocol/src/key_management/key_tree/mod.rs @@ -140,6 +140,12 @@ impl KeyTree { self.key_map.remove(&chain_index) } + /// Populates tree with children. + /// + /// For given `depth` adds children to a tree such that their `ChainIndex::depth(&self) < + /// depth`. + /// + /// Tree must be empty before start pub fn generate_tree_for_depth(&mut self, depth: u32) { let mut id_stack = vec![ChainIndex::root()]; @@ -157,6 +163,14 @@ impl KeyTree { } impl KeyTree { + /// Cleanup of all non-initialized accounts in a private tree + /// + /// For given `depth` checks children to a tree such that their `ChainIndex::depth(&self) < + /// depth`. + /// + /// If account is default, removes them. + /// + /// Chain must be parsed for accounts beforehand pub fn cleanup_tree_for_depth(&mut self, depth: u32) { let mut id_stack = vec![ChainIndex::root()]; @@ -180,6 +194,12 @@ impl KeyTree { } impl KeyTree { + /// Cleanup of all non-initialized accounts in a public tree + /// + /// For given `depth` checks children to a tree such that their `ChainIndex::depth(&self) < + /// depth`. + /// + /// If account is default, removes them. pub async fn cleanup_tree_for_depth( &mut self, depth: u32, From 2453ae095f9f5e650bc6a6e2df986dca70866e03 Mon Sep 17 00:00:00 2001 From: Pravdyvy Date: Thu, 4 Dec 2025 14:31:52 +0200 Subject: [PATCH 04/10] fix: cleanup fixing --- .../src/key_management/key_tree/mod.rs | 160 +++++++++++++++--- key_protocol/src/key_protocol_core/mod.rs | 10 +- 2 files changed, 141 insertions(+), 29 deletions(-) diff --git a/key_protocol/src/key_management/key_tree/mod.rs b/key_protocol/src/key_management/key_tree/mod.rs index 3d55516..7b85d8e 100644 --- a/key_protocol/src/key_management/key_tree/mod.rs +++ b/key_protocol/src/key_management/key_tree/mod.rs @@ -95,27 +95,29 @@ impl KeyTree { right = (left_border + right) / 2; } (None, Some(_)) => { - unreachable!(); + break Some(right); } } } } - pub fn generate_new_node(&mut self, parent_cci: ChainIndex) -> Option { - let father_keys = self.key_map.get(&parent_cci)?; + pub fn generate_new_node( + &mut self, + parent_cci: &ChainIndex, + ) -> Option<(nssa::AccountId, ChainIndex)> { + let father_keys = self.key_map.get(parent_cci)?; let next_child_id = self - .find_next_last_child_of_id(&parent_cci) + .find_next_last_child_of_id(parent_cci) .expect("Can be None only if parent is not present"); let next_cci = parent_cci.nth_child(next_child_id); let child_keys = father_keys.nth_child(next_child_id); - let account_id = child_keys.account_id(); self.key_map.insert(next_cci.clone(), child_keys); - self.account_id_map.insert(account_id, next_cci); + self.account_id_map.insert(account_id, next_cci.clone()); - Some(account_id) + Some((account_id, next_cci)) } pub fn get_node(&self, account_id: nssa::AccountId) -> Option<&N> { @@ -150,11 +152,10 @@ impl KeyTree { let mut id_stack = vec![ChainIndex::root()]; while let Some(curr_id) = id_stack.pop() { - self.generate_new_node(curr_id.clone()); - let mut next_id = curr_id.nth_child(0); - while (next_id.depth()) < depth - 1 { + while (next_id.depth()) < depth { + self.generate_new_node(&curr_id); id_stack.push(next_id.clone()); next_id = next_id.next_in_line(); } @@ -185,7 +186,7 @@ impl KeyTree { let mut next_id = curr_id.nth_child(0); - while (next_id.depth()) < depth - 1 { + while (next_id.depth()) < depth { id_stack.push(next_id.clone()); next_id = next_id.next_in_line(); } @@ -219,7 +220,7 @@ impl KeyTree { let mut next_id = curr_id.nth_child(0); - while (next_id.depth()) < depth - 1 { + while (next_id.depth()) < depth { id_stack.push(next_id.clone()); next_id = next_id.next_in_line(); } @@ -268,7 +269,7 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 0); - tree.generate_new_node(ChainIndex::root()).unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); assert!( tree.key_map @@ -281,12 +282,12 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 1); - tree.generate_new_node(ChainIndex::root()).unwrap(); - tree.generate_new_node(ChainIndex::root()).unwrap(); - tree.generate_new_node(ChainIndex::root()).unwrap(); - tree.generate_new_node(ChainIndex::root()).unwrap(); - tree.generate_new_node(ChainIndex::root()).unwrap(); - tree.generate_new_node(ChainIndex::root()).unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); let next_last_child_for_parent_id = tree .find_next_last_child_of_id(&ChainIndex::root()) @@ -307,7 +308,7 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 0); - tree.generate_new_node(ChainIndex::root()).unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); assert!( tree.key_map @@ -320,7 +321,7 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 1); - let key_opt = tree.generate_new_node(ChainIndex::from_str("/3").unwrap()); + let key_opt = tree.generate_new_node(&ChainIndex::from_str("/3").unwrap()); assert_eq!(key_opt, None); } @@ -337,7 +338,7 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 0); - tree.generate_new_node(ChainIndex::root()).unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); assert!( tree.key_map @@ -350,7 +351,7 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 1); - tree.generate_new_node(ChainIndex::root()).unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); assert!( tree.key_map @@ -363,7 +364,7 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 2); - tree.generate_new_node(ChainIndex::from_str("/0").unwrap()) + tree.generate_new_node(&ChainIndex::from_str("/0").unwrap()) .unwrap(); let next_last_child_for_parent_id = tree @@ -377,7 +378,7 @@ mod tests { .contains_key(&ChainIndex::from_str("/0/0").unwrap()) ); - tree.generate_new_node(ChainIndex::from_str("/0").unwrap()) + tree.generate_new_node(&ChainIndex::from_str("/0").unwrap()) .unwrap(); let next_last_child_for_parent_id = tree @@ -391,7 +392,7 @@ mod tests { .contains_key(&ChainIndex::from_str("/0/1").unwrap()) ); - tree.generate_new_node(ChainIndex::from_str("/0").unwrap()) + tree.generate_new_node(&ChainIndex::from_str("/0").unwrap()) .unwrap(); let next_last_child_for_parent_id = tree @@ -405,7 +406,7 @@ mod tests { .contains_key(&ChainIndex::from_str("/0/2").unwrap()) ); - tree.generate_new_node(ChainIndex::from_str("/0/1").unwrap()) + tree.generate_new_node(&ChainIndex::from_str("/0/1").unwrap()) .unwrap(); assert!( @@ -419,4 +420,109 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 1); } + + #[test] + fn test_cleanup_leftovers() { + let mut tree = KeyTreePrivate::new(&seed_holder_for_tests()); + + tree.generate_tree_for_depth(5); + + for (chain_id, keys) in &tree.key_map { + println!("{chain_id} : {}", keys.account_id()); + } + + let acc_1 = tree + .key_map + .get_mut(&ChainIndex::from_str("/1").unwrap()) + .unwrap(); + acc_1.value.1.balance = 100; + + let acc_3 = tree + .key_map + .get_mut(&ChainIndex::from_str("/3").unwrap()) + .unwrap(); + acc_3.value.1.balance = 100; + + tree.cleanup_tree_for_depth(5); + + println!("TREE AFTER CLEANUP"); + + for (chain_id, keys) in &tree.key_map { + println!("{chain_id} : {}", keys.account_id()); + } + + let next_last_child_of_root1 = tree + .find_next_last_child_of_id(&ChainIndex::root()) + .unwrap(); + + println!("next_last_child_of_root {next_last_child_of_root1}"); + + let (account_id, chain_id) = tree.generate_new_node(&ChainIndex::root()).unwrap(); + println!("{chain_id} : {account_id}"); + + let next_last_child_of_root2 = tree + .find_next_last_child_of_id(&ChainIndex::root()) + .unwrap(); + + println!("next_last_child_of_root {next_last_child_of_root2}"); + + let (account_id, chain_id) = tree.generate_new_node(&ChainIndex::root()).unwrap(); + println!("{chain_id} : {account_id}"); + + let next_last_child_of_root3 = tree + .find_next_last_child_of_id(&ChainIndex::root()) + .unwrap(); + + println!("next_last_child_of_root {next_last_child_of_root3}"); + + let (account_id, chain_id) = tree.generate_new_node(&ChainIndex::root()).unwrap(); + println!("{chain_id} : {account_id}"); + + let acc_5 = tree + .key_map + .get_mut(&ChainIndex::from_str("/5").unwrap()) + .unwrap(); + acc_5.value.1.balance = 100; + + tree.cleanup_tree_for_depth(10); + + println!("TREE AFTER CLEANUP"); + + for (chain_id, keys) in &tree.key_map { + println!("{chain_id} : {}", keys.account_id()); + } + + let next_last_child_of_root1 = tree + .find_next_last_child_of_id(&ChainIndex::root()) + .unwrap(); + + println!("next_last_child_of_root {next_last_child_of_root1}"); + + let (account_id, chain_id) = tree.generate_new_node(&ChainIndex::root()).unwrap(); + println!("{chain_id} : {account_id}"); + + let next_last_child_of_root2 = tree + .find_next_last_child_of_id(&ChainIndex::root()) + .unwrap(); + + println!("next_last_child_of_root {next_last_child_of_root2}"); + + let (account_id, chain_id) = tree.generate_new_node(&ChainIndex::root()).unwrap(); + println!("{chain_id} : {account_id}"); + + let next_last_child_of_root3 = tree + .find_next_last_child_of_id(&ChainIndex::root()) + .unwrap(); + + println!("next_last_child_of_root {next_last_child_of_root3}"); + + let (account_id, chain_id) = tree.generate_new_node(&ChainIndex::root()).unwrap(); + println!("{chain_id} : {account_id}"); + + println!("TREE AFTER MANIPULATIONS"); + + for (chain_id, keys) in &tree.key_map { + println!("{chain_id} : {}", keys.account_id()); + } + } } diff --git a/key_protocol/src/key_protocol_core/mod.rs b/key_protocol/src/key_protocol_core/mod.rs index ac5ee48..ca0aa54 100644 --- a/key_protocol/src/key_protocol_core/mod.rs +++ b/key_protocol/src/key_protocol_core/mod.rs @@ -91,7 +91,10 @@ impl NSSAUserData { &mut self, parent_cci: ChainIndex, ) -> nssa::AccountId { - self.public_key_tree.generate_new_node(parent_cci).unwrap() + self.public_key_tree + .generate_new_node(&parent_cci) + .unwrap() + .0 } /// Returns the signing key for public transaction signatures @@ -115,7 +118,10 @@ impl NSSAUserData { &mut self, parent_cci: ChainIndex, ) -> nssa::AccountId { - self.private_key_tree.generate_new_node(parent_cci).unwrap() + self.private_key_tree + .generate_new_node(&parent_cci) + .unwrap() + .0 } /// Returns the signing key for public transaction signatures From 926a292c9c34c61cc7634f8f9bf56e12a8132959 Mon Sep 17 00:00:00 2001 From: Pravdyvy Date: Thu, 4 Dec 2025 16:49:10 +0200 Subject: [PATCH 05/10] fix: constraint added --- key_protocol/Cargo.toml | 1 + .../key_management/key_tree/chain_index.rs | 62 ++++ .../src/key_management/key_tree/mod.rs | 341 +++++++++++++----- key_protocol/src/key_protocol_core/mod.rs | 17 +- wallet/src/cli/account.rs | 2 +- wallet/src/lib.rs | 5 +- 6 files changed, 316 insertions(+), 112 deletions(-) diff --git a/key_protocol/Cargo.toml b/key_protocol/Cargo.toml index a562515..4f94c79 100644 --- a/key_protocol/Cargo.toml +++ b/key_protocol/Cargo.toml @@ -22,3 +22,4 @@ path = "../common" [dependencies.nssa] path = "../nssa" +features = ["no_docker"] diff --git a/key_protocol/src/key_management/key_tree/chain_index.rs b/key_protocol/src/key_management/key_tree/chain_index.rs index 31a6a07..31445e7 100644 --- a/key_protocol/src/key_management/key_tree/chain_index.rs +++ b/key_protocol/src/key_management/key_tree/chain_index.rs @@ -77,6 +77,23 @@ impl ChainIndex { ChainIndex(chain) } + pub fn previous_in_line(&self) -> Option { + let mut chain = self.0.clone(); + if let Some(last_p) = chain.last_mut() { + *last_p = last_p.checked_sub(1)?; + } + + Some(ChainIndex(chain)) + } + + pub fn parent(&self) -> Option { + if self.0.is_empty() { + None + } else { + Some(ChainIndex(self.0[..(self.0.len() - 1)].to_vec())) + } + } + pub fn nth_child(&self, child_id: u32) -> ChainIndex { let mut chain = self.0.clone(); chain.push(child_id); @@ -155,4 +172,49 @@ mod tests { assert_eq!(string_index, "/5/7/8".to_string()); } + + #[test] + fn test_prev_in_line() { + let chain_id = ChainIndex(vec![1, 7, 3]); + + let prev_chain_id = chain_id.previous_in_line().unwrap(); + + assert_eq!(prev_chain_id, ChainIndex(vec![1, 7, 2])) + } + + #[test] + fn test_prev_in_line_no_prev() { + let chain_id = ChainIndex(vec![1, 7, 0]); + + let prev_chain_id = chain_id.previous_in_line(); + + assert_eq!(prev_chain_id, None) + } + + #[test] + fn test_parent() { + let chain_id = ChainIndex(vec![1, 7, 3]); + + let parent_chain_id = chain_id.parent().unwrap(); + + assert_eq!(parent_chain_id, ChainIndex(vec![1, 7])) + } + + #[test] + fn test_parent_no_parent() { + let chain_id = ChainIndex(vec![]); + + let parent_chain_id = chain_id.parent(); + + assert_eq!(parent_chain_id, None) + } + + #[test] + fn test_parent_root() { + let chain_id = ChainIndex(vec![1]); + + let parent_chain_id = chain_id.parent().unwrap(); + + assert_eq!(parent_chain_id, ChainIndex::root()) + } } diff --git a/key_protocol/src/key_management/key_tree/mod.rs b/key_protocol/src/key_management/key_tree/mod.rs index 7b85d8e..5f2616d 100644 --- a/key_protocol/src/key_management/key_tree/mod.rs +++ b/key_protocol/src/key_management/key_tree/mod.rs @@ -4,7 +4,7 @@ use std::{ }; use anyhow::Result; -use common::sequencer_client::SequencerClient; +use common::{error::SequencerClientError, sequencer_client::SequencerClient}; use serde::{Deserialize, Serialize}; use crate::key_management::{ @@ -29,6 +29,16 @@ pub struct KeyTree { pub type KeyTreePublic = KeyTree; pub type KeyTreePrivate = KeyTree; +#[derive(thiserror::Error, Debug)] +pub enum KeyTreeGenerationError { + #[error("Parent chain id {0} not present in tree")] + ParentChainIdNotFound(ChainIndex), + #[error("Parent or left relative of {0} is not initialized")] + PredecesorsNotInitialized(ChainIndex), + #[error("Sequencer client error {0:#?}")] + SequencerClientError(#[from] SequencerClientError), +} + impl KeyTree { pub fn new(seed: &SeedHolder) -> Self { let seed_fit: [u8; 64] = seed @@ -95,13 +105,13 @@ impl KeyTree { right = (left_border + right) / 2; } (None, Some(_)) => { - break Some(right); + unreachable!(); } } } } - pub fn generate_new_node( + fn generate_new_node_unconstrained( &mut self, parent_cci: &ChainIndex, ) -> Option<(nssa::AccountId, ChainIndex)> { @@ -155,7 +165,7 @@ impl KeyTree { let mut next_id = curr_id.nth_child(0); while (next_id.depth()) < depth { - self.generate_new_node(&curr_id); + self.generate_new_node_unconstrained(&curr_id); id_stack.push(next_id.clone()); next_id = next_id.next_in_line(); } @@ -164,6 +174,45 @@ impl KeyTree { } impl KeyTree { + pub fn generate_new_node( + &mut self, + parent_cci: &ChainIndex, + ) -> Result<(nssa::AccountId, ChainIndex), KeyTreeGenerationError> { + let father_keys = + self.key_map + .get(parent_cci) + .ok_or(KeyTreeGenerationError::ParentChainIdNotFound( + parent_cci.clone(), + ))?; + let next_child_id = self + .find_next_last_child_of_id(parent_cci) + .expect("Can be None only if parent is not present"); + let next_cci = parent_cci.nth_child(next_child_id); + + if let Some(prev_cci) = next_cci.previous_in_line() { + let prev_keys = self.key_map.get(&prev_cci).expect( + format!("Constraint violated, previous child with id {prev_cci} is missing") + .as_str(), + ); + + if prev_keys.value.1 == nssa::Account::default() { + return Err(KeyTreeGenerationError::PredecesorsNotInitialized(next_cci)); + } + } else if *parent_cci != ChainIndex::root() { + if father_keys.value.1 == nssa::Account::default() { + return Err(KeyTreeGenerationError::PredecesorsNotInitialized(next_cci)); + } + } + + let child_keys = father_keys.nth_child(next_child_id); + let account_id = child_keys.account_id(); + + self.key_map.insert(next_cci.clone(), child_keys); + self.account_id_map.insert(account_id, next_cci.clone()); + + Ok((account_id, next_cci)) + } + /// Cleanup of all non-initialized accounts in a private tree /// /// For given `depth` checks children to a tree such that their `ChainIndex::depth(&self) < @@ -195,6 +244,55 @@ impl KeyTree { } impl KeyTree { + pub async fn generate_new_node( + &mut self, + parent_cci: &ChainIndex, + client: Arc, + ) -> Result<(nssa::AccountId, ChainIndex), KeyTreeGenerationError> { + let father_keys = + self.key_map + .get(parent_cci) + .ok_or(KeyTreeGenerationError::ParentChainIdNotFound( + parent_cci.clone(), + ))?; + let next_child_id = self + .find_next_last_child_of_id(parent_cci) + .expect("Can be None only if parent is not present"); + let next_cci = parent_cci.nth_child(next_child_id); + + if let Some(prev_cci) = next_cci.previous_in_line() { + let prev_keys = self.key_map.get(&prev_cci).expect( + format!("Constraint violated, previous child with id {prev_cci} is missing") + .as_str(), + ); + let prev_acc = client + .get_account(prev_keys.account_id().to_string()) + .await? + .account; + + if prev_acc == nssa::Account::default() { + return Err(KeyTreeGenerationError::PredecesorsNotInitialized(next_cci)); + } + } else if *parent_cci != ChainIndex::root() { + let parent_acc = client + .get_account(father_keys.account_id().to_string()) + .await? + .account; + + if parent_acc == nssa::Account::default() { + return Err(KeyTreeGenerationError::PredecesorsNotInitialized(next_cci)); + } + } + + let child_keys = father_keys.nth_child(next_child_id); + let account_id = child_keys.account_id(); + + self.key_map.insert(next_cci.clone(), child_keys); + self.account_id_map.insert(account_id, next_cci.clone()); + + Ok((account_id, next_cci)) + } + /// Cleanup of all non-initialized accounts in a public tree /// /// For given `depth` checks children to a tree such that their `ChainIndex::depth(&self) < @@ -232,7 +330,7 @@ impl KeyTree { #[cfg(test)] mod tests { - use std::str::FromStr; + use std::{collections::HashSet, str::FromStr}; use nssa::AccountId; @@ -261,7 +359,7 @@ mod tests { fn test_small_key_tree() { let seed_holder = seed_holder_for_tests(); - let mut tree = KeyTreePublic::new(&seed_holder); + let mut tree = KeyTreePrivate::new(&seed_holder); let next_last_child_for_parent_id = tree .find_next_last_child_of_id(&ChainIndex::root()) @@ -269,7 +367,8 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 0); - tree.generate_new_node(&ChainIndex::root()).unwrap(); + tree.generate_new_node_unconstrained(&ChainIndex::root()) + .unwrap(); assert!( tree.key_map @@ -282,12 +381,18 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 1); - tree.generate_new_node(&ChainIndex::root()).unwrap(); - tree.generate_new_node(&ChainIndex::root()).unwrap(); - tree.generate_new_node(&ChainIndex::root()).unwrap(); - tree.generate_new_node(&ChainIndex::root()).unwrap(); - tree.generate_new_node(&ChainIndex::root()).unwrap(); - tree.generate_new_node(&ChainIndex::root()).unwrap(); + tree.generate_new_node_unconstrained(&ChainIndex::root()) + .unwrap(); + tree.generate_new_node_unconstrained(&ChainIndex::root()) + .unwrap(); + tree.generate_new_node_unconstrained(&ChainIndex::root()) + .unwrap(); + tree.generate_new_node_unconstrained(&ChainIndex::root()) + .unwrap(); + tree.generate_new_node_unconstrained(&ChainIndex::root()) + .unwrap(); + tree.generate_new_node_unconstrained(&ChainIndex::root()) + .unwrap(); let next_last_child_for_parent_id = tree .find_next_last_child_of_id(&ChainIndex::root()) @@ -300,7 +405,7 @@ mod tests { fn test_key_tree_can_not_make_child_keys() { let seed_holder = seed_holder_for_tests(); - let mut tree = KeyTreePublic::new(&seed_holder); + let mut tree = KeyTreePrivate::new(&seed_holder); let next_last_child_for_parent_id = tree .find_next_last_child_of_id(&ChainIndex::root()) @@ -308,7 +413,8 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 0); - tree.generate_new_node(&ChainIndex::root()).unwrap(); + tree.generate_new_node_unconstrained(&ChainIndex::root()) + .unwrap(); assert!( tree.key_map @@ -321,7 +427,7 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 1); - let key_opt = tree.generate_new_node(&ChainIndex::from_str("/3").unwrap()); + let key_opt = tree.generate_new_node_unconstrained(&ChainIndex::from_str("/3").unwrap()); assert_eq!(key_opt, None); } @@ -338,7 +444,8 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 0); - tree.generate_new_node(&ChainIndex::root()).unwrap(); + tree.generate_new_node_unconstrained(&ChainIndex::root()) + .unwrap(); assert!( tree.key_map @@ -351,7 +458,8 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 1); - tree.generate_new_node(&ChainIndex::root()).unwrap(); + tree.generate_new_node_unconstrained(&ChainIndex::root()) + .unwrap(); assert!( tree.key_map @@ -364,7 +472,7 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 2); - tree.generate_new_node(&ChainIndex::from_str("/0").unwrap()) + tree.generate_new_node_unconstrained(&ChainIndex::from_str("/0").unwrap()) .unwrap(); let next_last_child_for_parent_id = tree @@ -378,7 +486,7 @@ mod tests { .contains_key(&ChainIndex::from_str("/0/0").unwrap()) ); - tree.generate_new_node(&ChainIndex::from_str("/0").unwrap()) + tree.generate_new_node_unconstrained(&ChainIndex::from_str("/0").unwrap()) .unwrap(); let next_last_child_for_parent_id = tree @@ -392,7 +500,7 @@ mod tests { .contains_key(&ChainIndex::from_str("/0/1").unwrap()) ); - tree.generate_new_node(&ChainIndex::from_str("/0").unwrap()) + tree.generate_new_node_unconstrained(&ChainIndex::from_str("/0").unwrap()) .unwrap(); let next_last_child_for_parent_id = tree @@ -406,7 +514,7 @@ mod tests { .contains_key(&ChainIndex::from_str("/0/2").unwrap()) ); - tree.generate_new_node(&ChainIndex::from_str("/0/1").unwrap()) + tree.generate_new_node_unconstrained(&ChainIndex::from_str("/0/1").unwrap()) .unwrap(); assert!( @@ -422,107 +530,144 @@ mod tests { } #[test] - fn test_cleanup_leftovers() { - let mut tree = KeyTreePrivate::new(&seed_holder_for_tests()); + fn test_key_generation_constraint() { + let seed_holder = seed_holder_for_tests(); - tree.generate_tree_for_depth(5); + let mut tree = KeyTreePrivate::new(&seed_holder); - for (chain_id, keys) in &tree.key_map { - println!("{chain_id} : {}", keys.account_id()); - } + let (_, chain_id) = tree.generate_new_node(&ChainIndex::root()).unwrap(); - let acc_1 = tree + assert_eq!(chain_id, ChainIndex::from_str("/0").unwrap()); + + let res = tree.generate_new_node(&ChainIndex::from_str("/").unwrap()); + + assert!(matches!( + res, + Err(KeyTreeGenerationError::PredecesorsNotInitialized(_)) + )); + + let res = tree.generate_new_node(&ChainIndex::from_str("/0").unwrap()); + + assert!(matches!( + res, + Err(KeyTreeGenerationError::PredecesorsNotInitialized(_)) + )); + + let acc = tree + .key_map + .get_mut(&ChainIndex::from_str("/0").unwrap()) + .unwrap(); + acc.value.1.balance = 1; + + let (_, chain_id) = tree + .generate_new_node(&ChainIndex::from_str("/").unwrap()) + .unwrap(); + + assert_eq!(chain_id, ChainIndex::from_str("/1").unwrap()); + + let (_, chain_id) = tree + .generate_new_node(&ChainIndex::from_str("/0").unwrap()) + .unwrap(); + + assert_eq!(chain_id, ChainIndex::from_str("/0/0").unwrap()); + } + + #[test] + fn test_cleanup() { + let seed_holder = seed_holder_for_tests(); + + let mut tree = KeyTreePrivate::new(&seed_holder); + tree.generate_tree_for_depth(10); + + let acc = tree + .key_map + .get_mut(&ChainIndex::from_str("/0").unwrap()) + .unwrap(); + acc.value.1.balance = 1; + + let acc = tree .key_map .get_mut(&ChainIndex::from_str("/1").unwrap()) .unwrap(); - acc_1.value.1.balance = 100; + acc.value.1.balance = 2; - let acc_3 = tree + let acc = tree .key_map - .get_mut(&ChainIndex::from_str("/3").unwrap()) + .get_mut(&ChainIndex::from_str("/2").unwrap()) .unwrap(); - acc_3.value.1.balance = 100; + acc.value.1.balance = 3; - tree.cleanup_tree_for_depth(5); - - println!("TREE AFTER CLEANUP"); - - for (chain_id, keys) in &tree.key_map { - println!("{chain_id} : {}", keys.account_id()); - } - - let next_last_child_of_root1 = tree - .find_next_last_child_of_id(&ChainIndex::root()) - .unwrap(); - - println!("next_last_child_of_root {next_last_child_of_root1}"); - - let (account_id, chain_id) = tree.generate_new_node(&ChainIndex::root()).unwrap(); - println!("{chain_id} : {account_id}"); - - let next_last_child_of_root2 = tree - .find_next_last_child_of_id(&ChainIndex::root()) - .unwrap(); - - println!("next_last_child_of_root {next_last_child_of_root2}"); - - let (account_id, chain_id) = tree.generate_new_node(&ChainIndex::root()).unwrap(); - println!("{chain_id} : {account_id}"); - - let next_last_child_of_root3 = tree - .find_next_last_child_of_id(&ChainIndex::root()) - .unwrap(); - - println!("next_last_child_of_root {next_last_child_of_root3}"); - - let (account_id, chain_id) = tree.generate_new_node(&ChainIndex::root()).unwrap(); - println!("{chain_id} : {account_id}"); - - let acc_5 = tree + let acc = tree .key_map - .get_mut(&ChainIndex::from_str("/5").unwrap()) + .get_mut(&ChainIndex::from_str("/0/0").unwrap()) .unwrap(); - acc_5.value.1.balance = 100; + acc.value.1.balance = 4; + + let acc = tree + .key_map + .get_mut(&ChainIndex::from_str("/0/1").unwrap()) + .unwrap(); + acc.value.1.balance = 5; + + let acc = tree + .key_map + .get_mut(&ChainIndex::from_str("/1/0").unwrap()) + .unwrap(); + acc.value.1.balance = 6; tree.cleanup_tree_for_depth(10); - println!("TREE AFTER CLEANUP"); + let mut key_set_res = HashSet::new(); + key_set_res.insert("/0".to_string()); + key_set_res.insert("/1".to_string()); + key_set_res.insert("/2".to_string()); + key_set_res.insert("/".to_string()); + key_set_res.insert("/0/0".to_string()); + key_set_res.insert("/0/1".to_string()); + key_set_res.insert("/1/0".to_string()); - for (chain_id, keys) in &tree.key_map { - println!("{chain_id} : {}", keys.account_id()); + let mut key_set = HashSet::new(); + + for key in tree.key_map.keys() { + key_set.insert(key.to_string()); } - let next_last_child_of_root1 = tree - .find_next_last_child_of_id(&ChainIndex::root()) + assert_eq!(key_set, key_set_res); + + let acc = tree + .key_map + .get(&ChainIndex::from_str("/0").unwrap()) .unwrap(); + assert_eq!(acc.value.1.balance, 1); - println!("next_last_child_of_root {next_last_child_of_root1}"); - - let (account_id, chain_id) = tree.generate_new_node(&ChainIndex::root()).unwrap(); - println!("{chain_id} : {account_id}"); - - let next_last_child_of_root2 = tree - .find_next_last_child_of_id(&ChainIndex::root()) + let acc = tree + .key_map + .get(&ChainIndex::from_str("/1").unwrap()) .unwrap(); + assert_eq!(acc.value.1.balance, 2); - println!("next_last_child_of_root {next_last_child_of_root2}"); - - let (account_id, chain_id) = tree.generate_new_node(&ChainIndex::root()).unwrap(); - println!("{chain_id} : {account_id}"); - - let next_last_child_of_root3 = tree - .find_next_last_child_of_id(&ChainIndex::root()) + let acc = tree + .key_map + .get(&ChainIndex::from_str("/2").unwrap()) .unwrap(); + assert_eq!(acc.value.1.balance, 3); - println!("next_last_child_of_root {next_last_child_of_root3}"); + let acc = tree + .key_map + .get(&ChainIndex::from_str("/0/0").unwrap()) + .unwrap(); + assert_eq!(acc.value.1.balance, 4); - let (account_id, chain_id) = tree.generate_new_node(&ChainIndex::root()).unwrap(); - println!("{chain_id} : {account_id}"); + let acc = tree + .key_map + .get(&ChainIndex::from_str("/0/1").unwrap()) + .unwrap(); + assert_eq!(acc.value.1.balance, 5); - println!("TREE AFTER MANIPULATIONS"); - - for (chain_id, keys) in &tree.key_map { - println!("{chain_id} : {}", keys.account_id()); - } + let acc = tree + .key_map + .get(&ChainIndex::from_str("/1/0").unwrap()) + .unwrap(); + assert_eq!(acc.value.1.balance, 6); } } diff --git a/key_protocol/src/key_protocol_core/mod.rs b/key_protocol/src/key_protocol_core/mod.rs index ca0aa54..c474df0 100644 --- a/key_protocol/src/key_protocol_core/mod.rs +++ b/key_protocol/src/key_protocol_core/mod.rs @@ -1,6 +1,7 @@ -use std::collections::HashMap; +use std::{collections::HashMap, sync::Arc}; use anyhow::Result; +use common::sequencer_client::SequencerClient; use k256::AffinePoint; use serde::{Deserialize, Serialize}; @@ -87,12 +88,14 @@ impl NSSAUserData { /// Generated new private key for public transaction signatures /// /// Returns the account_id of new account - pub fn generate_new_public_transaction_private_key( + pub async fn generate_new_public_transaction_private_key( &mut self, parent_cci: ChainIndex, + sequencer_client: Arc, ) -> nssa::AccountId { self.public_key_tree - .generate_new_node(&parent_cci) + .generate_new_node(&parent_cci, sequencer_client) + .await .unwrap() .0 } @@ -175,17 +178,9 @@ mod tests { fn test_new_account() { let mut user_data = NSSAUserData::default(); - let account_id_pub = - user_data.generate_new_public_transaction_private_key(ChainIndex::root()); let account_id_private = user_data.generate_new_privacy_preserving_transaction_key_chain(ChainIndex::root()); - let is_private_key_generated = user_data - .get_pub_account_signing_key(&account_id_pub) - .is_some(); - - assert!(is_private_key_generated); - let is_key_chain_generated = user_data.get_private_account(&account_id_private).is_some(); assert!(is_key_chain_generated); diff --git a/wallet/src/cli/account.rs b/wallet/src/cli/account.rs index 5b23b2b..1aa2bd3 100644 --- a/wallet/src/cli/account.rs +++ b/wallet/src/cli/account.rs @@ -115,7 +115,7 @@ impl WalletSubcommand for NewSubcommand { ) -> Result { match self { NewSubcommand::Public { cci } => { - let account_id = wallet_core.create_new_account_public(cci); + let account_id = wallet_core.create_new_account_public(cci).await; println!("Generated new account with account_id Public/{account_id}"); diff --git a/wallet/src/lib.rs b/wallet/src/lib.rs index f79d947..538076e 100644 --- a/wallet/src/lib.rs +++ b/wallet/src/lib.rs @@ -108,10 +108,11 @@ impl WalletCore { Ok(config_path) } - pub fn create_new_account_public(&mut self, chain_index: ChainIndex) -> AccountId { + pub async fn create_new_account_public(&mut self, chain_index: ChainIndex) -> AccountId { self.storage .user_data - .generate_new_public_transaction_private_key(chain_index) + .generate_new_public_transaction_private_key(chain_index, self.sequencer_client.clone()) + .await } pub fn create_new_account_private(&mut self, chain_index: ChainIndex) -> AccountId { From 29c3737704605b728227ec1bfb8544e1b542de7a Mon Sep 17 00:00:00 2001 From: Pravdyvy Date: Fri, 5 Dec 2025 07:46:59 +0200 Subject: [PATCH 06/10] fix: lint fix comments addressed --- key_protocol/Cargo.toml | 1 - .../key_management/key_tree/chain_index.rs | 8 +--- .../src/key_management/key_tree/mod.rs | 44 +++++++++---------- wallet/src/cli/mod.rs | 4 +- 4 files changed, 25 insertions(+), 32 deletions(-) diff --git a/key_protocol/Cargo.toml b/key_protocol/Cargo.toml index 4f94c79..a562515 100644 --- a/key_protocol/Cargo.toml +++ b/key_protocol/Cargo.toml @@ -22,4 +22,3 @@ path = "../common" [dependencies.nssa] path = "../nssa" -features = ["no_docker"] diff --git a/key_protocol/src/key_management/key_tree/chain_index.rs b/key_protocol/src/key_management/key_tree/chain_index.rs index 31445e7..8b28327 100644 --- a/key_protocol/src/key_management/key_tree/chain_index.rs +++ b/key_protocol/src/key_management/key_tree/chain_index.rs @@ -102,13 +102,7 @@ impl ChainIndex { } pub fn depth(&self) -> u32 { - let mut res = 0; - - for cci in &self.0 { - res += cci + 1; - } - - res + self.0.iter().map(|cci| cci + 1).sum() } } diff --git a/key_protocol/src/key_management/key_tree/mod.rs b/key_protocol/src/key_management/key_tree/mod.rs index 5f2616d..b75671e 100644 --- a/key_protocol/src/key_management/key_tree/mod.rs +++ b/key_protocol/src/key_management/key_tree/mod.rs @@ -115,13 +115,13 @@ impl KeyTree { &mut self, parent_cci: &ChainIndex, ) -> Option<(nssa::AccountId, ChainIndex)> { - let father_keys = self.key_map.get(parent_cci)?; + let parent_keys = self.key_map.get(parent_cci)?; let next_child_id = self .find_next_last_child_of_id(parent_cci) .expect("Can be None only if parent is not present"); let next_cci = parent_cci.nth_child(next_child_id); - let child_keys = father_keys.nth_child(next_child_id); + let child_keys = parent_keys.nth_child(next_child_id); let account_id = child_keys.account_id(); self.key_map.insert(next_cci.clone(), child_keys); @@ -174,11 +174,12 @@ impl KeyTree { } impl KeyTree { + #[allow(clippy::result_large_err)] pub fn generate_new_node( &mut self, parent_cci: &ChainIndex, ) -> Result<(nssa::AccountId, ChainIndex), KeyTreeGenerationError> { - let father_keys = + let parent_keys = self.key_map .get(parent_cci) .ok_or(KeyTreeGenerationError::ParentChainIdNotFound( @@ -190,21 +191,20 @@ impl KeyTree { let next_cci = parent_cci.nth_child(next_child_id); if let Some(prev_cci) = next_cci.previous_in_line() { - let prev_keys = self.key_map.get(&prev_cci).expect( - format!("Constraint violated, previous child with id {prev_cci} is missing") - .as_str(), - ); + let prev_keys = self.key_map.get(&prev_cci).unwrap_or_else(|| { + panic!("Constraint violated, previous child with id {prev_cci} is missing") + }); if prev_keys.value.1 == nssa::Account::default() { return Err(KeyTreeGenerationError::PredecesorsNotInitialized(next_cci)); } - } else if *parent_cci != ChainIndex::root() { - if father_keys.value.1 == nssa::Account::default() { - return Err(KeyTreeGenerationError::PredecesorsNotInitialized(next_cci)); - } + } else if *parent_cci != ChainIndex::root() + && parent_keys.value.1 == nssa::Account::default() + { + return Err(KeyTreeGenerationError::PredecesorsNotInitialized(next_cci)); } - let child_keys = father_keys.nth_child(next_child_id); + let child_keys = parent_keys.nth_child(next_child_id); let account_id = child_keys.account_id(); self.key_map.insert(next_cci.clone(), child_keys); @@ -221,7 +221,7 @@ impl KeyTree { /// If account is default, removes them. /// /// Chain must be parsed for accounts beforehand - pub fn cleanup_tree_for_depth(&mut self, depth: u32) { + pub fn cleanup_tree_remove_ininit_for_depth(&mut self, depth: u32) { let mut id_stack = vec![ChainIndex::root()]; while let Some(curr_id) = id_stack.pop() { @@ -244,12 +244,13 @@ impl KeyTree { } impl KeyTree { + #[allow(clippy::result_large_err)] pub async fn generate_new_node( &mut self, parent_cci: &ChainIndex, client: Arc, ) -> Result<(nssa::AccountId, ChainIndex), KeyTreeGenerationError> { - let father_keys = + let parent_keys = self.key_map .get(parent_cci) .ok_or(KeyTreeGenerationError::ParentChainIdNotFound( @@ -261,10 +262,9 @@ impl KeyTree { let next_cci = parent_cci.nth_child(next_child_id); if let Some(prev_cci) = next_cci.previous_in_line() { - let prev_keys = self.key_map.get(&prev_cci).expect( - format!("Constraint violated, previous child with id {prev_cci} is missing") - .as_str(), - ); + let prev_keys = self.key_map.get(&prev_cci).unwrap_or_else(|| { + panic!("Constraint violated, previous child with id {prev_cci} is missing") + }); let prev_acc = client .get_account(prev_keys.account_id().to_string()) .await? @@ -275,7 +275,7 @@ impl KeyTree { } } else if *parent_cci != ChainIndex::root() { let parent_acc = client - .get_account(father_keys.account_id().to_string()) + .get_account(parent_keys.account_id().to_string()) .await? .account; @@ -284,7 +284,7 @@ impl KeyTree { } } - let child_keys = father_keys.nth_child(next_child_id); + let child_keys = parent_keys.nth_child(next_child_id); let account_id = child_keys.account_id(); self.key_map.insert(next_cci.clone(), child_keys); @@ -299,7 +299,7 @@ impl KeyTree { /// depth`. /// /// If account is default, removes them. - pub async fn cleanup_tree_for_depth( + pub async fn cleanup_tree_remove_ininit_for_depth( &mut self, depth: u32, client: Arc, @@ -615,7 +615,7 @@ mod tests { .unwrap(); acc.value.1.balance = 6; - tree.cleanup_tree_for_depth(10); + tree.cleanup_tree_remove_ininit_for_depth(10); let mut key_set_res = HashSet::new(); key_set_res.insert("/0".to_string()); diff --git a/wallet/src/cli/mod.rs b/wallet/src/cli/mod.rs index 1f0e53b..07bf252 100644 --- a/wallet/src/cli/mod.rs +++ b/wallet/src/cli/mod.rs @@ -231,7 +231,7 @@ pub async fn execute_keys_restoration(password: String, depth: u32) -> Result<() .storage .user_data .public_key_tree - .cleanup_tree_for_depth(depth, wallet_core.sequencer_client.clone()) + .cleanup_tree_remove_ininit_for_depth(depth, wallet_core.sequencer_client.clone()) .await?; println!("Public tree cleaned up"); @@ -258,7 +258,7 @@ pub async fn execute_keys_restoration(password: String, depth: u32) -> Result<() .storage .user_data .private_key_tree - .cleanup_tree_for_depth(depth); + .cleanup_tree_remove_ininit_for_depth(depth); println!("Private tree cleaned up"); From d2dc255bc51b02d996a58868d4c77c3454dc5a6e Mon Sep 17 00:00:00 2001 From: Pravdyvy Date: Fri, 5 Dec 2025 09:36:53 +0200 Subject: [PATCH 07/10] fix: merge fix --- wallet/src/cli/mod.rs | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/wallet/src/cli/mod.rs b/wallet/src/cli/mod.rs index 9d207c7..cda599b 100644 --- a/wallet/src/cli/mod.rs +++ b/wallet/src/cli/mod.rs @@ -232,13 +232,7 @@ pub async fn execute_keys_restoration(password: String, depth: u32) -> Result<() println!("Last block is {last_block}"); - parse_block_range( - 1, - last_block, - wallet_core.sequencer_client.clone(), - &mut wallet_core, - ) - .await?; + wallet_core.sync_to_block(last_block).await?; println!("Private tree clean up start"); From 40991cf6d1be3540ccc292edeffb8043f6f67907 Mon Sep 17 00:00:00 2001 From: Pravdyvy Date: Mon, 8 Dec 2025 12:11:11 +0200 Subject: [PATCH 08/10] fix: layered cleanup --- key_protocol/Cargo.toml | 2 + .../key_management/key_tree/chain_index.rs | 70 ++++- .../src/key_management/key_tree/mod.rs | 264 ++++++------------ key_protocol/src/key_protocol_core/mod.rs | 9 +- wallet/src/cli/mod.rs | 4 +- wallet/src/lib.rs | 3 +- 6 files changed, 157 insertions(+), 195 deletions(-) diff --git a/key_protocol/Cargo.toml b/key_protocol/Cargo.toml index a562515..24b92c0 100644 --- a/key_protocol/Cargo.toml +++ b/key_protocol/Cargo.toml @@ -16,9 +16,11 @@ bip39.workspace = true hmac-sha512.workspace = true thiserror.workspace = true nssa-core = { path = "../nssa/core", features = ["host"] } +itertools.workspace = true [dependencies.common] path = "../common" [dependencies.nssa] path = "../nssa" +features = ["no_docker"] diff --git a/key_protocol/src/key_management/key_tree/chain_index.rs b/key_protocol/src/key_management/key_tree/chain_index.rs index 8b28327..d2c9c3b 100644 --- a/key_protocol/src/key_management/key_tree/chain_index.rs +++ b/key_protocol/src/key_management/key_tree/chain_index.rs @@ -1,8 +1,9 @@ use std::{fmt::Display, str::FromStr}; +use itertools::Itertools; use serde::{Deserialize, Serialize}; -#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize)] +#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, Hash)] pub struct ChainIndex(Vec); #[derive(thiserror::Error, Debug)] @@ -104,6 +105,39 @@ impl ChainIndex { pub fn depth(&self) -> u32 { self.0.iter().map(|cci| cci + 1).sum() } + + fn collapse_back(&self) -> Option { + let mut res = self.parent()?; + + let last_mut = res.0.last_mut()?; + *last_mut += *(self.0.last()?) + 1; + + Some(res) + } + + fn shuffle_iter(&self) -> impl Iterator { + self.0 + .iter() + .permutations(self.0.len()) + .unique() + .map(|item| ChainIndex(item.into_iter().cloned().collect())) + } + + pub fn chain_ids_at_depth(depth: usize) -> impl Iterator { + let mut stack = vec![ChainIndex(vec![0; depth])]; + let mut cumulative_stack = vec![ChainIndex(vec![0; depth])]; + + while let Some(id) = stack.pop() { + if let Some(collapsed_id) = id.collapse_back() { + for id in collapsed_id.shuffle_iter() { + stack.push(id.clone()); + cumulative_stack.push(id); + } + } + } + + cumulative_stack.into_iter().unique() + } } #[cfg(test)] @@ -211,4 +245,38 @@ mod tests { assert_eq!(parent_chain_id, ChainIndex::root()) } + + #[test] + fn test_collapse_back() { + let chain_id = ChainIndex(vec![1, 1]); + + let collapsed = chain_id.collapse_back().unwrap(); + + assert_eq!(collapsed, ChainIndex(vec![3])) + } + + #[test] + fn test_collapse_back_one() { + let chain_id = ChainIndex(vec![1]); + + let collapsed = chain_id.collapse_back(); + + assert_eq!(collapsed, None) + } + + #[test] + fn test_collapse_back_root() { + let chain_id = ChainIndex(vec![]); + + let collapsed = chain_id.collapse_back(); + + assert_eq!(collapsed, None) + } + + #[test] + fn test_shuffle() { + for id in ChainIndex::chain_ids_at_depth(5) { + println!("{id}"); + } + } } diff --git a/key_protocol/src/key_management/key_tree/mod.rs b/key_protocol/src/key_management/key_tree/mod.rs index b75671e..d964af4 100644 --- a/key_protocol/src/key_management/key_tree/mod.rs +++ b/key_protocol/src/key_management/key_tree/mod.rs @@ -111,7 +111,7 @@ impl KeyTree { } } - fn generate_new_node_unconstrained( + pub fn generate_new_node( &mut self, parent_cci: &ChainIndex, ) -> Option<(nssa::AccountId, ChainIndex)> { @@ -165,7 +165,7 @@ impl KeyTree { let mut next_id = curr_id.nth_child(0); while (next_id.depth()) < depth { - self.generate_new_node_unconstrained(&curr_id); + self.generate_new_node(&curr_id); id_stack.push(next_id.clone()); next_id = next_id.next_in_line(); } @@ -174,45 +174,6 @@ impl KeyTree { } impl KeyTree { - #[allow(clippy::result_large_err)] - pub fn generate_new_node( - &mut self, - parent_cci: &ChainIndex, - ) -> Result<(nssa::AccountId, ChainIndex), KeyTreeGenerationError> { - let parent_keys = - self.key_map - .get(parent_cci) - .ok_or(KeyTreeGenerationError::ParentChainIdNotFound( - parent_cci.clone(), - ))?; - let next_child_id = self - .find_next_last_child_of_id(parent_cci) - .expect("Can be None only if parent is not present"); - let next_cci = parent_cci.nth_child(next_child_id); - - if let Some(prev_cci) = next_cci.previous_in_line() { - let prev_keys = self.key_map.get(&prev_cci).unwrap_or_else(|| { - panic!("Constraint violated, previous child with id {prev_cci} is missing") - }); - - if prev_keys.value.1 == nssa::Account::default() { - return Err(KeyTreeGenerationError::PredecesorsNotInitialized(next_cci)); - } - } else if *parent_cci != ChainIndex::root() - && parent_keys.value.1 == nssa::Account::default() - { - return Err(KeyTreeGenerationError::PredecesorsNotInitialized(next_cci)); - } - - let child_keys = parent_keys.nth_child(next_child_id); - let account_id = child_keys.account_id(); - - self.key_map.insert(next_cci.clone(), child_keys); - self.account_id_map.insert(account_id, next_cci.clone()); - - Ok((account_id, next_cci)) - } - /// Cleanup of all non-initialized accounts in a private tree /// /// For given `depth` checks children to a tree such that their `ChainIndex::depth(&self) < @@ -221,7 +182,9 @@ impl KeyTree { /// If account is default, removes them. /// /// Chain must be parsed for accounts beforehand - pub fn cleanup_tree_remove_ininit_for_depth(&mut self, depth: u32) { + /// + /// Fast, leaves gaps between accounts + pub fn cleanup_tree_remove_uninit_for_depth(&mut self, depth: u32) { let mut id_stack = vec![ChainIndex::root()]; while let Some(curr_id) = id_stack.pop() { @@ -241,64 +204,42 @@ impl KeyTree { } } } + + /// Cleanup of non-initialized accounts in a private tree + /// + /// If account is default, removes them, stops at first non-default account. + /// + /// Walks through tree in lairs of same depth using `ChainIndex::chain_ids_at_depth()` + /// + /// Chain must be parsed for accounts beforehand + /// + /// Slow, maintains tree consistency. + pub fn cleanup_tree_remove_uninit_layered(&mut self, depth: u32) { + 'outer: for i in (1..(depth as usize)).rev() { + println!("Cleanup of tree at depth {i}"); + for id in ChainIndex::chain_ids_at_depth(i) { + if let Some(node) = self.key_map.get(&id) { + if node.value.1 == nssa::Account::default() { + let addr = node.account_id(); + self.remove(addr); + } else { + break 'outer; + } + } + } + } + } } impl KeyTree { - #[allow(clippy::result_large_err)] - pub async fn generate_new_node( - &mut self, - parent_cci: &ChainIndex, - client: Arc, - ) -> Result<(nssa::AccountId, ChainIndex), KeyTreeGenerationError> { - let parent_keys = - self.key_map - .get(parent_cci) - .ok_or(KeyTreeGenerationError::ParentChainIdNotFound( - parent_cci.clone(), - ))?; - let next_child_id = self - .find_next_last_child_of_id(parent_cci) - .expect("Can be None only if parent is not present"); - let next_cci = parent_cci.nth_child(next_child_id); - - if let Some(prev_cci) = next_cci.previous_in_line() { - let prev_keys = self.key_map.get(&prev_cci).unwrap_or_else(|| { - panic!("Constraint violated, previous child with id {prev_cci} is missing") - }); - let prev_acc = client - .get_account(prev_keys.account_id().to_string()) - .await? - .account; - - if prev_acc == nssa::Account::default() { - return Err(KeyTreeGenerationError::PredecesorsNotInitialized(next_cci)); - } - } else if *parent_cci != ChainIndex::root() { - let parent_acc = client - .get_account(parent_keys.account_id().to_string()) - .await? - .account; - - if parent_acc == nssa::Account::default() { - return Err(KeyTreeGenerationError::PredecesorsNotInitialized(next_cci)); - } - } - - let child_keys = parent_keys.nth_child(next_child_id); - let account_id = child_keys.account_id(); - - self.key_map.insert(next_cci.clone(), child_keys); - self.account_id_map.insert(account_id, next_cci.clone()); - - Ok((account_id, next_cci)) - } - /// Cleanup of all non-initialized accounts in a public tree /// /// For given `depth` checks children to a tree such that their `ChainIndex::depth(&self) < /// depth`. /// /// If account is default, removes them. + /// + /// Fast, leaves gaps between accounts pub async fn cleanup_tree_remove_ininit_for_depth( &mut self, depth: u32, @@ -326,6 +267,38 @@ impl KeyTree { Ok(()) } + + /// Cleanup of non-initialized accounts in a public tree + /// + /// If account is default, removes them, stops at first non-default account. + /// + /// Walks through tree in lairs of same depth using `ChainIndex::chain_ids_at_depth()` + /// + /// Slow, maintains tree consistency. + pub async fn cleanup_tree_remove_uninit_layered( + &mut self, + depth: u32, + client: Arc, + ) -> Result<()> { + 'outer: for i in (1..(depth as usize)).rev() { + println!("Cleanup of tree at depth {i}"); + for id in ChainIndex::chain_ids_at_depth(i) { + if let Some(node) = self.key_map.get(&id) { + let address = node.account_id(); + let node_acc = client.get_account(address.to_string()).await?.account; + + if node_acc == nssa::Account::default() { + let addr = node.account_id(); + self.remove(addr); + } else { + break 'outer; + } + } + } + } + + Ok(()) + } } #[cfg(test)] @@ -367,8 +340,7 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 0); - tree.generate_new_node_unconstrained(&ChainIndex::root()) - .unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); assert!( tree.key_map @@ -381,18 +353,12 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 1); - tree.generate_new_node_unconstrained(&ChainIndex::root()) - .unwrap(); - tree.generate_new_node_unconstrained(&ChainIndex::root()) - .unwrap(); - tree.generate_new_node_unconstrained(&ChainIndex::root()) - .unwrap(); - tree.generate_new_node_unconstrained(&ChainIndex::root()) - .unwrap(); - tree.generate_new_node_unconstrained(&ChainIndex::root()) - .unwrap(); - tree.generate_new_node_unconstrained(&ChainIndex::root()) - .unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); let next_last_child_for_parent_id = tree .find_next_last_child_of_id(&ChainIndex::root()) @@ -413,8 +379,7 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 0); - tree.generate_new_node_unconstrained(&ChainIndex::root()) - .unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); assert!( tree.key_map @@ -427,7 +392,7 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 1); - let key_opt = tree.generate_new_node_unconstrained(&ChainIndex::from_str("/3").unwrap()); + let key_opt = tree.generate_new_node(&ChainIndex::from_str("/3").unwrap()); assert_eq!(key_opt, None); } @@ -444,8 +409,7 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 0); - tree.generate_new_node_unconstrained(&ChainIndex::root()) - .unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); assert!( tree.key_map @@ -458,8 +422,7 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 1); - tree.generate_new_node_unconstrained(&ChainIndex::root()) - .unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); assert!( tree.key_map @@ -472,7 +435,7 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 2); - tree.generate_new_node_unconstrained(&ChainIndex::from_str("/0").unwrap()) + tree.generate_new_node(&ChainIndex::from_str("/0").unwrap()) .unwrap(); let next_last_child_for_parent_id = tree @@ -486,7 +449,7 @@ mod tests { .contains_key(&ChainIndex::from_str("/0/0").unwrap()) ); - tree.generate_new_node_unconstrained(&ChainIndex::from_str("/0").unwrap()) + tree.generate_new_node(&ChainIndex::from_str("/0").unwrap()) .unwrap(); let next_last_child_for_parent_id = tree @@ -500,7 +463,7 @@ mod tests { .contains_key(&ChainIndex::from_str("/0/1").unwrap()) ); - tree.generate_new_node_unconstrained(&ChainIndex::from_str("/0").unwrap()) + tree.generate_new_node(&ChainIndex::from_str("/0").unwrap()) .unwrap(); let next_last_child_for_parent_id = tree @@ -514,7 +477,7 @@ mod tests { .contains_key(&ChainIndex::from_str("/0/2").unwrap()) ); - tree.generate_new_node_unconstrained(&ChainIndex::from_str("/0/1").unwrap()) + tree.generate_new_node(&ChainIndex::from_str("/0/1").unwrap()) .unwrap(); assert!( @@ -529,49 +492,6 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 1); } - #[test] - fn test_key_generation_constraint() { - let seed_holder = seed_holder_for_tests(); - - let mut tree = KeyTreePrivate::new(&seed_holder); - - let (_, chain_id) = tree.generate_new_node(&ChainIndex::root()).unwrap(); - - assert_eq!(chain_id, ChainIndex::from_str("/0").unwrap()); - - let res = tree.generate_new_node(&ChainIndex::from_str("/").unwrap()); - - assert!(matches!( - res, - Err(KeyTreeGenerationError::PredecesorsNotInitialized(_)) - )); - - let res = tree.generate_new_node(&ChainIndex::from_str("/0").unwrap()); - - assert!(matches!( - res, - Err(KeyTreeGenerationError::PredecesorsNotInitialized(_)) - )); - - let acc = tree - .key_map - .get_mut(&ChainIndex::from_str("/0").unwrap()) - .unwrap(); - acc.value.1.balance = 1; - - let (_, chain_id) = tree - .generate_new_node(&ChainIndex::from_str("/").unwrap()) - .unwrap(); - - assert_eq!(chain_id, ChainIndex::from_str("/1").unwrap()); - - let (_, chain_id) = tree - .generate_new_node(&ChainIndex::from_str("/0").unwrap()) - .unwrap(); - - assert_eq!(chain_id, ChainIndex::from_str("/0/0").unwrap()); - } - #[test] fn test_cleanup() { let seed_holder = seed_holder_for_tests(); @@ -579,12 +499,6 @@ mod tests { let mut tree = KeyTreePrivate::new(&seed_holder); tree.generate_tree_for_depth(10); - let acc = tree - .key_map - .get_mut(&ChainIndex::from_str("/0").unwrap()) - .unwrap(); - acc.value.1.balance = 1; - let acc = tree .key_map .get_mut(&ChainIndex::from_str("/1").unwrap()) @@ -597,12 +511,6 @@ mod tests { .unwrap(); acc.value.1.balance = 3; - let acc = tree - .key_map - .get_mut(&ChainIndex::from_str("/0/0").unwrap()) - .unwrap(); - acc.value.1.balance = 4; - let acc = tree .key_map .get_mut(&ChainIndex::from_str("/0/1").unwrap()) @@ -615,7 +523,7 @@ mod tests { .unwrap(); acc.value.1.balance = 6; - tree.cleanup_tree_remove_ininit_for_depth(10); + tree.cleanup_tree_remove_uninit_layered(10); let mut key_set_res = HashSet::new(); key_set_res.insert("/0".to_string()); @@ -634,12 +542,6 @@ mod tests { assert_eq!(key_set, key_set_res); - let acc = tree - .key_map - .get(&ChainIndex::from_str("/0").unwrap()) - .unwrap(); - assert_eq!(acc.value.1.balance, 1); - let acc = tree .key_map .get(&ChainIndex::from_str("/1").unwrap()) @@ -652,12 +554,6 @@ mod tests { .unwrap(); assert_eq!(acc.value.1.balance, 3); - let acc = tree - .key_map - .get(&ChainIndex::from_str("/0/0").unwrap()) - .unwrap(); - assert_eq!(acc.value.1.balance, 4); - let acc = tree .key_map .get(&ChainIndex::from_str("/0/1").unwrap()) diff --git a/key_protocol/src/key_protocol_core/mod.rs b/key_protocol/src/key_protocol_core/mod.rs index c474df0..fc0a393 100644 --- a/key_protocol/src/key_protocol_core/mod.rs +++ b/key_protocol/src/key_protocol_core/mod.rs @@ -1,7 +1,6 @@ -use std::{collections::HashMap, sync::Arc}; +use std::collections::HashMap; use anyhow::Result; -use common::sequencer_client::SequencerClient; use k256::AffinePoint; use serde::{Deserialize, Serialize}; @@ -88,14 +87,12 @@ impl NSSAUserData { /// Generated new private key for public transaction signatures /// /// Returns the account_id of new account - pub async fn generate_new_public_transaction_private_key( + pub fn generate_new_public_transaction_private_key( &mut self, parent_cci: ChainIndex, - sequencer_client: Arc, ) -> nssa::AccountId { self.public_key_tree - .generate_new_node(&parent_cci, sequencer_client) - .await + .generate_new_node(&parent_cci) .unwrap() .0 } diff --git a/wallet/src/cli/mod.rs b/wallet/src/cli/mod.rs index cda599b..e53849e 100644 --- a/wallet/src/cli/mod.rs +++ b/wallet/src/cli/mod.rs @@ -219,7 +219,7 @@ pub async fn execute_keys_restoration(password: String, depth: u32) -> Result<() .storage .user_data .public_key_tree - .cleanup_tree_remove_ininit_for_depth(depth, wallet_core.sequencer_client.clone()) + .cleanup_tree_remove_uninit_layered(depth, wallet_core.sequencer_client.clone()) .await?; println!("Public tree cleaned up"); @@ -240,7 +240,7 @@ pub async fn execute_keys_restoration(password: String, depth: u32) -> Result<() .storage .user_data .private_key_tree - .cleanup_tree_remove_ininit_for_depth(depth); + .cleanup_tree_remove_uninit_layered(depth); println!("Private tree cleaned up"); diff --git a/wallet/src/lib.rs b/wallet/src/lib.rs index 56872db..d0a9014 100644 --- a/wallet/src/lib.rs +++ b/wallet/src/lib.rs @@ -115,8 +115,7 @@ impl WalletCore { pub async fn create_new_account_public(&mut self, chain_index: ChainIndex) -> AccountId { self.storage .user_data - .generate_new_public_transaction_private_key(chain_index, self.sequencer_client.clone()) - .await + .generate_new_public_transaction_private_key(chain_index) } pub fn create_new_account_private(&mut self, chain_index: ChainIndex) -> AccountId { From 01f7acb4d68e982d6d071092dbbdcc733b01873a Mon Sep 17 00:00:00 2001 From: Pravdyvy Date: Mon, 8 Dec 2025 12:13:03 +0200 Subject: [PATCH 09/10] fix: remove unused enum --- key_protocol/src/key_management/key_tree/mod.rs | 12 +----------- 1 file changed, 1 insertion(+), 11 deletions(-) diff --git a/key_protocol/src/key_management/key_tree/mod.rs b/key_protocol/src/key_management/key_tree/mod.rs index d964af4..9b5014f 100644 --- a/key_protocol/src/key_management/key_tree/mod.rs +++ b/key_protocol/src/key_management/key_tree/mod.rs @@ -4,7 +4,7 @@ use std::{ }; use anyhow::Result; -use common::{error::SequencerClientError, sequencer_client::SequencerClient}; +use common::sequencer_client::SequencerClient; use serde::{Deserialize, Serialize}; use crate::key_management::{ @@ -29,16 +29,6 @@ pub struct KeyTree { pub type KeyTreePublic = KeyTree; pub type KeyTreePrivate = KeyTree; -#[derive(thiserror::Error, Debug)] -pub enum KeyTreeGenerationError { - #[error("Parent chain id {0} not present in tree")] - ParentChainIdNotFound(ChainIndex), - #[error("Parent or left relative of {0} is not initialized")] - PredecesorsNotInitialized(ChainIndex), - #[error("Sequencer client error {0:#?}")] - SequencerClientError(#[from] SequencerClientError), -} - impl KeyTree { pub fn new(seed: &SeedHolder) -> Self { let seed_fit: [u8; 64] = seed From eed485bd2ca42db1193a8b560418e6ac89bf80e1 Mon Sep 17 00:00:00 2001 From: Pravdyvy Date: Tue, 9 Dec 2025 11:10:40 +0200 Subject: [PATCH 10/10] fix: tomil fix --- key_protocol/Cargo.toml | 1 - 1 file changed, 1 deletion(-) diff --git a/key_protocol/Cargo.toml b/key_protocol/Cargo.toml index 24b92c0..103a1de 100644 --- a/key_protocol/Cargo.toml +++ b/key_protocol/Cargo.toml @@ -23,4 +23,3 @@ path = "../common" [dependencies.nssa] path = "../nssa" -features = ["no_docker"]