diff --git a/shielder/contract/Cargo.lock b/shielder/contract/Cargo.lock index efa48fe..b475592 100644 --- a/shielder/contract/Cargo.lock +++ b/shielder/contract/Cargo.lock @@ -3997,6 +3997,7 @@ dependencies = [ "anyhow", "drink", "ink", + "rand", ] [[package]] diff --git a/shielder/contract/Cargo.toml b/shielder/contract/Cargo.toml index 04cf5f3..bf17068 100644 --- a/shielder/contract/Cargo.toml +++ b/shielder/contract/Cargo.toml @@ -12,6 +12,7 @@ ink = { version = "5.0.0-rc", default-features = false } [dev-dependencies] drink = "0.9.0" +rand = { version = "0.8.5", default-features = false } [lib] path = "lib.rs" diff --git a/shielder/contract/drink_tests/mod.rs b/shielder/contract/drink_tests/mod.rs index cb9b4bf..2221067 100644 --- a/shielder/contract/drink_tests/mod.rs +++ b/shielder/contract/drink_tests/mod.rs @@ -3,16 +3,17 @@ mod utils; use anyhow::Result; use drink::{session::Session, AccountId32}; -use crate::{ - drink_tests::utils::{chain::*, ops::*, psp22::*, shielder::*}, - test_utils::merkle::MerkleTree, -}; +use crate::drink_tests::utils::{chain::*, ops::*, psp22::*, shielder::*}; + +use rand::prelude::*; #[drink::contract_bundle_provider] pub enum BundleProvider {} #[drink::test] fn deploy_single_deposit_single_withdraw(mut session: Session) -> Result<()> { + let mut rng = StdRng::seed_from_u64(1); + let alice = init_alice(&mut session)?; let bob = init_bob(&mut session)?; @@ -21,14 +22,12 @@ fn deploy_single_deposit_single_withdraw(mut session: Session) -> Result<()> { let shielder_address = deploy_shielder(&mut session)?; let psp22_address = deploy_test_token(&mut session, 100)?; - let mut merkle_tree = MerkleTree::new(); - // CREATE ACCOUNT let user_shielded_data = create_shielder_account( &mut session, &shielder_address, &psp22_address, - &mut merkle_tree, + rng.gen::().into(), )?; // APPROVE TRANSFER @@ -40,7 +39,7 @@ fn deploy_single_deposit_single_withdraw(mut session: Session) -> Result<()> { &shielder_address, deposit_op(&psp22_address, &alice, 10), user_shielded_data, - &mut merkle_tree, + rng.gen::().into(), )?; let alice_psp22_balance = get_psp22_balance(&mut session, &psp22_address, &alice)?; @@ -58,7 +57,7 @@ fn deploy_single_deposit_single_withdraw(mut session: Session) -> Result<()> { &shielder_address, withdraw_op(&psp22_address, &bob, 1), user_shielded_data, - &mut merkle_tree, + rng.gen::().into(), )?; let bob_psp22_balance = get_psp22_balance(&mut session, &psp22_address, &bob)?; @@ -72,6 +71,8 @@ fn deploy_single_deposit_single_withdraw(mut session: Session) -> Result<()> { #[drink::test] fn deploy_single_deposit_multiple_withdraw(mut session: Session) -> Result<()> { + let mut rng = StdRng::seed_from_u64(2); + let alice = init_alice(&mut session)?; session = session.with_actor(alice.clone()); @@ -85,14 +86,12 @@ fn deploy_single_deposit_multiple_withdraw(mut session: Session) -> Result<()> { let shielder_address = deploy_shielder(&mut session)?; let psp22_address = deploy_test_token(&mut session, 100)?; - let mut merkle_tree = MerkleTree::new(); - // CREATE ACCOUNT let mut user_shielded_data = create_shielder_account( &mut session, &shielder_address, &psp22_address, - &mut merkle_tree, + rng.gen::().into(), )?; // APPROVE TRANSFER @@ -108,7 +107,7 @@ fn deploy_single_deposit_multiple_withdraw(mut session: Session) -> Result<()> { &shielder_address, deposit_op(&psp22_address, &alice, 50), user_shielded_data, - &mut merkle_tree, + rng.gen::().into(), )?; for withdrawer_addr in withdrawers { @@ -121,7 +120,7 @@ fn deploy_single_deposit_multiple_withdraw(mut session: Session) -> Result<()> { &shielder_address, withdraw_op(&psp22_address, &withdrawer_addr, 1), user_shielded_data, - &mut merkle_tree, + rng.gen::().into(), )?; let psp22_balance = get_psp22_balance(&mut session, &psp22_address, &withdrawer_addr)?; assert_eq!(psp22_balance, 1); @@ -132,3 +131,78 @@ fn deploy_single_deposit_multiple_withdraw(mut session: Session) -> Result<()> { Ok(()) } + +#[drink::test] +fn deploy_multiple_deposit_multiple_withdraw(mut session: Session) -> Result<()> { + let mut rng = StdRng::seed_from_u64(3); + + let alice = init_alice(&mut session)?; + session = session.with_actor(alice.clone()); + + let mut depositors = vec![]; + for i in 3..11 { + let acc = AccountId32::new([i as u8; 32]); + init_acc_with_balance(&mut session, &acc)?; + depositors.push(acc.clone()); + } + + let mut withdrawers = vec![]; + for i in 11..19 { + let acc = AccountId32::new([i as u8; 32]); + init_acc_with_balance(&mut session, &acc)?; + withdrawers.push(acc.clone()); + } + + let shielder_address = deploy_shielder(&mut session)?; + let psp22_address = deploy_test_token(&mut session, 800)?; + for depositor_addr in &depositors { + psp22_transfer(&mut session, &psp22_address, &depositor_addr, 100)?; + } + + let mut user_shielded_data = vec![]; + for (i, depositor_addr) in depositors.iter().enumerate() { + // SWITCH TO withdrawer + session = session.with_actor(depositor_addr.clone()); + + // CREATE ACCOUNT + user_shielded_data.push(create_shielder_account( + &mut session, + &shielder_address, + &psp22_address, + rng.gen::().into(), + )?); + + // APPROVE TRANSFER + psp22_approve(&mut session, &psp22_address, &shielder_address, 50)?; + + // DEPOSIT + user_shielded_data[i] = shielder_update( + &mut session, + &shielder_address, + deposit_op(&psp22_address, &depositor_addr, 50), + user_shielded_data[i], + rng.gen::().into(), + )?; + } + + for (i, withdrawer_addr) in withdrawers.iter().enumerate() { + // SWITCH TO withdrawer + session = session.with_actor(withdrawer_addr.clone()); + + // WITHDRAW + user_shielded_data[i] = shielder_update( + &mut session, + &shielder_address, + withdraw_op(&psp22_address, &withdrawer_addr, 1), + user_shielded_data[i], + rng.gen::().into(), + )?; + let psp22_balance = get_psp22_balance(&mut session, &psp22_address, &withdrawer_addr)?; + assert_eq!(psp22_balance, 1); + } + let shielder_psp22_balance = + get_psp22_balance(&mut session, &psp22_address, &shielder_address)?; + assert_eq!(shielder_psp22_balance, 400 - 8); + + Ok(()) +} diff --git a/shielder/contract/drink_tests/utils/psp22.rs b/shielder/contract/drink_tests/utils/psp22.rs index 5187e08..fc66846 100644 --- a/shielder/contract/drink_tests/utils/psp22.rs +++ b/shielder/contract/drink_tests/utils/psp22.rs @@ -9,16 +9,13 @@ pub fn deploy_test_token( session: &mut Session, supply: u128, ) -> Result { - let formatted_supply = format!("{}", supply); - let psp22_bundle = ContractBundle::load(std::path::Path::new("../PSP22/target/ink/psp22.contract"))?; - let res = session.deploy_bundle( psp22_bundle, "new", &[ - formatted_supply.as_str(), + format!("{}", supply).as_str(), "Some(\"TST\")", "Some(\"TST\")", "9", @@ -64,11 +61,30 @@ pub fn psp22_approve( to: &AccountId32, amount: u128, ) -> Result<()> { - let formatted_amount = format!("{}", amount); session.call_with_address( token.clone(), "PSP22::approve", - &[to.to_string(), formatted_amount], + &[to.to_string(), format!("{}", amount)], + NO_ENDOWMENT, + )??; + Ok(()) +} + +pub fn psp22_transfer( + session: &mut Session, + token: &AccountId32, + to: &AccountId32, + amount: u128, +) -> Result<()> { + let empty_arr: [u8; 0] = []; + session.call_with_address( + token.clone(), + "PSP22::transfer", + &[ + to.to_string(), + format!("{}", amount), + format!("{:?}", empty_arr), + ], NO_ENDOWMENT, )??; Ok(()) diff --git a/shielder/contract/drink_tests/utils/shielder.rs b/shielder/contract/drink_tests/utils/shielder.rs index 495e2f6..680013a 100644 --- a/shielder/contract/drink_tests/utils/shielder.rs +++ b/shielder/contract/drink_tests/utils/shielder.rs @@ -6,7 +6,9 @@ use drink::{ }; use crate::{ + contract::DEPTH, drink_tests::{BundleProvider, UpdateOperation}, + errors::ShielderError, mocked_zk::{ account::Account, note::Note, @@ -15,10 +17,10 @@ use crate::{ traits::Hashable, TOKENS_NUMBER, }, - test_utils::merkle::MerkleTree, types::Scalar, }; +#[derive(Clone, Copy, Debug)] pub struct ShielderUserEnv { pub proof: ZkProof, pub nullifier: Scalar, @@ -40,7 +42,7 @@ pub fn create_shielder_account( session: &mut Session, shielder_address: &AccountId32, token: &AccountId32, - merkle_tree: &mut MerkleTree, + nullifier: Scalar, ) -> Result { let mut tokens: [Scalar; TOKENS_NUMBER] = [0_u128.into(); TOKENS_NUMBER]; tokens[0] = Scalar::from_bytes(*((*token).as_ref())); @@ -48,7 +50,6 @@ pub fn create_shielder_account( let acc = Account::new(tokens); let id = 0_128.into(); - let nullifier = 0_u128.into(); let trapdoor = 0_u128.into(); let op_priv = OpPriv { user: 0_u128.into(), @@ -58,19 +59,17 @@ pub fn create_shielder_account( let h_note_new = Note::new(id, trapdoor, nullifier, acc.hash()).hash(); - session.call_with_address( + let leaf_id_res: Result = session.call_with_address( shielder_address.clone(), "add_note", &[format!("{:?}", h_note_new), format!("{:?}", proof)], NO_ENDOWMENT, )??; - merkle_tree.add_leaf(h_note_new).unwrap(); - Ok(ShielderUserEnv { proof, nullifier, - tree_leaf_id: 0, + tree_leaf_id: leaf_id_res.unwrap(), }) } @@ -79,13 +78,22 @@ pub fn shielder_update( shielder_address: &AccountId32, upd_op: UpdateOperation, user_shielded_data: ShielderUserEnv, - merkle_tree: &mut MerkleTree, + nullifier: Scalar, ) -> Result { - let merkle_root = merkle_tree.root(); - let merkle_proof = merkle_tree - .gen_proof(user_shielded_data.tree_leaf_id as usize) - .unwrap(); - let nullifier_new = (u128::from(user_shielded_data.nullifier) + 1).into(); + let merkle_root_res: Result = session.call_with_address( + shielder_address.clone(), + "notes_merkle_root", + NO_ARGS, + NO_ENDOWMENT, + )??; + let merkle_root = merkle_root_res.unwrap(); + let merkle_proof_res: Result<[Scalar; DEPTH], ShielderError> = session.call_with_address( + shielder_address.clone(), + "notes_merkle_path", + &[format!("{:?}", user_shielded_data.tree_leaf_id)], + NO_ENDOWMENT, + )??; + let merkle_proof = merkle_proof_res.unwrap(); let trapdoor_new = 1_u128.into(); let op_pub = upd_op.op_pub; @@ -96,14 +104,13 @@ pub fn shielder_update( .update_account( operation, trapdoor_new, - nullifier_new, + nullifier, merkle_proof, user_shielded_data.tree_leaf_id, ) .unwrap(); - merkle_tree.add_leaf(note_hash).unwrap(); - session.call_with_address( + let new_leaf_id_res: Result = session.call_with_address( shielder_address.clone(), "update_note", &[ @@ -118,7 +125,7 @@ pub fn shielder_update( Ok(ShielderUserEnv { proof: new_proof, - nullifier: nullifier_new, - tree_leaf_id: user_shielded_data.tree_leaf_id + 1, + nullifier, + tree_leaf_id: new_leaf_id_res.unwrap(), }) } diff --git a/shielder/contract/errors.rs b/shielder/contract/errors.rs index 5fbd8aa..58b70a8 100644 --- a/shielder/contract/errors.rs +++ b/shielder/contract/errors.rs @@ -8,6 +8,7 @@ pub enum ShielderError { MerkleTreeVerificationFail, MerkleTreeLimitExceeded, MerkleTreeProofGenFail, + MerkleTreeNonExistingNode, ZkpVerificationFail, ArithmeticError, } diff --git a/shielder/contract/lib.rs b/shielder/contract/lib.rs index 4e069e1..d220c30 100755 --- a/shielder/contract/lib.rs +++ b/shielder/contract/lib.rs @@ -8,8 +8,6 @@ mod drink_tests; mod errors; mod merkle; pub mod mocked_zk; -#[cfg(test)] -pub mod test_utils; mod traits; mod types; @@ -49,12 +47,14 @@ pub mod contract { }, } + pub const DEPTH: usize = 10; + /// Contract storage #[ink(storage)] #[derive(Default)] pub struct Contract { nullifier_set: Set, - notes: MerkleTree, + notes: MerkleTree<{ DEPTH }>, } impl Contract { @@ -66,19 +66,20 @@ pub mod contract { /// Adds empty note to shielder storage /// Registers new account with empty balance + /// Returns id of the note in shielder's storage #[ink(message)] pub fn add_note( &mut self, h_note_new: Scalar, proof: ZkProof, - ) -> Result<(), ShielderError> { + ) -> Result { proof.verify_creation(h_note_new)?; - self.notes.add_leaf(h_note_new)?; - Ok(()) + self.notes.add_leaf(h_note_new) } /// Updates existing note /// Applies operation to private account stored in shielder + /// Returns id of the note in shielder's storage #[ink(message)] pub fn update_note( &mut self, @@ -87,13 +88,13 @@ pub mod contract { merkle_root: Scalar, nullifier_old: Scalar, proof: ZkProof, - ) -> Result<(), ShielderError> { + ) -> Result { self.notes.is_historical_root(merkle_root)?; self.nullify(nullifier_old)?; proof.verify_update(op_pub, h_note_new, merkle_root, nullifier_old)?; - self.notes.add_leaf(h_note_new)?; + let leaf_id = self.notes.add_leaf(h_note_new)?; self.process_operation(op_pub)?; - Ok(()) + Ok(leaf_id) } fn process_operation(&mut self, op_pub: OpPub) -> Result<(), ShielderError> { @@ -123,6 +124,19 @@ pub mod contract { Ok(()) } + /// Returns merkle root of notes storage + #[ink(message)] + pub fn notes_merkle_root(&self) -> Result { + self.notes.root() + } + + /// Returns merkle path + /// WARNING: that might expose identity of caller! + #[ink(message)] + pub fn notes_merkle_path(&self, note_id: u32) -> Result<[Scalar; DEPTH], ShielderError> { + self.notes.gen_proof(note_id) + } + fn nullify(&mut self, nullifier: Scalar) -> Result<(), ShielderError> { self.nullifier_set .insert(nullifier, &()) diff --git a/shielder/contract/merkle.rs b/shielder/contract/merkle.rs index 54ac933..c5069c3 100644 --- a/shielder/contract/merkle.rs +++ b/shielder/contract/merkle.rs @@ -9,11 +9,10 @@ use crate::{ }; /// depth of the tree -pub const DEPTH: usize = 10; #[ink::storage_item] #[derive(Debug)] -pub struct MerkleTree { +pub struct MerkleTree { /// mapping of tree indexes to values held in nodes nodes: Mapping, /// set of historical roots (nodes[1]) of tree @@ -30,7 +29,7 @@ pub fn compute_hash(first: Scalar, second: Scalar) -> Scalar { Scalar::from_bytes(res) } -impl Default for MerkleTree { +impl Default for MerkleTree { fn default() -> Self { Self { nodes: Default::default(), @@ -41,8 +40,14 @@ impl Default for MerkleTree { } } -impl MerkleTree { - pub fn add_leaf(&mut self, leaf_value: Scalar) -> Result<(), ShielderError> { +impl MerkleTree { + fn node_value(&self, id: u32) -> Result { + self.nodes + .get(id) + .ok_or(ShielderError::MerkleTreeNonExistingNode) + } + + pub fn add_leaf(&mut self, leaf_value: Scalar) -> Result { if self.next_leaf_idx == self.size { return Err(ShielderError::MerkleTreeLimitExceeded); } @@ -50,17 +55,20 @@ impl MerkleTree { .next_leaf_idx .checked_add(self.size) .ok_or(ShielderError::ArithmeticError)?; + let cur_leaf_id = self.next_leaf_idx; self.nodes.insert(id, &leaf_value); id /= 2; while id > 0 { let id_mul_2 = id.checked_mul(2).ok_or(ShielderError::ArithmeticError)?; - let left_n = self.node_value(id_mul_2); - let right_n = self.node_value( - id_mul_2 - .checked_add(1) - .ok_or(ShielderError::ArithmeticError)?, - ); + let left_n = self.node_value(id_mul_2).unwrap_or(0_u128.into()); + let right_n = self + .node_value( + id_mul_2 + .checked_add(1) + .ok_or(ShielderError::ArithmeticError)?, + ) + .unwrap_or(0_u128.into()); let hash = compute_hash(left_n, right_n); self.nodes.insert(id, &hash); id /= 2; @@ -69,8 +77,8 @@ impl MerkleTree { .next_leaf_idx .checked_add(1) .ok_or(ShielderError::ArithmeticError)?; - self.roots_log.insert(self.node_value(1), &()); - Ok(()) + self.roots_log.insert(self.root()?, &()); + Ok(cur_leaf_id) } pub fn is_historical_root(&self, merkle_root_possible: Scalar) -> Result<(), ShielderError> { @@ -80,7 +88,81 @@ impl MerkleTree { .ok_or(ShielderError::MerkleTreeVerificationFail) } - fn node_value(&self, id: u32) -> Scalar { - self.nodes.get(id).unwrap_or_default() + pub fn gen_proof(&self, leaf_id: u32) -> Result<[Scalar; DEPTH], ShielderError> { + let mut res = [Scalar::from_bytes([0x0; 32]); DEPTH]; + if self.next_leaf_idx == self.size { + return Err(ShielderError::MerkleTreeProofGenFail); + } + let mut id = leaf_id + .checked_add(self.size) + .ok_or(ShielderError::ArithmeticError)?; + for node in res.iter_mut().take(DEPTH) { + *node = self.node_value(id ^ 1).unwrap_or(0_u128.into()); + id /= 2; + } + Ok(res) + } + + pub fn root(&self) -> Result { + self.node_value(1) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use ink::primitives::AccountId; + + #[test] + fn add_two_leaves_and_root() { + ink::env::test::set_callee::(AccountId::from([0x1; 32])); + let mut merkle_tree = MerkleTree::<10>::default(); + let leaf0_id = merkle_tree.add_leaf(1_u128.into()).unwrap(); + assert_eq!(leaf0_id, 0); + let leaf1_id = merkle_tree.add_leaf(2_u128.into()).unwrap(); + assert_eq!(leaf1_id, 1); + + let mut hash_left = compute_hash(1_u128.into(), 2_u128.into()); + let mut hash_right = compute_hash(0_u128.into(), 0_u128.into()); + for _i in 1..10 { + hash_left = compute_hash(hash_left, 0_u128.into()); + hash_right = compute_hash(hash_right, hash_right); + } + + assert_eq!(hash_left, merkle_tree.root().unwrap()); + } + + #[test] + fn size_limit() { + ink::env::test::set_callee::(AccountId::from([0x1; 32])); + let mut merkle_tree = MerkleTree::<10>::default(); + for i in 0..(1 << 10) { + merkle_tree.add_leaf((i as u128).into()).unwrap(); + } + assert!(merkle_tree.add_leaf(0_u128.into()).is_err()); + } + + #[test] + fn historical_root() { + ink::env::test::set_callee::(AccountId::from([0x1; 32])); + let mut merkle_tree = MerkleTree::<10>::default(); + let mut roots = vec![]; + let leaves_num = 10; + for i in 0..leaves_num { + merkle_tree.add_leaf((i as u128).into()).unwrap(); + roots.push(merkle_tree.root().unwrap()); + } + // redeploy + ink::env::test::set_callee::(AccountId::from([0x2; 32])); + let mut merkle_tree = MerkleTree::<10>::default(); + for i in 0..leaves_num { + for j in 0..i { + assert!(merkle_tree.is_historical_root(roots[j]).is_ok()); + } + for j in i..leaves_num { + assert!(merkle_tree.is_historical_root(roots[j]).is_err()); + } + merkle_tree.add_leaf((i as u128).into()).unwrap(); + } } } diff --git a/shielder/contract/mocked_zk/relations.rs b/shielder/contract/mocked_zk/relations.rs index c5d620e..4e2e6d5 100644 --- a/shielder/contract/mocked_zk/relations.rs +++ b/shielder/contract/mocked_zk/relations.rs @@ -5,9 +5,9 @@ use super::{ traits::Hashable, }; use crate::{ - contract::OpPub, + contract::{OpPub, DEPTH}, errors::ShielderError, - merkle::{self, DEPTH}, + merkle::{self}, types::Scalar, }; @@ -24,7 +24,7 @@ pub struct ZkProof { acc_old: Account, acc_new: Account, op_priv: OpPriv, - merkle_proof: [Scalar; merkle::DEPTH], + merkle_proof: [Scalar; DEPTH], merkle_proof_leaf_id: u32, } diff --git a/shielder/contract/mocked_zk/tests.rs b/shielder/contract/mocked_zk/tests.rs index 0600171..199554e 100644 --- a/shielder/contract/mocked_zk/tests.rs +++ b/shielder/contract/mocked_zk/tests.rs @@ -8,10 +8,11 @@ use super::{ use crate::{ contract::OpPub, errors::ShielderError, + merkle::MerkleTree, mocked_zk::{mocked_user, MOCKED_TOKEN, TOKENS_NUMBER}, - test_utils::merkle::MerkleTree, types::Scalar, }; +use ink::primitives::AccountId; fn create_empty_note_proof(id: Scalar, nullifier: Scalar, trapdoor: Scalar) -> (Scalar, ZkProof) { let mut tokens: [Scalar; TOKENS_NUMBER] = [0_u128.into(); TOKENS_NUMBER]; @@ -80,16 +81,20 @@ fn test_create_note_fails() -> Result<(), ShielderError> { #[test] fn test_update_note() -> Result<(), ShielderError> { + // need this because MerkleTree is called + ink::env::test::set_callee::(AccountId::from([0x1; 32])); + let id = 0_u128.into(); - let mut merkle_tree = MerkleTree::new(); + let mut merkle_tree = MerkleTree::default(); let nullifier = 0_u128.into(); let trapdoor = 0_u128.into(); let (h_new_note, proof) = create_empty_note_proof(id, nullifier, trapdoor); proof.verify_creation(h_new_note)?; - let merkle_root = merkle_tree.add_leaf(h_new_note)?; + merkle_tree.add_leaf(h_new_note)?; + let merkle_root = merkle_tree.root()?; let merkle_proof = merkle_tree.gen_proof(0)?; let nullifier_new = 1_u128.into(); @@ -110,16 +115,20 @@ fn test_update_note() -> Result<(), ShielderError> { #[test] fn test_update_note_fail_op_priv() -> Result<(), ShielderError> { + // need this because merkle tree is called + ink::env::test::set_callee::(AccountId::from([0x1; 32])); + let id = 0_u128.into(); - let mut merkle_tree = MerkleTree::new(); + let mut merkle_tree = MerkleTree::default(); let nullifier = 0_u128.into(); let trapdoor = 0_u128.into(); let (h_new_note, proof) = create_empty_note_proof(id, nullifier, trapdoor); proof.verify_creation(h_new_note)?; - let merkle_root = merkle_tree.add_leaf(h_new_note)?; + merkle_tree.add_leaf(h_new_note)?; + let merkle_root = merkle_tree.root()?; let merkle_proof = merkle_tree.gen_proof(0)?; let nullifier_new = 1_u128.into(); diff --git a/shielder/contract/test_utils/merkle.rs b/shielder/contract/test_utils/merkle.rs deleted file mode 100644 index 7aa0951..0000000 --- a/shielder/contract/test_utils/merkle.rs +++ /dev/null @@ -1,65 +0,0 @@ -use ink::env::hash::{CryptoHash, Sha2x256}; - -use crate::{errors::ShielderError, types::Scalar}; - -/// depth of the tree -pub const DEPTH: usize = 10; - -#[derive(Default)] -pub struct MerkleTree { - nodes: Vec, - next_leaf_idx: usize, - size: usize, -} - -pub fn compute_hash(first: Scalar, second: Scalar) -> Scalar { - let mut res = [0x0; 32]; - Sha2x256::hash([first.bytes, second.bytes].concat().as_slice(), &mut res); - Scalar::from_bytes(res) -} - -impl MerkleTree { - pub fn new() -> Self { - Self { - nodes: vec![Scalar::from_bytes([0x0; 32]); 1 << (DEPTH + 1)], - next_leaf_idx: 0, - size: (1 << DEPTH), - } - } - - pub fn add_leaf(&mut self, leaf_value: Scalar) -> Result { - if self.next_leaf_idx == self.size { - return Err(ShielderError::MerkleTreeLimitExceeded); - } - let mut id = self.next_leaf_idx + self.size; - self.nodes[id] = leaf_value; - - id /= 2; - while id > 0 { - let left_n = self.nodes[id * 2]; - let right_n = self.nodes[id * 2 + 1]; - let hash = compute_hash(left_n, right_n); - self.nodes[id] = hash; - id /= 2; - } - self.next_leaf_idx += 1; - Ok(self.nodes[1]) - } - - pub fn root(&self) -> Scalar { - self.nodes[1] - } - - pub fn gen_proof(&self, leaf_id: usize) -> Result<[Scalar; DEPTH], ShielderError> { - let mut res = [Scalar::from_bytes([0x0; 32]); DEPTH]; - if self.next_leaf_idx == self.size { - return Err(ShielderError::MerkleTreeProofGenFail); - } - let mut id = leaf_id + self.size; - for i in 0..DEPTH { - res[i] = self.nodes[id ^ 1]; - id /= 2; - } - Ok(res) - } -} diff --git a/shielder/contract/test_utils/mod.rs b/shielder/contract/test_utils/mod.rs deleted file mode 100644 index bdf9eb8..0000000 --- a/shielder/contract/test_utils/mod.rs +++ /dev/null @@ -1 +0,0 @@ -pub mod merkle;