diff --git a/crates/common/src/lib.rs b/crates/common/src/lib.rs index 912f5a5e..c6e4d4a2 100644 --- a/crates/common/src/lib.rs +++ b/crates/common/src/lib.rs @@ -8,7 +8,5 @@ pub mod tree; #[macro_use] extern crate log; -#[cfg(feature = "test_utils")] -pub mod test_utils; #[cfg(feature = "test_utils")] pub mod transaction_builder; diff --git a/crates/common/src/test_utils.rs b/crates/common/src/test_utils.rs deleted file mode 100644 index 8005d492..00000000 --- a/crates/common/src/test_utils.rs +++ /dev/null @@ -1,275 +0,0 @@ -use crate::{ - digest::Digest, - hashchain::Hashchain, - hasher::Hasher, - operation::{ServiceChallenge, ServiceChallengeInput, SignatureBundle}, - transaction::Transaction, - tree::{ - HashchainResponse::*, InsertProof, KeyDirectoryTree, Proof, SnarkableTree, UpdateProof, - }, -}; -use anyhow::{anyhow, Result}; -use jmt::{mock::MockTreeStore, KeyHash}; -use prism_keys::{SigningKey, VerifyingKey}; -use rand::{rngs::StdRng, Rng}; -use std::{ - collections::{HashMap, HashSet}, - sync::Arc, -}; - -pub struct TestTreeState { - pub tree: KeyDirectoryTree, - pub signing_keys: HashMap, - inserted_keys: HashSet, - pub services: HashMap, -} - -#[derive(Clone)] -pub struct TestAccount { - pub id: String, - pub key_hash: KeyHash, - pub hashchain: Hashchain, -} - -#[derive(Clone)] -pub struct Service { - pub id: String, - pub sk: SigningKey, - pub vk: VerifyingKey, - pub registration: TestAccount, -} - -impl TestTreeState { - pub fn new() -> Self { - Self::default() - } - - pub fn register_service(&mut self, service_id: String) -> Service { - let service_challenge_key = SigningKey::new_ed25519(); - let service_signing_key = SigningKey::new_ed25519(); - let service_vk: VerifyingKey = service_signing_key.clone().into(); - - let mut hashchain = Hashchain::empty(); - - hashchain - .register_service( - service_id.clone(), - ServiceChallenge::from(service_challenge_key.clone()), - service_vk, - &service_signing_key, - ) - .unwrap(); - - let key_hash = KeyHash::with::(&service_id); - - Service { - id: service_id.clone(), - sk: service_challenge_key.clone(), - vk: service_challenge_key.into(), - registration: TestAccount { - id: service_id, - key_hash, - hashchain, - }, - } - } - - pub fn create_account(&mut self, id: String, service: Service) -> TestAccount { - let signing_key = SigningKey::new_ed25519(); - let vk: VerifyingKey = signing_key.clone().into(); - self.signing_keys.insert(id.clone(), signing_key.clone()); - - // Simulate some external service signing account creation credentials - let hash = Digest::hash_items(&[id.as_bytes(), service.id.as_bytes(), &vk.to_bytes()]); - let signature = service.sk.sign(&hash.to_bytes()); - - let mut hashchain = Hashchain::empty(); - hashchain - .create_account( - id.clone(), - service.id.clone(), - ServiceChallengeInput::Signed(signature), - vk, - &signing_key, - ) - .unwrap(); - - let key_hash = KeyHash::with::(&id); - - TestAccount { - id, - key_hash, - hashchain, - } - } - - pub fn insert_account(&mut self, account: TestAccount) -> Result { - if self.inserted_keys.contains(&account.id) { - return Err(anyhow!("{:?} already contained in tree", account.id)); - } - - let transaction = Transaction { - id: account.id.clone(), - entry: account.hashchain.last().unwrap().clone(), - }; - - let proof = self.tree.process_transaction(transaction)?; - if let Proof::Insert(insert_proof) = proof { - self.inserted_keys.insert(account.id); - return Ok(*insert_proof); - } - Err(anyhow!("Insert proof not returned")) - } - - pub fn update_account(&mut self, account: TestAccount) -> Result { - if !self.inserted_keys.contains(&account.id) { - return Err(anyhow!("{:?} not found in tree", account.id)); - } - - let transaction = Transaction { - id: account.id.clone(), - entry: account.hashchain.last().unwrap().clone(), - }; - - let proof = self.tree.process_transaction(transaction)?; - if let Proof::Update(update_proof) = proof { - return Ok(*update_proof); - } - Err(anyhow!("Update proof not returned")) - } - - pub fn add_key_to_account(&mut self, account: &mut TestAccount) -> Result<(), anyhow::Error> { - let signing_key_to_add = SigningKey::new_ed25519(); - let key_to_add = signing_key_to_add.into(); - - account - .hashchain - .add_key(key_to_add, self.signing_keys.get(&account.id).unwrap(), 0) - .unwrap(); - Ok(()) - } - - pub fn add_unsigned_data_to_account( - &mut self, - data: &[u8], - account: &mut TestAccount, - ) -> Result<()> { - self.add_data_to_account(data, account, None) - } - - pub fn add_signed_data_to_account( - &mut self, - data: &[u8], - account: &mut TestAccount, - ) -> Result<()> { - let random_signing_key = SigningKey::new_ed25519(); - self.add_data_to_account(data, account, Some(&random_signing_key)) - } - - fn add_data_to_account( - &mut self, - data: &[u8], - account: &mut TestAccount, - signing_key: Option<&SigningKey>, - ) -> Result<()> { - let signature_bundle = signing_key.map(|sk| SignatureBundle { - verifying_key: sk.clone().into(), - signature: sk.sign(data), - }); - - let signing_key = self.signing_keys.get(&account.id).unwrap(); - - account.hashchain.add_data(data.to_vec(), signature_bundle, signing_key, 0)?; - Ok(()) - } -} - -impl Default for TestTreeState { - fn default() -> Self { - let store = Arc::new(MockTreeStore::default()); - let tree = KeyDirectoryTree::new(store); - Self { - tree, - inserted_keys: HashSet::new(), - signing_keys: HashMap::new(), - services: HashMap::new(), - } - } -} - -pub fn create_random_insert(state: &mut TestTreeState, rng: &mut StdRng) -> InsertProof { - loop { - let random_string: String = - (0..10).map(|_| rng.sample(rand::distributions::Alphanumeric) as char).collect(); - let sk = SigningKey::new_ed25519(); - - let (_, service) = - state.services.iter().nth(rng.gen_range(0..state.services.len())).unwrap(); - - let vk: VerifyingKey = sk.clone().into(); - - // Simulate some external service signing account creation credentials - let hash = Digest::hash_items(&[ - random_string.as_bytes(), - service.id.as_bytes(), - &vk.to_bytes(), - ]); - let signature = service.sk.sign(&hash.to_bytes()); - - let key_hash = KeyHash::with::(&random_string); - - let entry = Hashchain::empty() - .create_account( - random_string.clone(), - service.id.clone(), - ServiceChallengeInput::Signed(signature), - vk, - &sk, - ) - .unwrap(); - - if !state.inserted_keys.contains(&random_string) { - let proof = state.tree.insert(key_hash, entry).expect("Insert should succeed"); - state.inserted_keys.insert(random_string.clone()); - state.signing_keys.insert(random_string, sk); - return proof; - } - } -} - -pub fn create_random_update(state: &mut TestTreeState, rng: &mut StdRng) -> UpdateProof { - if state.inserted_keys.is_empty() { - panic!("No keys have been inserted yet. Cannot perform update."); - } - - let key = state.inserted_keys.iter().nth(rng.gen_range(0..state.inserted_keys.len())).unwrap(); - - let key_hash = KeyHash::with::(key); - - let Found(mut hc, _) = state.tree.get(key_hash).unwrap() else { - panic!("No response found for key. Cannot perform update."); - }; - - let signing_key = SigningKey::new_ed25519(); - let verifying_key = signing_key.into(); - - let signer = state - .signing_keys - .get(key) - .ok_or_else(|| anyhow::anyhow!("Signing key not found for hashchain")) - .unwrap(); - - let entry = hc.add_key(verifying_key, signer, 0).unwrap(); - let transaction = Transaction { - id: key.clone(), - entry, - }; - - let Proof::Update(update_proof) = - state.tree.process_transaction(transaction).expect("Processing transaction should succeed") - else { - panic!("No update proof returned."); - }; - - *update_proof -} diff --git a/crates/common/src/transaction_builder.rs b/crates/common/src/transaction_builder.rs index d749ece4..8cba3807 100644 --- a/crates/common/src/transaction_builder.rs +++ b/crates/common/src/transaction_builder.rs @@ -1,14 +1,10 @@ -use std::{collections::HashMap, sync::Arc}; - -use jmt::{mock::MockTreeStore, KeyHash}; +use std::collections::HashMap; use crate::{ digest::Digest, - hashchain::HashchainEntry, - hasher::Hasher, + hashchain::{Hashchain, HashchainEntry}, operation::{ServiceChallenge, ServiceChallengeInput, SignatureBundle}, transaction::Transaction, - tree::{HashchainResponse::*, KeyDirectoryTree, SnarkableTree}, }; use prism_keys::{SigningKey, VerifyingKey}; enum PostCommitAction { @@ -27,10 +23,14 @@ impl UncommittedTransaction<'_> { /// Commits and returns a transaction, updating the builder. Subsequent transactions /// built with the same builder will have the correct previous hash. pub fn commit(self) -> Transaction { - self.builder - .tree - .process_transaction(self.transaction.clone()) - .expect("Processing transaction should work"); + let hc = self + .builder + .hashchains + .entry(self.transaction.id.clone()) + .or_insert_with(Hashchain::empty); + + hc.add_entry(self.transaction.entry.clone()) + .expect("Adding transaction entry to hashchain should work"); match self.post_commit_action { PostCommitAction::UpdateStorageOnly => (), @@ -54,7 +54,7 @@ impl UncommittedTransaction<'_> { pub struct TransactionBuilder { /// Simulated hashchain storage that is mutated when transactions are applied - tree: Box, + hashchains: HashMap, /// Remembers private keys of services to simulate account creation via an external service service_keys: HashMap, /// Remembers private keys of accounts to simulate actions on behalf of these accounts @@ -63,13 +63,12 @@ pub struct TransactionBuilder { impl Default for TransactionBuilder { fn default() -> Self { - let store = Arc::new(MockTreeStore::default()); - let tree = Box::new(KeyDirectoryTree::new(store)); + let hashchains = HashMap::new(); let service_keys = HashMap::new(); let account_keys = HashMap::new(); Self { - tree, + hashchains, service_keys, account_keys, } @@ -81,6 +80,10 @@ impl TransactionBuilder { Self::default() } + pub fn get_hashchain(&self, id: &str) -> Option<&Hashchain> { + self.hashchains.get(id) + } + pub fn register_service_with_random_keys(&mut self, id: &str) -> UncommittedTransaction { let random_service_challenge_key = SigningKey::new_ed25519(); let random_service_signing_key = SigningKey::new_ed25519(); @@ -111,27 +114,47 @@ impl TransactionBuilder { } } - pub fn create_account_with_random_key( + pub fn create_account_with_random_key_signed( &mut self, id: &str, service_id: &str, ) -> UncommittedTransaction { - let random_signing_key = SigningKey::new_ed25519(); - self.create_account(id, service_id, random_signing_key) + let account_signing_key = SigningKey::new_ed25519(); + self.create_account_signed(id, service_id, account_signing_key) } - pub fn create_account( + pub fn create_account_signed( &mut self, id: &str, service_id: &str, signing_key: SigningKey, ) -> UncommittedTransaction { - let Some(service_signing_key) = self.service_keys.get(service_id) else { + let Some(service_signing_key) = self.service_keys.get(service_id).cloned() else { panic!("No existing service found for {}", service_id) }; - let vk: VerifyingKey = signing_key.clone().into(); + self.create_account(id, service_id, &service_signing_key, signing_key) + } + + pub fn create_account_with_random_key( + &mut self, + id: &str, + service_id: &str, + service_signing_key: &SigningKey, + ) -> UncommittedTransaction { + let account_signing_key = SigningKey::new_ed25519(); + self.create_account(id, service_id, service_signing_key, account_signing_key) + } + + pub fn create_account( + &mut self, + id: &str, + service_id: &str, + service_signing_key: &SigningKey, + signing_key: SigningKey, + ) -> UncommittedTransaction { // Simulate some external service signing account creation credentials + let vk = signing_key.verifying_key(); let hash = Digest::hash_items(&[id.as_bytes(), service_id.as_bytes(), &vk.to_bytes()]); let signature = service_signing_key.sign(&hash.to_bytes()); @@ -190,13 +213,9 @@ impl TransactionBuilder { signing_key: &SigningKey, key_idx: usize, ) -> UncommittedTransaction { - let key_hash = KeyHash::with::(id); - - let Ok(Found(hc, _)) = self.tree.get(key_hash) else { - panic!("No existing hashchain found for {}", id) - }; + let last_hash = self.hashchains.get(id).map_or(Digest::zero(), Hashchain::last_hash); - let entry = HashchainEntry::new_add_key(key, hc.last_hash(), signing_key, key_idx); + let entry = HashchainEntry::new_add_key(key, last_hash, signing_key, key_idx); UncommittedTransaction { transaction: Transaction { @@ -227,13 +246,9 @@ impl TransactionBuilder { signing_key: &SigningKey, key_idx: usize, ) -> UncommittedTransaction { - let key_hash = KeyHash::with::(id); + let last_hash = self.hashchains.get(id).map_or(Digest::zero(), Hashchain::last_hash); - let Ok(Found(hc, _)) = self.tree.get(key_hash) else { - panic!("No existing hashchain found for {}", id) - }; - - let entry = HashchainEntry::new_revoke_key(key, hc.last_hash(), signing_key, key_idx); + let entry = HashchainEntry::new_revoke_key(key, last_hash, signing_key, key_idx); UncommittedTransaction { transaction: Transaction { @@ -245,7 +260,55 @@ impl TransactionBuilder { } } + pub fn add_randomly_signed_data( + &mut self, + id: &str, + value: Vec, + signing_key: &SigningKey, + key_idx: usize, + ) -> UncommittedTransaction { + let value_signing_key = SigningKey::new_ed25519(); + self.add_signed_data(id, value, &value_signing_key, signing_key, key_idx) + } + + pub fn add_randomly_signed_data_verified_with_root( + &mut self, + id: &str, + value: Vec, + ) -> UncommittedTransaction { + let value_signing_key = SigningKey::new_ed25519(); + self.add_signed_data_verified_with_root(id, value, &value_signing_key) + } + pub fn add_signed_data( + &mut self, + id: &str, + value: Vec, + value_signing_key: &SigningKey, + signing_key: &SigningKey, + key_idx: usize, + ) -> UncommittedTransaction { + let value_signature_bundle = SignatureBundle { + verifying_key: value_signing_key.verifying_key(), + signature: value_signing_key.sign(&value), + }; + self.add_pre_signed_data(id, value, value_signature_bundle, signing_key, key_idx) + } + + pub fn add_signed_data_verified_with_root( + &mut self, + id: &str, + value: Vec, + value_signing_key: &SigningKey, + ) -> UncommittedTransaction { + let value_signature_bundle = SignatureBundle { + verifying_key: value_signing_key.verifying_key(), + signature: value_signing_key.sign(&value), + }; + self.add_pre_signed_data_verified_with_root(id, value, value_signature_bundle) + } + + pub fn add_pre_signed_data( &mut self, id: &str, value: Vec, @@ -256,7 +319,7 @@ impl TransactionBuilder { self.add_data(id, value, Some(value_signature), signing_key, key_idx) } - pub fn add_signed_data_verified_with_root( + pub fn add_pre_signed_data_verified_with_root( &mut self, id: &str, value: Vec, @@ -304,19 +367,10 @@ impl TransactionBuilder { signing_key: &SigningKey, key_idx: usize, ) -> UncommittedTransaction { - let key_hash = KeyHash::with::(id); + let last_hash = self.hashchains.get(id).map_or(Digest::zero(), Hashchain::last_hash); - let Ok(Found(hc, _)) = self.tree.get(key_hash) else { - panic!("No existing hashchain found for {}", id) - }; - - let entry = HashchainEntry::new_add_data( - data, - data_signature, - hc.last_hash(), - signing_key, - key_idx, - ); + let entry = + HashchainEntry::new_add_data(data, data_signature, last_hash, signing_key, key_idx); UncommittedTransaction { transaction: Transaction { diff --git a/crates/common/src/tree/mod.rs b/crates/common/src/tree/mod.rs index aae20067..f53806a3 100644 --- a/crates/common/src/tree/mod.rs +++ b/crates/common/src/tree/mod.rs @@ -20,113 +20,169 @@ pub enum HashchainResponse { #[cfg(all(test, feature = "test_utils"))] mod tests { - use super::*; - use crate::{test_utils::TestTreeState, tree::HashchainResponse::*}; - use jmt::KeyHash; + use std::sync::Arc; + + use jmt::{mock::MockTreeStore, KeyHash}; use prism_keys::SigningKey; - use crate::hasher::Hasher; + use super::{HashchainResponse::*, *}; + use crate::{digest::Digest, hasher::Hasher, transaction_builder::TransactionBuilder}; #[test] fn test_insert_and_get() { - let mut tree_state = TestTreeState::default(); - let service = tree_state.register_service("service_1".to_string()); - let account = tree_state.create_account("key_1".to_string(), service.clone()); + let mut tree = KeyDirectoryTree::new(Arc::new(MockTreeStore::default())); + let mut tx_builder = TransactionBuilder::new(); - let insert_proof = tree_state.insert_account(service.registration.clone()).unwrap(); + let service_tx = tx_builder.register_service_with_random_keys("service_1").commit(); + let Proof::Insert(insert_proof) = tree.process_transaction(service_tx).unwrap() else { + panic!("Processing transaction did not return the expected insert proof"); + }; assert!(insert_proof.verify().is_ok()); - let insert_proof = tree_state.insert_account(account.clone()).unwrap(); + let account_tx = + tx_builder.create_account_with_random_key_signed("acc_1", "service_1").commit(); + + let Proof::Insert(insert_proof) = tree.process_transaction(account_tx).unwrap() else { + panic!("Processing transaction did not return the expected insert proof"); + }; assert!(insert_proof.verify().is_ok()); - let Found(hashchain, membership_proof) = tree_state.tree.get(account.key_hash).unwrap() + let Found(hashchain, membership_proof) = + tree.get(KeyHash::with::("acc_1")).unwrap() else { panic!("Expected hashchain to be found, but was not found.") }; - assert_eq!(hashchain, account.hashchain); + let test_hashchain = + tx_builder.get_hashchain("acc_1").expect("Getting builder hashchain should work"); + + assert_eq!(&hashchain, test_hashchain); assert!(membership_proof.verify().is_ok()); } #[test] fn test_insert_for_nonexistent_service_fails() { - let mut tree_state = TestTreeState::default(); - let service = tree_state.register_service("service_1".to_string()); - let account = tree_state.create_account("key_1".to_string(), service.clone()); + let mut tree = KeyDirectoryTree::new(Arc::new(MockTreeStore::default())); + let mut tx_builder = TransactionBuilder::new(); - let insert_proof = tree_state.insert_account(account.clone()); - assert!(insert_proof.is_err()); + let service_signing_key = SigningKey::new_ed25519(); + + let invalid_account_tx = tx_builder + .create_account_with_random_key( + "acc_1", + "service_id_that_does_not_exist", + &service_signing_key, + ) + .build(); + + let insertion_result = tree.process_transaction(invalid_account_tx); + assert!(insertion_result.is_err()); } #[test] fn test_insert_with_invalid_service_challenge_fails() { - let mut tree_state = TestTreeState::default(); - let service = tree_state.register_service("service_1".to_string()); + let mut tree = KeyDirectoryTree::new(Arc::new(MockTreeStore::default())); + let mut tx_builder = TransactionBuilder::new(); + + let service_tx = tx_builder.register_service_with_random_keys("service_1").commit(); + + // The correct way was to use the key from service registration, + // but here we want things to break + let incorrect_service_signing_key = SigningKey::new_ed25519(); - let mut falsified_service = service.clone(); - falsified_service.sk = SigningKey::new_ed25519(); + let initial_acc_signing_key = SigningKey::new_ed25519(); - let account = tree_state.create_account("key_1".to_string(), falsified_service.clone()); + let acc_with_invalid_challenge_tx = tx_builder + .create_account( + "key_1", + "service_1", + &incorrect_service_signing_key, + initial_acc_signing_key, + ) + .build(); - let insert_proof = tree_state.insert_account(service.registration.clone()).unwrap(); + let Proof::Insert(insert_proof) = tree.process_transaction(service_tx).unwrap() else { + panic!("Processing service registration failed") + }; assert!(insert_proof.verify().is_ok()); - let insert_proof = tree_state.insert_account(account.clone()); - assert!(insert_proof.is_err()); + let create_account_result = tree.process_transaction(acc_with_invalid_challenge_tx); + assert!(create_account_result.is_err()); } #[test] fn test_insert_duplicate_key() { - let mut tree_state = TestTreeState::default(); - let service = tree_state.register_service("service_1".to_string()); - let account = tree_state.create_account("key_1".to_string(), service.clone()); + let mut tree = KeyDirectoryTree::new(Arc::new(MockTreeStore::default())); + let mut tx_builder = TransactionBuilder::new(); + + let service_tx = tx_builder.register_service_with_random_keys("service_1").commit(); + let account_tx = + tx_builder.create_account_with_random_key_signed("acc_1", "service_1").commit(); + let account_with_same_id_tx = + tx_builder.create_account_with_random_key_signed("acc_1", "service_1").build(); - let insert_proof = tree_state.insert_account(service.registration.clone()).unwrap(); + let Proof::Insert(insert_proof) = tree.process_transaction(service_tx).unwrap() else { + panic!("Processing service registration failed") + }; assert!(insert_proof.verify().is_ok()); - tree_state.insert_account(account.clone()).unwrap(); + let Proof::Insert(insert_proof) = tree.process_transaction(account_tx).unwrap() else { + panic!("Processing Account creation failed") + }; + assert!(insert_proof.verify().is_ok()); - let result = tree_state.insert_account(account.clone()); - assert!(result.is_err()); + let create_acc_with_same_id_result = tree.process_transaction(account_with_same_id_tx); + assert!(create_acc_with_same_id_result.is_err()); } #[test] fn test_update_existing_key() { - let mut tree_state = TestTreeState::default(); + let mut tree = KeyDirectoryTree::new(Arc::new(MockTreeStore::default())); + let mut tx_builder = TransactionBuilder::new(); - let service = tree_state.register_service("service_1".to_string()); - let mut account = tree_state.create_account("key_1".to_string(), service.clone()); - tree_state.insert_account(service.registration.clone()).unwrap(); - tree_state.insert_account(account.clone()).unwrap(); + let service_tx = tx_builder.register_service_with_random_keys("service_1").commit(); + let acc_tx = + tx_builder.create_account_with_random_key_signed("acc_1", "service_1").commit(); - // Add a new key - tree_state.add_key_to_account(&mut account).unwrap(); + tree.process_transaction(service_tx).unwrap(); + tree.process_transaction(acc_tx).unwrap(); - // Update the account using the correct key index - let update_proof = tree_state.update_account(account.clone()).unwrap(); + let key_tx = tx_builder.add_random_key_verified_with_root("acc_1").commit(); + + let Proof::Update(update_proof) = tree.process_transaction(key_tx).unwrap() else { + panic!("Processing key update failed") + }; assert!(update_proof.verify().is_ok()); - let get_result = tree_state.tree.get(account.key_hash); - assert!(matches!(get_result.unwrap(), Found(hc, _) if hc == account.hashchain)); + let get_result = tree.get(KeyHash::with::("acc_1")).unwrap(); + let test_hashchain = tx_builder.get_hashchain("acc_1").unwrap(); + + assert!(matches!(get_result, Found(hc, _) if &hc == test_hashchain)); } #[test] fn test_update_non_existing_key() { - let mut tree_state = TestTreeState::default(); - let service = tree_state.register_service("service_1".to_string()); - let account = tree_state.create_account("key_1".to_string(), service.clone()); - tree_state.insert_account(service.registration.clone()).unwrap(); + let mut tree = KeyDirectoryTree::new(Arc::new(MockTreeStore::default())); + let mut tx_builder = TransactionBuilder::new(); + + let service_tx = tx_builder.register_service_with_random_keys("service_1").commit(); + + tree.process_transaction(service_tx).unwrap(); + + // This is a signing key not known to the storage yet + let random_signing_key = SigningKey::new_ed25519(); + // This transaction shall be invalid, because it is signed with an unknown key + let invalid_key_tx = tx_builder.add_random_key("acc_1", &random_signing_key, 0).build(); - let result = tree_state.update_account(account); + let result = tree.process_transaction(invalid_key_tx); assert!(result.is_err()); } #[test] fn test_get_non_existing_key() { - let tree_state = TestTreeState::default(); - let key = KeyHash::with::(b"non_existing_key"); + let tree = KeyDirectoryTree::new(Arc::new(MockTreeStore::default())); - let result = tree_state.tree.get(key).unwrap(); + let result = tree.get(KeyHash::with::("non_existing_id")).unwrap(); let NotFound(non_membership_proof) = result else { panic!("Hashchain found for key while it was expected to be missing"); @@ -137,127 +193,140 @@ mod tests { #[test] fn test_multiple_inserts_and_updates() { - let mut tree_state = TestTreeState::default(); + let mut tree = KeyDirectoryTree::new(Arc::new(MockTreeStore::default())); + let mut tx_builder = TransactionBuilder::new(); - let service = tree_state.register_service("service_1".to_string()); - let mut account1 = tree_state.create_account("key_1".to_string(), service.clone()); - let mut account2 = tree_state.create_account("key_2".to_string(), service.clone()); + let service_tx = tx_builder.register_service_with_random_keys("service_1").commit(); + let acc1_tx = + tx_builder.create_account_with_random_key_signed("acc_1", "service_1").commit(); + let acc2_tx = + tx_builder.create_account_with_random_key_signed("acc_2", "service_1").commit(); - tree_state.insert_account(service.registration).unwrap(); + tree.process_transaction(service_tx).unwrap(); - tree_state.insert_account(account1.clone()).unwrap(); - tree_state.insert_account(account2.clone()).unwrap(); + tree.process_transaction(acc1_tx).unwrap(); + tree.process_transaction(acc2_tx).unwrap(); // Do insert and update accounts using the correct key indices - tree_state.add_key_to_account(&mut account1).unwrap(); - tree_state.update_account(account1.clone()).unwrap(); + let key_1_tx = tx_builder.add_random_key_verified_with_root("acc_1").commit(); + tree.process_transaction(key_1_tx).unwrap(); - tree_state.add_unsigned_data_to_account(b"unsigned", &mut account2).unwrap(); - tree_state.update_account(account2.clone()).unwrap(); - tree_state.add_signed_data_to_account(b"signed", &mut account2).unwrap(); - tree_state.update_account(account2.clone()).unwrap(); + let data_1_tx = + tx_builder.add_unsigned_data_verified_with_root("acc_2", b"unsigned".to_vec()).commit(); + tree.process_transaction(data_1_tx).unwrap(); - let get_result1 = tree_state.tree.get(account1.key_hash); - let get_result2 = tree_state.tree.get(account2.key_hash); + let data_2_tx = tx_builder + .add_randomly_signed_data_verified_with_root("acc_2", b"signed".to_vec()) + .commit(); + tree.process_transaction(data_2_tx).unwrap(); - assert!(matches!(get_result1.unwrap(), Found(hc, _) if hc == account1.hashchain)); - assert!(matches!(get_result2.unwrap(), Found(hc, _) if hc == account2.hashchain)); + let get_result1 = tree.get(KeyHash::with::("acc_1")).unwrap(); + let get_result2 = tree.get(KeyHash::with::("acc_2")).unwrap(); + + let test_hashchain_acc1 = tx_builder.get_hashchain("acc_1").unwrap(); + let test_hashchain_acc2 = tx_builder.get_hashchain("acc_2").unwrap(); + + assert!(matches!(get_result1, Found(hc, _) if &hc == test_hashchain_acc1)); + assert!(matches!(get_result2, Found(hc, _) if &hc == test_hashchain_acc2)); } #[test] fn test_interleaved_inserts_and_updates() { - let mut test_tree = TestTreeState::default(); + let mut tree = KeyDirectoryTree::new(Arc::new(MockTreeStore::default())); + let mut tx_builder = TransactionBuilder::new(); - let service = test_tree.register_service("service_1".to_string()); - let mut account_1 = test_tree.create_account("key_1".to_string(), service.clone()); - let mut account_2 = test_tree.create_account("key_2".to_string(), service.clone()); + let service_tx = tx_builder.register_service_with_random_keys("service_1").commit(); + let acc1_tx = + tx_builder.create_account_with_random_key_signed("acc_1", "service_1").commit(); + let acc2_tx = + tx_builder.create_account_with_random_key_signed("acc_2", "service_1").commit(); - test_tree.insert_account(service.registration).unwrap(); + tree.process_transaction(service_tx).unwrap(); + tree.process_transaction(acc1_tx).unwrap(); - test_tree.insert_account(account_1.clone()).unwrap(); + let add_key_to_1_tx = tx_builder.add_random_key_verified_with_root("acc_1").commit(); + tree.process_transaction(add_key_to_1_tx).unwrap(); - test_tree.add_key_to_account(&mut account_1).unwrap(); - // Update account_1 using the correct key index - test_tree.update_account(account_1.clone()).unwrap(); + tree.process_transaction(acc2_tx).unwrap(); - test_tree.insert_account(account_2.clone()).unwrap(); - - test_tree.add_key_to_account(&mut account_2).unwrap(); + let add_key_to_2_tx = tx_builder.add_random_key_verified_with_root("acc_2").commit(); + let last_proof = tree.process_transaction(add_key_to_2_tx).unwrap(); // Update account_2 using the correct key index - let last_proof = test_tree.update_account(account_2.clone()).unwrap(); + let Proof::Update(update_proof) = last_proof else { + panic!("Expetced insert proof for transaction"); + }; - let get_result1 = test_tree.tree.get(account_1.key_hash); - let get_result2 = test_tree.tree.get(account_2.key_hash); + let get_result1 = tree.get(KeyHash::with::("acc_1")).unwrap(); + let get_result2 = tree.get(KeyHash::with::("acc_2")).unwrap(); - assert!(matches!(get_result1.unwrap(), Found(hc, _) if hc == account_1.hashchain)); - assert!(matches!(get_result2.unwrap(), Found(hc, _) if hc == account_2.hashchain)); + let test_hashchain_acc1 = tx_builder.get_hashchain("acc_1").unwrap(); + let test_hashchain_acc2 = tx_builder.get_hashchain("acc_2").unwrap(); + + assert!(matches!(get_result1, Found(hc, _) if &hc == test_hashchain_acc1)); + assert!(matches!(get_result2, Found(hc, _) if &hc == test_hashchain_acc2)); assert_eq!( - last_proof.new_root, - test_tree.tree.get_current_root().unwrap() + Digest::from(update_proof.new_root), + tree.get_commitment().unwrap() ); } #[test] fn test_root_hash_changes() { - let mut tree_state = TestTreeState::default(); - let service = tree_state.register_service("service_1".to_string()); - let account = tree_state.create_account("key_1".to_string(), service.clone()); + let mut tree = KeyDirectoryTree::new(Arc::new(MockTreeStore::default())); + let mut tx_builder = TransactionBuilder::new(); + + let service_tx = tx_builder.register_service_with_random_keys("service_1").commit(); + let account1_tx = + tx_builder.create_account_with_random_key_signed("acc_1", "service_1").commit(); - tree_state.insert_account(service.registration).unwrap(); + tree.process_transaction(service_tx).unwrap(); - let root_before = tree_state.tree.get_current_root().unwrap(); - tree_state.insert_account(account).unwrap(); - let root_after = tree_state.tree.get_current_root().unwrap(); + let root_before = tree.get_current_root().unwrap(); + tree.process_transaction(account1_tx).unwrap(); + let root_after = tree.get_current_root().unwrap(); assert_ne!(root_before, root_after); } #[test] fn test_batch_writing() { - let mut tree_state = TestTreeState::default(); - let service = tree_state.register_service("service_1".to_string()); + let mut tree = KeyDirectoryTree::new(Arc::new(MockTreeStore::default())); + let mut tx_builder = TransactionBuilder::new(); - let account1 = tree_state.create_account("key_1".to_string(), service.clone()); - let account2 = tree_state.create_account("key_2".to_string(), service.clone()); - tree_state.insert_account(service.registration).unwrap(); + let service_tx = tx_builder.register_service_with_random_keys("service_1").commit(); + let account1_tx = + tx_builder.create_account_with_random_key_signed("acc_1", "service_1").commit(); + let account2_tx = + tx_builder.create_account_with_random_key_signed("acc_2", "service_1").commit(); - println!("Inserting key1: {:?}", account1.key_hash); - tree_state.insert_account(account1.clone()).unwrap(); + tree.process_transaction(service_tx).unwrap(); - println!( - "Tree state after first insert: {:?}", - tree_state.tree.get_commitment() - ); - println!( - "Tree state after first write_batch: {:?}", - tree_state.tree.get_commitment() - ); + println!("Inserting acc_1"); + tree.process_transaction(account1_tx).unwrap(); + + println!("Tree state after first insert: {:?}", tree.get_commitment()); // Try to get the first value immediately - let get_result1 = tree_state.tree.get(account1.key_hash); + let get_result1 = tree.get(KeyHash::with::("acc_1")); println!("Get result for key1 after first write: {:?}", get_result1); - println!("Inserting key2: {:?}", account2.key_hash); - tree_state.insert_account(account2.clone()).unwrap(); + println!("Inserting acc_2"); + tree.process_transaction(account2_tx).unwrap(); - println!( - "Tree state after second insert: {:?}", - tree_state.tree.get_commitment() - ); - println!( - "Tree state after second write_batch: {:?}", - tree_state.tree.get_commitment() - ); + println!("Tree state after 2nd insert: {:?}", tree.get_commitment()); // Try to get both values - let get_result1 = tree_state.tree.get(account1.key_hash); - let get_result2 = tree_state.tree.get(account2.key_hash); + let get_result1 = tree.get(KeyHash::with::("acc_1")).unwrap(); + let get_result2 = tree.get(KeyHash::with::("acc_2")).unwrap(); println!("Final get result for key1: {:?}", get_result1); println!("Final get result for key2: {:?}", get_result2); - assert!(matches!(get_result1.unwrap(), Found(hc, _) if hc == account1.hashchain)); - assert!(matches!(get_result2.unwrap(), Found(hc, _) if hc == account2.hashchain)); + let test_hashchain_acc1 = tx_builder.get_hashchain("acc_1").unwrap(); + let test_hashchain_acc2 = tx_builder.get_hashchain("acc_2").unwrap(); + + assert!(matches!(get_result1, Found(hc, _) if &hc == test_hashchain_acc1)); + assert!(matches!(get_result2, Found(hc, _) if &hc == test_hashchain_acc2)); } } diff --git a/crates/node_types/prover/src/prover/tests.rs b/crates/node_types/prover/src/prover/tests.rs index 61560dbe..a136d38b 100644 --- a/crates/node_types/prover/src/prover/tests.rs +++ b/crates/node_types/prover/src/prover/tests.rs @@ -22,10 +22,10 @@ fn create_mock_transactions(service_id: String) -> Vec { vec![ transaction_builder.register_service_with_random_keys(&service_id).commit(), transaction_builder - .create_account_with_random_key("user1@example.com", &service_id) + .create_account_with_random_key_signed("user1@example.com", &service_id) .commit(), transaction_builder - .create_account_with_random_key("user2@example.com", &service_id) + .create_account_with_random_key_signed("user2@example.com", &service_id) .commit(), transaction_builder.add_random_key_verified_with_root("user1@example.com").commit(), ] @@ -54,8 +54,9 @@ async fn test_process_transactions() { let mut transaction_builder = TransactionBuilder::new(); let register_service_transaction = transaction_builder.register_service_with_random_keys("test_service").commit(); - let create_account_transaction = - transaction_builder.create_account_with_random_key("test_account", "test_service").commit(); + let create_account_transaction = transaction_builder + .create_account_with_random_key_signed("test_account", "test_service") + .commit(); let proof = prover.process_transaction(register_service_transaction).await.unwrap(); assert!(matches!(proof, Proof::Insert(_))); @@ -96,7 +97,7 @@ async fn test_execute_block_with_invalid_tx() { let transactions = vec![ tx_builder.register_service_with_random_keys("service_id").commit(), - tx_builder.create_account_with_random_key("account_id", "service_id").commit(), + tx_builder.create_account_with_random_key_signed("account_id", "service_id").commit(), // add new key, so it will be index = 1 tx_builder.add_key_verified_with_root("account_id", new_key_vk.clone()).commit(), // revoke new key again diff --git a/crates/tests/src/lib.rs b/crates/tests/src/lib.rs index 88d58c1d..555e7bc0 100644 --- a/crates/tests/src/lib.rs +++ b/crates/tests/src/lib.rs @@ -17,8 +17,6 @@ use rand::{rngs::StdRng, Rng, SeedableRng}; use std::sync::Arc; use tokio::{spawn, time::Duration}; -use prism_common::test_utils::TestTreeState; - use tempfile::TempDir; fn setup_db() -> Arc> { @@ -50,8 +48,6 @@ async fn test_light_client_prover_talking() -> Result<()> { let signing_key = create_signing_key(); let pubkey = signing_key.verification_key(); - let mut test_state = TestTreeState::new(); - let _service = test_state.register_service("test_service".to_string()); let prover_cfg = prism_prover::Config { signing_key, ..prism_prover::Config::default() @@ -95,7 +91,7 @@ async fn test_light_client_prover_talking() -> Result<()> { for _ in 0..num_new_accounts { let random_user_id = format!("{}@gmail.com", i); let new_acc = transaction_builder - .create_account_with_random_key(random_user_id.as_str(), "test_service") + .create_account_with_random_key_signed(random_user_id.as_str(), "test_service") .commit(); match prover.clone().validate_and_queue_update(new_acc).await { Ok(_) => {