From a40e78d1222fd36823449fbe41559be7e185c9f4 Mon Sep 17 00:00:00 2001 From: Maksym Date: Tue, 20 Feb 2024 03:27:05 +0100 Subject: [PATCH 1/7] Update merkle tree tests --- shielder/contract/Cargo.lock | 1 + shielder/contract/Cargo.toml | 1 + shielder/contract/drink_tests/mod.rs | 102 +++++++++++++++--- shielder/contract/drink_tests/utils/psp22.rs | 17 +++ .../contract/drink_tests/utils/shielder.rs | 40 +++---- shielder/contract/lib.rs | 27 +++-- shielder/contract/merkle.rs | 41 ++++++- shielder/contract/mocked_zk/tests.rs | 19 +++- shielder/contract/test_utils/merkle.rs | 65 ----------- shielder/contract/test_utils/mod.rs | 1 - 10 files changed, 198 insertions(+), 116 deletions(-) delete mode 100644 shielder/contract/test_utils/merkle.rs delete mode 100644 shielder/contract/test_utils/mod.rs 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..457cc15 100644 --- a/shielder/contract/drink_tests/utils/psp22.rs +++ b/shielder/contract/drink_tests/utils/psp22.rs @@ -73,3 +73,20 @@ pub fn psp22_approve( )??; Ok(()) } + +pub fn psp22_transfer( + session: &mut Session, + token: &AccountId32, + to: &AccountId32, + amount: u128, +) -> Result<()> { + let formatted_amount = format!("{}", amount); + let empty_arr: [u8; 0] = []; + session.call_with_address( + token.clone(), + "PSP22::transfer", + &[to.to_string(), formatted_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..0d08d40 100644 --- a/shielder/contract/drink_tests/utils/shielder.rs +++ b/shielder/contract/drink_tests/utils/shielder.rs @@ -15,10 +15,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 +40,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 +48,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 +57,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 +76,21 @@ 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: Scalar = session.call_with_address( + shielder_address.clone(), + "notes_merkle_root", + NO_ARGS, + NO_ENDOWMENT, + )??; + let merkle_proof_res: Result<[Scalar; 10], ()> = 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 +101,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 +122,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/lib.rs b/shielder/contract/lib.rs index 4e069e1..2a32c87 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; @@ -19,7 +17,7 @@ pub mod contract { use crate::{ errors::ShielderError, - merkle::MerkleTree, + merkle::{MerkleTree, DEPTH}, mocked_zk::relations::ZkProof, traits::psp22::PSP22, types::{Scalar, Set}, @@ -71,10 +69,9 @@ pub mod contract { &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 @@ -87,13 +84,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 +120,18 @@ pub mod contract { Ok(()) } + /// Returns merkle root of notes storage + #[ink(message)] + pub fn notes_merkle_root(&self) -> Scalar { + self.notes.root() + } + + /// Returns merkle path + #[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..2ad2236 100644 --- a/shielder/contract/merkle.rs +++ b/shielder/contract/merkle.rs @@ -42,7 +42,11 @@ impl Default for MerkleTree { } impl MerkleTree { - pub fn add_leaf(&mut self, leaf_value: Scalar) -> Result<(), ShielderError> { + fn node_value(&self, id: u32) -> Scalar { + self.nodes.get(id).unwrap_or_default() + } + + pub fn add_leaf(&mut self, leaf_value: Scalar) -> Result { if self.next_leaf_idx == self.size { return Err(ShielderError::MerkleTreeLimitExceeded); } @@ -50,6 +54,7 @@ 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; @@ -70,7 +75,7 @@ impl MerkleTree { .checked_add(1) .ok_or(ShielderError::ArithmeticError)?; self.roots_log.insert(self.node_value(1), &()); - Ok(()) + Ok(cur_leaf_id) } pub fn is_historical_root(&self, merkle_root_possible: Scalar) -> Result<(), ShielderError> { @@ -80,7 +85,35 @@ 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); + id /= 2; + } + Ok(res) + } + + pub fn root(&self) -> Scalar { + self.node_value(1) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use ink::primitives::AccountId; + + #[test] + fn merkle_tree_test() { + ink::env::test::set_callee::(AccountId::from([0x1; 32])); + let mut merkle_tree = MerkleTree::default(); + merkle_tree.add_leaf(0_u128.into()).unwrap(); } } diff --git a/shielder/contract/mocked_zk/tests.rs b/shielder/contract/mocked_zk/tests.rs index 0600171..1eb50a2 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; From eab9ee7e781195b97d077fc1ef262f9e48e70dcb Mon Sep 17 00:00:00 2001 From: Maksym Date: Tue, 20 Feb 2024 03:58:36 +0100 Subject: [PATCH 2/7] upd unit for merkle --- shielder/contract/merkle.rs | 50 +++++++++++++++++++++++++++++++++++-- 1 file changed, 48 insertions(+), 2 deletions(-) diff --git a/shielder/contract/merkle.rs b/shielder/contract/merkle.rs index 2ad2236..cf918ae 100644 --- a/shielder/contract/merkle.rs +++ b/shielder/contract/merkle.rs @@ -111,9 +111,55 @@ mod tests { use ink::primitives::AccountId; #[test] - fn merkle_tree_test() { + fn add_two_leaves_and_root() { ink::env::test::set_callee::(AccountId::from([0x1; 32])); let mut merkle_tree = MerkleTree::default(); - merkle_tree.add_leaf(0_u128.into()).unwrap(); + 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..DEPTH { + hash_left = compute_hash(hash_left, 0_u128.into()); + hash_right = compute_hash(hash_right, hash_right); + } + + assert_eq!(hash_left, merkle_tree.root()); + } + + #[test] + fn size_limit() { + ink::env::test::set_callee::(AccountId::from([0x1; 32])); + let mut merkle_tree = MerkleTree::default(); + for i in 0..(1 << DEPTH) { + 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::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()); + } + // redeploy + ink::env::test::set_callee::(AccountId::from([0x2; 32])); + let mut merkle_tree = MerkleTree::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(); + } } } From 9b0cf2cc82fe0cc9d05d6e1fe6daa5667c243595 Mon Sep 17 00:00:00 2001 From: Maksym Date: Wed, 21 Feb 2024 10:58:05 +0100 Subject: [PATCH 3/7] fix --- shielder/contract/drink_tests/utils/psp22.rs | 14 +++++---- .../contract/drink_tests/utils/shielder.rs | 10 ++++--- shielder/contract/errors.rs | 1 + shielder/contract/lib.rs | 4 ++- shielder/contract/merkle.rs | 30 +++++++++++-------- shielder/contract/mocked_zk/tests.rs | 4 +-- 6 files changed, 37 insertions(+), 26 deletions(-) diff --git a/shielder/contract/drink_tests/utils/psp22.rs b/shielder/contract/drink_tests/utils/psp22.rs index 457cc15..c6a755f 100644 --- a/shielder/contract/drink_tests/utils/psp22.rs +++ b/shielder/contract/drink_tests/utils/psp22.rs @@ -9,7 +9,7 @@ pub fn deploy_test_token( session: &mut Session, supply: u128, ) -> Result { - let formatted_supply = format!("{}", supply); + // let formatted_supply = ; let psp22_bundle = ContractBundle::load(std::path::Path::new("../PSP22/target/ink/psp22.contract"))?; @@ -18,7 +18,7 @@ pub fn deploy_test_token( psp22_bundle, "new", &[ - formatted_supply.as_str(), + format!("{}", supply).as_str(), "Some(\"TST\")", "Some(\"TST\")", "9", @@ -64,11 +64,10 @@ 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(()) @@ -80,12 +79,15 @@ pub fn psp22_transfer( to: &AccountId32, amount: u128, ) -> Result<()> { - let formatted_amount = format!("{}", amount); let empty_arr: [u8; 0] = []; session.call_with_address( token.clone(), "PSP22::transfer", - &[to.to_string(), formatted_amount, format!("{:?}", empty_arr)], + &[ + 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 0d08d40..23519bb 100644 --- a/shielder/contract/drink_tests/utils/shielder.rs +++ b/shielder/contract/drink_tests/utils/shielder.rs @@ -7,6 +7,7 @@ use drink::{ use crate::{ drink_tests::{BundleProvider, UpdateOperation}, + errors::ShielderError, mocked_zk::{ account::Account, note::Note, @@ -57,7 +58,7 @@ pub fn create_shielder_account( let h_note_new = Note::new(id, trapdoor, nullifier, acc.hash()).hash(); - let leaf_id_res: Result = session.call_with_address( + let leaf_id_res: Result = session.call_with_address( shielder_address.clone(), "add_note", &[format!("{:?}", h_note_new), format!("{:?}", proof)], @@ -78,13 +79,14 @@ pub fn shielder_update( user_shielded_data: ShielderUserEnv, nullifier: Scalar, ) -> Result { - let merkle_root: Scalar = session.call_with_address( + let merkle_root_res: Result = session.call_with_address( shielder_address.clone(), "notes_merkle_root", NO_ARGS, NO_ENDOWMENT, )??; - let merkle_proof_res: Result<[Scalar; 10], ()> = session.call_with_address( + let merkle_root = merkle_root_res.unwrap(); + let merkle_proof_res: Result<[Scalar; 10], ShielderError> = session.call_with_address( shielder_address.clone(), "notes_merkle_path", &[format!("{:?}", user_shielded_data.tree_leaf_id)], @@ -107,7 +109,7 @@ pub fn shielder_update( ) .unwrap(); - let new_leaf_id_res: Result = session.call_with_address( + let new_leaf_id_res: Result = session.call_with_address( shielder_address.clone(), "update_note", &[ 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 2a32c87..7505e01 100755 --- a/shielder/contract/lib.rs +++ b/shielder/contract/lib.rs @@ -64,6 +64,7 @@ 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, @@ -76,6 +77,7 @@ pub mod contract { /// 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, @@ -122,7 +124,7 @@ pub mod contract { /// Returns merkle root of notes storage #[ink(message)] - pub fn notes_merkle_root(&self) -> Scalar { + pub fn notes_merkle_root(&self) -> Result { self.notes.root() } diff --git a/shielder/contract/merkle.rs b/shielder/contract/merkle.rs index cf918ae..828b15f 100644 --- a/shielder/contract/merkle.rs +++ b/shielder/contract/merkle.rs @@ -42,8 +42,10 @@ impl Default for MerkleTree { } impl MerkleTree { - fn node_value(&self, id: u32) -> Scalar { - self.nodes.get(id).unwrap_or_default() + 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 { @@ -60,12 +62,14 @@ impl MerkleTree { 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; @@ -74,7 +78,7 @@ impl MerkleTree { .next_leaf_idx .checked_add(1) .ok_or(ShielderError::ArithmeticError)?; - self.roots_log.insert(self.node_value(1), &()); + self.roots_log.insert(self.root()?, &()); Ok(cur_leaf_id) } @@ -94,13 +98,13 @@ impl MerkleTree { .checked_add(self.size) .ok_or(ShielderError::ArithmeticError)?; for node in res.iter_mut().take(DEPTH) { - *node = self.node_value(id ^ 1); + *node = self.node_value(id ^ 1).unwrap_or(0_u128.into()); id /= 2; } Ok(res) } - pub fn root(&self) -> Scalar { + pub fn root(&self) -> Result { self.node_value(1) } } @@ -126,7 +130,7 @@ mod tests { hash_right = compute_hash(hash_right, hash_right); } - assert_eq!(hash_left, merkle_tree.root()); + assert_eq!(hash_left, merkle_tree.root().unwrap()); } #[test] @@ -147,7 +151,7 @@ mod tests { let leaves_num = 10; for i in 0..leaves_num { merkle_tree.add_leaf((i as u128).into()).unwrap(); - roots.push(merkle_tree.root()); + roots.push(merkle_tree.root().unwrap()); } // redeploy ink::env::test::set_callee::(AccountId::from([0x2; 32])); diff --git a/shielder/contract/mocked_zk/tests.rs b/shielder/contract/mocked_zk/tests.rs index 1eb50a2..199554e 100644 --- a/shielder/contract/mocked_zk/tests.rs +++ b/shielder/contract/mocked_zk/tests.rs @@ -94,7 +94,7 @@ fn test_update_note() -> Result<(), ShielderError> { let (h_new_note, proof) = create_empty_note_proof(id, nullifier, trapdoor); proof.verify_creation(h_new_note)?; merkle_tree.add_leaf(h_new_note)?; - let merkle_root = merkle_tree.root(); + let merkle_root = merkle_tree.root()?; let merkle_proof = merkle_tree.gen_proof(0)?; let nullifier_new = 1_u128.into(); @@ -128,7 +128,7 @@ fn test_update_note_fail_op_priv() -> Result<(), ShielderError> { let (h_new_note, proof) = create_empty_note_proof(id, nullifier, trapdoor); proof.verify_creation(h_new_note)?; merkle_tree.add_leaf(h_new_note)?; - let merkle_root = merkle_tree.root(); + let merkle_root = merkle_tree.root()?; let merkle_proof = merkle_tree.gen_proof(0)?; let nullifier_new = 1_u128.into(); From b3aff02f7b26d23b4332f35881bc99aca80bf4d7 Mon Sep 17 00:00:00 2001 From: Maksym Date: Wed, 21 Feb 2024 10:59:54 +0100 Subject: [PATCH 4/7] upd --- shielder/contract/lib.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/shielder/contract/lib.rs b/shielder/contract/lib.rs index 7505e01..610b9f1 100755 --- a/shielder/contract/lib.rs +++ b/shielder/contract/lib.rs @@ -129,6 +129,7 @@ pub mod contract { } /// 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) From ecc5f13c66e815066a98137cde7e8db2c303aa6d Mon Sep 17 00:00:00 2001 From: Maksym Date: Wed, 21 Feb 2024 11:09:58 +0100 Subject: [PATCH 5/7] make depth generic --- .../contract/drink_tests/utils/shielder.rs | 3 ++- shielder/contract/lib.rs | 6 ++++-- shielder/contract/merkle.rs | 20 +++++++++---------- shielder/contract/mocked_zk/relations.rs | 6 +++--- 4 files changed, 19 insertions(+), 16 deletions(-) diff --git a/shielder/contract/drink_tests/utils/shielder.rs b/shielder/contract/drink_tests/utils/shielder.rs index 23519bb..680013a 100644 --- a/shielder/contract/drink_tests/utils/shielder.rs +++ b/shielder/contract/drink_tests/utils/shielder.rs @@ -6,6 +6,7 @@ use drink::{ }; use crate::{ + contract::DEPTH, drink_tests::{BundleProvider, UpdateOperation}, errors::ShielderError, mocked_zk::{ @@ -86,7 +87,7 @@ pub fn shielder_update( NO_ENDOWMENT, )??; let merkle_root = merkle_root_res.unwrap(); - let merkle_proof_res: Result<[Scalar; 10], ShielderError> = session.call_with_address( + 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)], diff --git a/shielder/contract/lib.rs b/shielder/contract/lib.rs index 610b9f1..d220c30 100755 --- a/shielder/contract/lib.rs +++ b/shielder/contract/lib.rs @@ -17,7 +17,7 @@ pub mod contract { use crate::{ errors::ShielderError, - merkle::{MerkleTree, DEPTH}, + merkle::MerkleTree, mocked_zk::relations::ZkProof, traits::psp22::PSP22, types::{Scalar, Set}, @@ -47,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 { diff --git a/shielder/contract/merkle.rs b/shielder/contract/merkle.rs index 828b15f..f451ef1 100644 --- a/shielder/contract/merkle.rs +++ b/shielder/contract/merkle.rs @@ -9,11 +9,11 @@ use crate::{ }; /// depth of the tree -pub const DEPTH: usize = 10; +// 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 +30,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,7 +41,7 @@ impl Default for MerkleTree { } } -impl MerkleTree { +impl MerkleTree { fn node_value(&self, id: u32) -> Result { self.nodes .get(id) @@ -117,7 +117,7 @@ mod tests { #[test] fn add_two_leaves_and_root() { ink::env::test::set_callee::(AccountId::from([0x1; 32])); - let mut merkle_tree = MerkleTree::default(); + 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(); @@ -125,7 +125,7 @@ mod tests { 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..DEPTH { + for _i in 1..10 { hash_left = compute_hash(hash_left, 0_u128.into()); hash_right = compute_hash(hash_right, hash_right); } @@ -136,8 +136,8 @@ mod tests { #[test] fn size_limit() { ink::env::test::set_callee::(AccountId::from([0x1; 32])); - let mut merkle_tree = MerkleTree::default(); - for i in 0..(1 << DEPTH) { + 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()); @@ -146,7 +146,7 @@ mod tests { #[test] fn historical_root() { ink::env::test::set_callee::(AccountId::from([0x1; 32])); - let mut merkle_tree = MerkleTree::default(); + let mut merkle_tree = MerkleTree::<10>::default(); let mut roots = vec![]; let leaves_num = 10; for i in 0..leaves_num { @@ -155,7 +155,7 @@ mod tests { } // redeploy ink::env::test::set_callee::(AccountId::from([0x2; 32])); - let mut merkle_tree = MerkleTree::default(); + 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()); 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, } From f32c39c6481db3f4acc1d855ed8c8bf13a005a62 Mon Sep 17 00:00:00 2001 From: Maksym Date: Thu, 22 Feb 2024 11:29:27 +0100 Subject: [PATCH 6/7] fix --- shielder/contract/drink_tests/utils/psp22.rs | 3 --- 1 file changed, 3 deletions(-) diff --git a/shielder/contract/drink_tests/utils/psp22.rs b/shielder/contract/drink_tests/utils/psp22.rs index c6a755f..fc66846 100644 --- a/shielder/contract/drink_tests/utils/psp22.rs +++ b/shielder/contract/drink_tests/utils/psp22.rs @@ -9,11 +9,8 @@ pub fn deploy_test_token( session: &mut Session, supply: u128, ) -> Result { - // let formatted_supply = ; - let psp22_bundle = ContractBundle::load(std::path::Path::new("../PSP22/target/ink/psp22.contract"))?; - let res = session.deploy_bundle( psp22_bundle, "new", From 0a5a15a145d85b074c28c03865028f3dee9f704e Mon Sep 17 00:00:00 2001 From: Maksym Date: Thu, 22 Feb 2024 11:30:13 +0100 Subject: [PATCH 7/7] fix --- shielder/contract/merkle.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/shielder/contract/merkle.rs b/shielder/contract/merkle.rs index f451ef1..c5069c3 100644 --- a/shielder/contract/merkle.rs +++ b/shielder/contract/merkle.rs @@ -9,7 +9,6 @@ use crate::{ }; /// depth of the tree -// pub const DEPTH: usize = 10; #[ink::storage_item] #[derive(Debug)]