From 926a292c9c34c61cc7634f8f9bf56e12a8132959 Mon Sep 17 00:00:00 2001 From: Pravdyvy Date: Thu, 4 Dec 2025 16:49:10 +0200 Subject: [PATCH] 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 {