From e92ad2132f769f60cd4ee7b885df0bf67fc2874e Mon Sep 17 00:00:00 2001 From: Pravdyvy Date: Wed, 12 Nov 2025 08:26:25 +0200 Subject: [PATCH] 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;