diff --git a/integration_tests/src/test_suite_map.rs b/integration_tests/src/test_suite_map.rs index 7ee3ef4..40cf34f 100644 --- a/integration_tests/src/test_suite_map.rs +++ b/integration_tests/src/test_suite_map.rs @@ -2,6 +2,7 @@ use std::{ collections::HashMap, path::PathBuf, pin::Pin, + str::FromStr, time::{Duration, Instant}, }; @@ -1672,6 +1673,217 @@ pub fn prepare_function_map() -> HashMap { info!("Success!"); } + #[nssa_integration_test] + pub async fn test_keys_restoration() { + info!("########## test_keys_restoration ##########"); + let from: AccountId = ACC_SENDER_PRIVATE.parse().unwrap(); + + let command = Command::Account(AccountSubcommand::New(NewSubcommand::Private { + cci: ChainIndex::root(), + })); + + let sub_ret = wallet::cli::execute_subcommand(command).await.unwrap(); + let SubcommandReturnValue::RegisterAccount { + account_id: to_account_id1, + } = 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::cli::execute_subcommand(command).await.unwrap(); + let SubcommandReturnValue::RegisterAccount { + account_id: to_account_id2, + } = 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_account_id1.to_string(), + )), + to_npk: None, + to_ipk: None, + amount: 100, + }); + + wallet::cli::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_account_id2.to_string(), + )), + to_npk: None, + to_ipk: None, + amount: 101, + }); + + wallet::cli::execute_subcommand(command).await.unwrap(); + + let from: AccountId = ACC_SENDER.parse().unwrap(); + + let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public { + cci: ChainIndex::root(), + })); + + let sub_ret = wallet::cli::execute_subcommand(command).await.unwrap(); + let SubcommandReturnValue::RegisterAccount { + account_id: to_account_id3, + } = 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::cli::execute_subcommand(command).await.unwrap(); + let SubcommandReturnValue::RegisterAccount { + account_id: to_account_id4, + } = 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_account_id3.to_string(), + )), + to_npk: None, + to_ipk: None, + amount: 102, + }); + + wallet::cli::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_account_id4.to_string(), + )), + to_npk: None, + to_ipk: None, + amount: 103, + }); + + wallet::cli::execute_subcommand(command).await.unwrap(); + + info!("########## PREPARATION END ##########"); + + wallet::cli::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(); + + 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_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::cli::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::cli::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!"); + } + println!("{function_map:#?}"); function_map diff --git a/key_protocol/Cargo.toml b/key_protocol/Cargo.toml index a562515..103a1de 100644 --- a/key_protocol/Cargo.toml +++ b/key_protocol/Cargo.toml @@ -16,6 +16,7 @@ bip39.workspace = true hmac-sha512.workspace = true thiserror.workspace = true nssa-core = { path = "../nssa/core", features = ["host"] } +itertools.workspace = true [dependencies.common] path = "../common" 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 e46fc0f..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)] @@ -77,12 +78,66 @@ 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); ChainIndex(chain) } + + 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)] @@ -145,4 +200,83 @@ 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()) + } + + #[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 1319a5a..9b5014f 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::{ @@ -96,21 +101,23 @@ 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 parent_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 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); + 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> { @@ -129,11 +136,164 @@ impl KeyTree { self.account_id_map.insert(account_id, chain_index.clone()); self.key_map.insert(chain_index, node); } + + pub fn remove(&mut self, addr: nssa::AccountId) -> Option { + let chain_index = self.account_id_map.remove(&addr).unwrap(); + 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()]; + + while let Some(curr_id) = id_stack.pop() { + let mut next_id = curr_id.nth_child(0); + + while (next_id.depth()) < depth { + self.generate_new_node(&curr_id); + id_stack.push(next_id.clone()); + next_id = next_id.next_in_line(); + } + } + } +} + +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 + /// + /// 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() { + if let Some(node) = self.key_map.get(&curr_id) + && node.value.1 == nssa::Account::default() + && curr_id != ChainIndex::root() + { + let addr = node.account_id(); + self.remove(addr); + } + + let mut next_id = curr_id.nth_child(0); + + while (next_id.depth()) < depth { + id_stack.push(next_id.clone()); + next_id = next_id.next_in_line(); + } + } + } + + /// 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 { + /// 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, + 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.account_id(); + 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.nth_child(0); + + while (next_id.depth()) < depth { + id_stack.push(next_id.clone()); + next_id = next_id.next_in_line(); + } + } + + 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)] mod tests { - use std::str::FromStr; + use std::{collections::HashSet, str::FromStr}; use nssa::AccountId; @@ -162,7 +322,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()) @@ -170,7 +330,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 @@ -183,12 +343,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()) @@ -201,7 +361,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()) @@ -209,7 +369,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 @@ -222,7 +382,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); } @@ -239,7 +399,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 @@ -252,7 +412,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 @@ -265,7 +425,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 @@ -279,7 +439,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 @@ -293,7 +453,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 @@ -307,7 +467,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!( @@ -321,4 +481,79 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 1); } + + #[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("/1").unwrap()) + .unwrap(); + acc.value.1.balance = 2; + + let acc = tree + .key_map + .get_mut(&ChainIndex::from_str("/2").unwrap()) + .unwrap(); + acc.value.1.balance = 3; + + 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_remove_uninit_layered(10); + + 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()); + + let mut key_set = HashSet::new(); + + for key in tree.key_map.keys() { + key_set.insert(key.to_string()); + } + + assert_eq!(key_set, key_set_res); + + let acc = tree + .key_map + .get(&ChainIndex::from_str("/1").unwrap()) + .unwrap(); + assert_eq!(acc.value.1.balance, 2); + + let acc = tree + .key_map + .get(&ChainIndex::from_str("/2").unwrap()) + .unwrap(); + assert_eq!(acc.value.1.balance, 3); + + let acc = tree + .key_map + .get(&ChainIndex::from_str("/0/1").unwrap()) + .unwrap(); + assert_eq!(acc.value.1.balance, 5); + + 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 ac5ee48..fc0a393 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 @@ -169,17 +175,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 2299f9e..2f5ef04 100644 --- a/wallet/src/cli/account.rs +++ b/wallet/src/cli/account.rs @@ -113,7 +113,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/cli/mod.rs b/wallet/src/cli/mod.rs index b3d40b8..86d8210 100644 --- a/wallet/src/cli/mod.rs +++ b/wallet/src/cli/mod.rs @@ -65,6 +65,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 @@ -206,3 +213,67 @@ pub async fn execute_setup_with_auth(password: String, auth: Option) -> Ok(()) } + +pub async fn execute_keys_restoration(password: String, depth: u32) -> Result<()> { + execute_keys_restoration_with_auth(password, depth, None).await +} + +pub async fn execute_keys_restoration_with_auth( + password: String, + depth: u32, + auth: Option, +) -> Result<()> { + let config = fetch_config().await?; + let config = merge_auth_config(config, auth)?; + 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); + + 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_remove_uninit_layered(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}"); + + wallet_core.sync_to_block(last_block).await?; + + println!("Private tree clean up start"); + + wallet_core + .storage + .user_data + .private_key_tree + .cleanup_tree_remove_uninit_layered(depth); + + println!("Private tree cleaned up"); + + wallet_core.store_persistent_data().await?; + + Ok(()) +} diff --git a/wallet/src/lib.rs b/wallet/src/lib.rs index 11f4a4a..e68e6a4 100644 --- a/wallet/src/lib.rs +++ b/wallet/src/lib.rs @@ -126,7 +126,7 @@ 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) diff --git a/wallet/src/main.rs b/wallet/src/main.rs index 27d102d..1430b60 100644 --- a/wallet/src/main.rs +++ b/wallet/src/main.rs @@ -2,8 +2,8 @@ use anyhow::Result; use clap::{CommandFactory as _, Parser as _}; use tokio::runtime::Builder; use wallet::cli::{ - Args, OverCommand, execute_continuous_run_with_auth, execute_setup_with_auth, - execute_subcommand_with_auth, + Args, OverCommand, execute_continuous_run_with_auth, execute_keys_restoration_with_auth, + execute_setup_with_auth, execute_subcommand_with_auth, }; pub const NUM_THREADS: usize = 2; @@ -31,6 +31,9 @@ fn main() -> Result<()> { let _output = execute_subcommand_with_auth(command, args.auth).await?; Ok(()) } + OverCommand::RestoreKeys { password, depth } => { + execute_keys_restoration_with_auth(password, depth, args.auth).await + } OverCommand::Setup { password } => { execute_setup_with_auth(password, args.auth).await }