From b81aafbc2334ba6c45c3aa0347ddf4789957d074 Mon Sep 17 00:00:00 2001 From: Pravdyvy Date: Wed, 3 Dec 2025 13:10:07 +0200 Subject: [PATCH 1/4] feat: UX improvements --- integration_tests/src/test_suite_map.rs | 85 ++++--------- .../src/key_management/key_tree/mod.rs | 114 ++++++++++++++---- key_protocol/src/key_protocol_core/mod.rs | 37 ++++-- wallet/src/cli/account.rs | 14 ++- wallet/src/cli/config.rs | 9 -- wallet/src/cli/mod.rs | 48 +++++--- wallet/src/lib.rs | 10 +- wallet/src/main.rs | 19 +-- 8 files changed, 192 insertions(+), 144 deletions(-) diff --git a/integration_tests/src/test_suite_map.rs b/integration_tests/src/test_suite_map.rs index 8eb849f..d5bed4c 100644 --- a/integration_tests/src/test_suite_map.rs +++ b/integration_tests/src/test_suite_map.rs @@ -87,9 +87,7 @@ pub fn prepare_function_map() -> HashMap { #[nssa_integration_test] pub async fn test_success_move_to_another_account() { info!("########## test_success_move_to_another_account ##########"); - let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: ChainIndex::root(), - })); + let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public { cci: None })); let wallet_config = fetch_config().await.unwrap(); @@ -293,9 +291,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: definition_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Public { - cci: ChainIndex::root(), - }, + NewSubcommand::Public { cci: None }, ))) .await .unwrap() @@ -306,9 +302,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: supply_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Public { - cci: ChainIndex::root(), - }, + NewSubcommand::Public { cci: None }, ))) .await .unwrap() @@ -319,9 +313,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: recipient_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Public { - cci: ChainIndex::root(), - }, + NewSubcommand::Public { cci: None }, ))) .await .unwrap() @@ -453,9 +445,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: definition_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Public { - cci: ChainIndex::root(), - }, + NewSubcommand::Public { cci: None }, ))) .await .unwrap() @@ -466,9 +456,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: supply_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Private { - cci: ChainIndex::root(), - }, + NewSubcommand::Private { cci: None }, ))) .await .unwrap() @@ -479,9 +467,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: recipient_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Private { - cci: ChainIndex::root(), - }, + NewSubcommand::Private { cci: None }, ))) .await .unwrap() @@ -614,9 +600,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: definition_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Public { - cci: ChainIndex::root(), - }, + NewSubcommand::Public { cci: None }, ))) .await .unwrap() @@ -627,9 +611,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: supply_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Private { - cci: ChainIndex::root(), - }, + NewSubcommand::Private { cci: None }, ))) .await .unwrap() @@ -640,9 +622,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: recipient_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Private { - cci: ChainIndex::root(), - }, + NewSubcommand::Private { cci: None }, ))) .await .unwrap() @@ -756,9 +736,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: definition_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Public { - cci: ChainIndex::root(), - }, + NewSubcommand::Public { cci: None }, ))) .await .unwrap() @@ -769,9 +747,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: supply_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Public { - cci: ChainIndex::root(), - }, + NewSubcommand::Public { cci: None }, ))) .await .unwrap() @@ -782,9 +758,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: recipient_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Private { - cci: ChainIndex::root(), - }, + NewSubcommand::Private { cci: None }, ))) .await .unwrap() @@ -898,9 +872,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: definition_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Public { - cci: ChainIndex::root(), - }, + NewSubcommand::Public { cci: None }, ))) .await .unwrap() @@ -911,9 +883,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: supply_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Private { - cci: ChainIndex::root(), - }, + NewSubcommand::Private { cci: None }, ))) .await .unwrap() @@ -924,9 +894,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: recipient_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Public { - cci: ChainIndex::root(), - }, + NewSubcommand::Public { cci: None }, ))) .await .unwrap() @@ -1127,9 +1095,8 @@ pub fn prepare_function_map() -> HashMap { ); let from: AccountId = ACC_SENDER_PRIVATE.parse().unwrap(); - let command = Command::Account(AccountSubcommand::New(NewSubcommand::Private { - cci: ChainIndex::root(), - })); + let command = + Command::Account(AccountSubcommand::New(NewSubcommand::Private { cci: None })); let sub_ret = wallet::cli::execute_subcommand(command).await.unwrap(); let SubcommandReturnValue::RegisterAccount { @@ -1490,9 +1457,7 @@ pub fn prepare_function_map() -> HashMap { #[nssa_integration_test] pub async fn test_authenticated_transfer_initialize_function() { info!("########## test initialize account for authenticated transfer ##########"); - let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: ChainIndex::root(), - })); + let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public { cci: None })); let SubcommandReturnValue::RegisterAccount { account_id } = wallet::cli::execute_subcommand(command).await.unwrap() else { @@ -1590,9 +1555,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: winner_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Private { - cci: ChainIndex::root(), - }, + NewSubcommand::Private { cci: None }, ))) .await .unwrap() @@ -1676,7 +1639,7 @@ pub fn prepare_function_map() -> HashMap { let from: AccountId = ACC_SENDER_PRIVATE.parse().unwrap(); let command = Command::Account(AccountSubcommand::New(NewSubcommand::Private { - cci: ChainIndex::root(), + cci: Some(ChainIndex::root()), })); let sub_ret = wallet::cli::execute_subcommand(command).await.unwrap(); @@ -1688,7 +1651,7 @@ pub fn prepare_function_map() -> HashMap { }; let command = Command::Account(AccountSubcommand::New(NewSubcommand::Private { - cci: ChainIndex::from_str("/0").unwrap(), + cci: Some(ChainIndex::from_str("/0").unwrap()), })); let sub_ret = wallet::cli::execute_subcommand(command).await.unwrap(); @@ -1726,7 +1689,7 @@ pub fn prepare_function_map() -> HashMap { let from: AccountId = ACC_SENDER.parse().unwrap(); let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: ChainIndex::root(), + cci: Some(ChainIndex::root()), })); let sub_ret = wallet::cli::execute_subcommand(command).await.unwrap(); @@ -1738,7 +1701,7 @@ pub fn prepare_function_map() -> HashMap { }; let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: ChainIndex::from_str("/0").unwrap(), + cci: Some(ChainIndex::from_str("/0").unwrap()), })); let sub_ret = wallet::cli::execute_subcommand(command).await.unwrap(); diff --git a/key_protocol/src/key_management/key_tree/mod.rs b/key_protocol/src/key_management/key_tree/mod.rs index 3d55516..beb0f95 100644 --- a/key_protocol/src/key_management/key_tree/mod.rs +++ b/key_protocol/src/key_management/key_tree/mod.rs @@ -1,5 +1,5 @@ use std::{ - collections::{BTreeMap, HashMap}, + collections::{BTreeMap, HashMap, VecDeque}, sync::Arc, }; @@ -20,6 +20,8 @@ pub mod keys_private; pub mod keys_public; pub mod traits; +pub const DEPTH_SOFT_CAP: u32 = 20; + #[derive(Debug, Serialize, Deserialize, Clone)] pub struct KeyTree { pub key_map: BTreeMap, @@ -101,10 +103,13 @@ impl KeyTree { } } - 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); @@ -113,9 +118,43 @@ impl KeyTree { 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)) + } + + fn have_child_slot_capped(&self, cci: &ChainIndex) -> bool { + let depth = cci.depth(); + + self.find_next_last_child_of_id(cci) + .map(|inn| inn + 1 + depth < DEPTH_SOFT_CAP) + .unwrap_or(false) + } + + pub fn search_new_parent_capped(&self) -> Option { + let mut parent_list = VecDeque::new(); + parent_list.push_front(ChainIndex::root()); + + let mut search_res = None; + + while let Some(next_parent) = parent_list.pop_back() { + if self.have_child_slot_capped(&next_parent) { + search_res = Some(next_parent); + break; + } else { + let last_child = self.find_next_last_child_of_id(&next_parent)?; + + for id in 0..last_child { + parent_list.push_front(next_parent.nth_child(id)); + } + } + } + + search_res + } + + pub fn generate_new_node_capped(&mut self) -> Option<(nssa::AccountId, ChainIndex)> { + self.generate_new_node(&self.search_new_parent_capped()?) } pub fn get_node(&self, account_id: nssa::AccountId) -> Option<&N> { @@ -150,7 +189,7 @@ impl KeyTree { let mut id_stack = vec![ChainIndex::root()]; while let Some(curr_id) = id_stack.pop() { - self.generate_new_node(curr_id.clone()); + self.generate_new_node(&curr_id); let mut next_id = curr_id.nth_child(0); @@ -268,7 +307,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 +320,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 +346,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 +359,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 +376,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 +389,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 +402,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 +416,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 +430,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 +444,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 +458,35 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 1); } + + #[test] + fn test_tree_balancing_automatic() { + let seed_holder = seed_holder_for_tests(); + + let mut tree = KeyTreePublic::new(&seed_holder); + + for _ in 0..19 { + tree.generate_new_node_capped().unwrap(); + } + + let next_suitable_parent = tree.search_new_parent_capped().unwrap(); + + assert_eq!(next_suitable_parent, ChainIndex::from_str("/0").unwrap()); + + for _ in 0..18 { + tree.generate_new_node_capped().unwrap(); + } + + let next_suitable_parent = tree.search_new_parent_capped().unwrap(); + + assert_eq!(next_suitable_parent, ChainIndex::from_str("/1").unwrap()); + + for _ in 0..17 { + tree.generate_new_node_capped().unwrap(); + } + + let next_suitable_parent = tree.search_new_parent_capped().unwrap(); + + assert_eq!(next_suitable_parent, ChainIndex::from_str("/2").unwrap()); + } } diff --git a/key_protocol/src/key_protocol_core/mod.rs b/key_protocol/src/key_protocol_core/mod.rs index ac5ee48..41a686b 100644 --- a/key_protocol/src/key_protocol_core/mod.rs +++ b/key_protocol/src/key_protocol_core/mod.rs @@ -89,9 +89,18 @@ impl NSSAUserData { /// Returns the account_id of new account pub fn generate_new_public_transaction_private_key( &mut self, - parent_cci: ChainIndex, - ) -> nssa::AccountId { - self.public_key_tree.generate_new_node(parent_cci).unwrap() + parent_cci: Option, + ) -> (nssa::AccountId, ChainIndex) { + match parent_cci { + Some(parent_cci) => self + .public_key_tree + .generate_new_node(&parent_cci) + .expect("Parent must be present in a tree"), + None => self + .public_key_tree + .generate_new_node_capped() + .expect("No slots left"), + } } /// Returns the signing key for public transaction signatures @@ -113,9 +122,18 @@ impl NSSAUserData { /// Returns the account_id of new account pub fn generate_new_privacy_preserving_transaction_key_chain( &mut self, - parent_cci: ChainIndex, - ) -> nssa::AccountId { - self.private_key_tree.generate_new_node(parent_cci).unwrap() + parent_cci: Option, + ) -> (nssa::AccountId, ChainIndex) { + match parent_cci { + Some(parent_cci) => self + .private_key_tree + .generate_new_node(&parent_cci) + .expect("Parent must be present in a tree"), + None => self + .private_key_tree + .generate_new_node_capped() + .expect("No slots left"), + } } /// Returns the signing key for public transaction signatures @@ -169,10 +187,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 (account_id_pub, _) = user_data.generate_new_public_transaction_private_key(None); + let (account_id_private, _) = + user_data.generate_new_privacy_preserving_transaction_key_chain(None); let is_private_key_generated = user_data .get_pub_account_signing_key(&account_id_pub) diff --git a/wallet/src/cli/account.rs b/wallet/src/cli/account.rs index 5b23b2b..b3c8d5c 100644 --- a/wallet/src/cli/account.rs +++ b/wallet/src/cli/account.rs @@ -98,13 +98,13 @@ pub enum NewSubcommand { Public { #[arg(long)] /// Chain index of a parent node - cci: ChainIndex, + cci: Option, }, /// Register new private account Private { #[arg(long)] /// Chain index of a parent node - cci: ChainIndex, + cci: Option, }, } @@ -115,9 +115,11 @@ impl WalletSubcommand for NewSubcommand { ) -> Result { match self { NewSubcommand::Public { cci } => { - let account_id = wallet_core.create_new_account_public(cci); + let (account_id, chain_index) = wallet_core.create_new_account_public(cci); - println!("Generated new account with account_id Public/{account_id}"); + println!( + "Generated new account with account_id Public/{account_id} at path {chain_index}" + ); let path = wallet_core.store_persistent_data().await?; @@ -126,7 +128,7 @@ impl WalletSubcommand for NewSubcommand { Ok(SubcommandReturnValue::RegisterAccount { account_id }) } NewSubcommand::Private { cci } => { - let account_id = wallet_core.create_new_account_private(cci); + let (account_id, chain_index) = wallet_core.create_new_account_private(cci); let (key, _) = wallet_core .storage @@ -135,7 +137,7 @@ impl WalletSubcommand for NewSubcommand { .unwrap(); println!( - "Generated new account with account_id Private/{}", + "Generated new account with account_id Private/{} at path {chain_index}", account_id.to_bytes().to_base58() ); println!("With npk {}", hex::encode(key.nullifer_public_key.0)); diff --git a/wallet/src/cli/config.rs b/wallet/src/cli/config.rs index 68670af..3026c29 100644 --- a/wallet/src/cli/config.rs +++ b/wallet/src/cli/config.rs @@ -9,10 +9,6 @@ use crate::{ /// Represents generic config CLI subcommand #[derive(Subcommand, Debug, Clone)] pub enum ConfigSubcommand { - /// Command to explicitly setup config and storage - /// - /// Does nothing in case if both already present - Setup {}, /// Getter of config fields Get { key: String }, /// Setter of config fields @@ -27,11 +23,6 @@ impl WalletSubcommand for ConfigSubcommand { wallet_core: &mut WalletCore, ) -> Result { match self { - ConfigSubcommand::Setup {} => { - let path = wallet_core.store_persistent_data().await?; - - println!("Stored persistent accounts at {path:#?}"); - } ConfigSubcommand::Get { key } => match key.as_str() { "all" => { let config_str = diff --git a/wallet/src/cli/mod.rs b/wallet/src/cli/mod.rs index 1f0e53b..e9f76bd 100644 --- a/wallet/src/cli/mod.rs +++ b/wallet/src/cli/mod.rs @@ -1,4 +1,4 @@ -use std::sync::Arc; +use std::{io::Write, sync::Arc}; use anyhow::Result; use clap::{Parser, Subcommand}; @@ -16,7 +16,7 @@ use crate::{ token::TokenProgramAgnosticSubcommand, }, }, - helperfunctions::{fetch_config, parse_block_range}, + helperfunctions::{fetch_config, fetch_persistent_storage, parse_block_range}, }; pub mod account; @@ -54,25 +54,12 @@ pub enum Command { /// Command to setup config, get and set config fields #[command(subcommand)] Config(ConfigSubcommand), -} - -/// Represents overarching CLI command for a wallet with setup included -#[derive(Debug, Subcommand, Clone)] -#[clap(about)] -pub enum OverCommand { - /// Represents CLI command for a wallet - #[command(subcommand)] - Command(Command), - /// Setup of a storage. Initializes rots for public and private trees from `password`. - Setup { - #[arg(short, long)] - password: String, - }, + /// Restoring keys from given password at given `depth` + /// /// !!!WARNING!!! will rewrite current storage RestoreKeys { #[arg(short, long)] - password: String, - #[arg(short, long)] + /// Indicates, how deep in tree accounts may be. Affects command complexity. depth: u32, }, } @@ -91,7 +78,7 @@ pub struct Args { pub continuous_run: bool, /// Wallet command #[command(subcommand)] - pub command: Option, + pub command: Option, } #[derive(Debug, Clone)] @@ -104,6 +91,13 @@ pub enum SubcommandReturnValue { } pub async fn execute_subcommand(command: Command) -> Result { + if fetch_persistent_storage().await.is_err() { + println!("Persistent storage not found, need to execute setup"); + + let password = read_password_from_stdin()?; + execute_setup(password).await?; + } + let wallet_config = fetch_config().await?; let mut wallet_core = WalletCore::start_from_config_update_chain(wallet_config).await?; @@ -164,6 +158,12 @@ pub async fn execute_subcommand(command: Command) -> Result { + let password = read_password_from_stdin()?; + execute_keys_restoration(password, depth).await?; + + SubcommandReturnValue::Empty + } }; Ok(subcommand_ret) @@ -197,6 +197,16 @@ pub async fn execute_continuous_run() -> Result<()> { } } +pub fn read_password_from_stdin() -> Result { + let mut password = String::new(); + + print!("Input password: "); + std::io::stdout().flush()?; + std::io::stdin().read_line(&mut password)?; + + Ok(password.trim().to_string()) +} + pub async fn execute_setup(password: String) -> Result<()> { let config = fetch_config().await?; let wallet_core = WalletCore::start_from_config_new_storage(config.clone(), password).await?; diff --git a/wallet/src/lib.rs b/wallet/src/lib.rs index f79d947..b2797ed 100644 --- a/wallet/src/lib.rs +++ b/wallet/src/lib.rs @@ -108,13 +108,19 @@ impl WalletCore { Ok(config_path) } - pub fn create_new_account_public(&mut self, chain_index: ChainIndex) -> AccountId { + pub fn create_new_account_public( + &mut self, + chain_index: Option, + ) -> (AccountId, ChainIndex) { self.storage .user_data .generate_new_public_transaction_private_key(chain_index) } - pub fn create_new_account_private(&mut self, chain_index: ChainIndex) -> AccountId { + pub fn create_new_account_private( + &mut self, + chain_index: Option, + ) -> (AccountId, ChainIndex) { self.storage .user_data .generate_new_privacy_preserving_transaction_key_chain(chain_index) diff --git a/wallet/src/main.rs b/wallet/src/main.rs index 00aa6d0..3fcac30 100644 --- a/wallet/src/main.rs +++ b/wallet/src/main.rs @@ -1,10 +1,7 @@ use anyhow::Result; use clap::{CommandFactory as _, Parser as _}; use tokio::runtime::Builder; -use wallet::cli::{ - Args, OverCommand, execute_continuous_run, execute_keys_restoration, execute_setup, - execute_subcommand, -}; +use wallet::cli::{Args, execute_continuous_run, execute_subcommand}; pub const NUM_THREADS: usize = 2; @@ -26,17 +23,9 @@ fn main() -> Result<()> { env_logger::init(); runtime.block_on(async move { - if let Some(over_command) = args.command { - match over_command { - OverCommand::Command(command) => { - let _output = execute_subcommand(command).await?; - Ok(()) - } - OverCommand::RestoreKeys { password, depth } => { - execute_keys_restoration(password, depth).await - } - OverCommand::Setup { password } => execute_setup(password).await, - } + if let Some(command) = args.command { + let _output = execute_subcommand(command).await?; + Ok(()) } else if args.continuous_run { execute_continuous_run().await } else { From 47ff7f0b64cc0e1c6e936fd8a02eaef993ea38fe Mon Sep 17 00:00:00 2001 From: Pravdyvy Date: Tue, 9 Dec 2025 11:09:15 +0200 Subject: [PATCH 2/4] fix: layered automatization --- key_protocol/Cargo.toml | 1 - .../src/key_management/key_tree/mod.rs | 74 +++++++------------ key_protocol/src/key_protocol_core/mod.rs | 8 +- 3 files changed, 31 insertions(+), 52 deletions(-) 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"] diff --git a/key_protocol/src/key_management/key_tree/mod.rs b/key_protocol/src/key_management/key_tree/mod.rs index 9f72ecf..324d6fe 100644 --- a/key_protocol/src/key_management/key_tree/mod.rs +++ b/key_protocol/src/key_management/key_tree/mod.rs @@ -1,5 +1,5 @@ use std::{ - collections::{BTreeMap, HashMap, VecDeque}, + collections::{BTreeMap, HashMap}, sync::Arc, }; @@ -122,38 +122,34 @@ impl KeyTree { Some((account_id, next_cci)) } - fn have_child_slot_capped(&self, cci: &ChainIndex) -> bool { - let depth = cci.depth(); + fn find_next_slot_layered(&self) -> ChainIndex { + let mut depth = 1; - self.find_next_last_child_of_id(cci) - .map(|inn| inn + 1 + depth < DEPTH_SOFT_CAP) - .unwrap_or(false) - } - - pub fn search_new_parent_capped(&self) -> Option { - let mut parent_list = VecDeque::new(); - parent_list.push_front(ChainIndex::root()); - - let mut search_res = None; - - while let Some(next_parent) = parent_list.pop_back() { - if self.have_child_slot_capped(&next_parent) { - search_res = Some(next_parent); - break; - } else { - let last_child = self.find_next_last_child_of_id(&next_parent)?; - - for id in 0..last_child { - parent_list.push_front(next_parent.nth_child(id)); + 'outer: loop { + for chain_id in ChainIndex::chain_ids_at_depth(depth) { + if self.key_map.get(&chain_id).is_none() { + break 'outer chain_id; } } + depth += 1; } - - search_res } - pub fn generate_new_node_capped(&mut self) -> Option<(nssa::AccountId, ChainIndex)> { - self.generate_new_node(&self.search_new_parent_capped()?) + pub fn fill_node(&mut self, chain_index: &ChainIndex) -> Option<(nssa::AccountId, ChainIndex)> { + let parent_keys = self.key_map.get(&chain_index.parent()?)?; + let child_id = *chain_index.chain().last()?; + + let child_keys = parent_keys.nth_child(child_id); + let account_id = child_keys.account_id(); + + self.key_map.insert(chain_index.clone(), child_keys); + self.account_id_map.insert(account_id, chain_index.clone()); + + Some((account_id, chain_index.clone())) + } + + pub fn generate_new_node_layered(&mut self) -> Option<(nssa::AccountId, ChainIndex)> { + self.fill_node(&self.find_next_slot_layered()) } pub fn get_node(&self, account_id: nssa::AccountId) -> Option<&N> { @@ -524,29 +520,13 @@ mod tests { let mut tree = KeyTreePublic::new(&seed_holder); - for _ in 0..19 { - tree.generate_new_node_capped().unwrap(); - } + let next_slot = tree.find_next_slot_layered(); - let next_suitable_parent = tree.search_new_parent_capped().unwrap(); + println!("NEXT SLOT {next_slot}"); - assert_eq!(next_suitable_parent, ChainIndex::from_str("/0").unwrap()); + let (acc_id, chain_id) = tree.generate_new_node_layered().unwrap(); - for _ in 0..18 { - tree.generate_new_node_capped().unwrap(); - } - - let next_suitable_parent = tree.search_new_parent_capped().unwrap(); - - assert_eq!(next_suitable_parent, ChainIndex::from_str("/1").unwrap()); - - for _ in 0..17 { - tree.generate_new_node_capped().unwrap(); - } - - let next_suitable_parent = tree.search_new_parent_capped().unwrap(); - - assert_eq!(next_suitable_parent, ChainIndex::from_str("/2").unwrap()); + println!("NEXT ACC {acc_id} at {chain_id}"); } #[test] diff --git a/key_protocol/src/key_protocol_core/mod.rs b/key_protocol/src/key_protocol_core/mod.rs index ce41d38..b46c46c 100644 --- a/key_protocol/src/key_protocol_core/mod.rs +++ b/key_protocol/src/key_protocol_core/mod.rs @@ -98,8 +98,8 @@ impl NSSAUserData { .expect("Parent must be present in a tree"), None => self .public_key_tree - .generate_new_node_capped() - .expect("No slots left"), + .generate_new_node_layered() + .expect("Search for new node slot failed"), } } @@ -131,8 +131,8 @@ impl NSSAUserData { .expect("Parent must be present in a tree"), None => self .private_key_tree - .generate_new_node_capped() - .expect("No slots left"), + .generate_new_node_layered() + .expect("Search for new node slot failed"), } } From e1eff074783a639893984b3597f0da2cdbd29a69 Mon Sep 17 00:00:00 2001 From: Pravdyvy Date: Tue, 9 Dec 2025 11:51:44 +0200 Subject: [PATCH 3/4] fix: layered autobalancing --- .../src/key_management/key_tree/chain_index.rs | 16 ++++++++++++++++ key_protocol/src/key_management/key_tree/mod.rs | 14 +++++++------- 2 files changed, 23 insertions(+), 7 deletions(-) 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 d2c9c3b..6dbaf9a 100644 --- a/key_protocol/src/key_management/key_tree/chain_index.rs +++ b/key_protocol/src/key_management/key_tree/chain_index.rs @@ -138,6 +138,22 @@ impl ChainIndex { cumulative_stack.into_iter().unique() } + + pub fn chain_ids_at_depth_rev(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().rev().unique() + } } #[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 324d6fe..389580b 100644 --- a/key_protocol/src/key_management/key_tree/mod.rs +++ b/key_protocol/src/key_management/key_tree/mod.rs @@ -126,8 +126,8 @@ impl KeyTree { let mut depth = 1; 'outer: loop { - for chain_id in ChainIndex::chain_ids_at_depth(depth) { - if self.key_map.get(&chain_id).is_none() { + for chain_id in ChainIndex::chain_ids_at_depth_rev(depth) { + if !self.key_map.contains_key(&chain_id) { break 'outer chain_id; } } @@ -520,13 +520,13 @@ mod tests { let mut tree = KeyTreePublic::new(&seed_holder); + for _ in 0..100 { + tree.generate_new_node_layered().unwrap(); + } + let next_slot = tree.find_next_slot_layered(); - println!("NEXT SLOT {next_slot}"); - - let (acc_id, chain_id) = tree.generate_new_node_layered().unwrap(); - - println!("NEXT ACC {acc_id} at {chain_id}"); + assert_eq!(next_slot, ChainIndex::from_str("/0/0/2/1").unwrap()); } #[test] From 45e3223d516578744ef7b41336d47df1c97b276b Mon Sep 17 00:00:00 2001 From: Pravdyvy Date: Wed, 10 Dec 2025 10:25:33 +0200 Subject: [PATCH 4/4] fix: merge fix --- integration_tests/src/test_suite_map.rs | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/integration_tests/src/test_suite_map.rs b/integration_tests/src/test_suite_map.rs index 1f0b7ec..582a093 100644 --- a/integration_tests/src/test_suite_map.rs +++ b/integration_tests/src/test_suite_map.rs @@ -604,7 +604,7 @@ pub fn prepare_function_map() -> HashMap { account_id: definition_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( NewSubcommand::Private { - cci: ChainIndex::root(), + cci: Some(ChainIndex::root()), }, ))) .await @@ -617,7 +617,7 @@ pub fn prepare_function_map() -> HashMap { account_id: supply_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( NewSubcommand::Public { - cci: ChainIndex::root(), + cci: Some(ChainIndex::root()), }, ))) .await @@ -666,8 +666,8 @@ pub fn prepare_function_map() -> HashMap { // The data of a token definition account has the following layout: // [ 0x00 || name (6 bytes) || total supply (little endian 16 bytes) ] assert_eq!( - supply_acc.data, - vec![ + supply_acc.data.as_ref(), + &[ 1, 128, 101, 5, 31, 43, 36, 97, 108, 164, 92, 25, 157, 173, 5, 14, 194, 121, 239, 84, 19, 160, 243, 47, 193, 2, 250, 247, 232, 253, 191, 232, 173, 37, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 @@ -689,7 +689,7 @@ pub fn prepare_function_map() -> HashMap { account_id: definition_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( NewSubcommand::Private { - cci: ChainIndex::root(), + cci: Some(ChainIndex::root()), }, ))) .await @@ -702,7 +702,7 @@ pub fn prepare_function_map() -> HashMap { account_id: supply_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( NewSubcommand::Private { - cci: ChainIndex::root(), + cci: Some(ChainIndex::root()), }, ))) .await @@ -756,8 +756,8 @@ pub fn prepare_function_map() -> HashMap { // The data of a token definition account has the following layout: // [ 0x00 || name (6 bytes) || total supply (little endian 16 bytes) ] assert_eq!( - definition_acc.data, - vec![ + definition_acc.data.as_ref(), + &[ 0, 65, 32, 78, 65, 77, 69, 37, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 ] ); @@ -766,8 +766,8 @@ pub fn prepare_function_map() -> HashMap { // The data of a token definition account has the following layout: // [ 0x00 || name (6 bytes) || total supply (little endian 16 bytes) ] assert_eq!( - supply_acc.data, - vec![ + supply_acc.data.as_ref(), + &[ 1, 128, 101, 5, 31, 43, 36, 97, 108, 164, 92, 25, 157, 173, 5, 14, 194, 121, 239, 84, 19, 160, 243, 47, 193, 2, 250, 247, 232, 253, 191, 232, 173, 37, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0