From 780a4df7a6edcbeaf3bbd2650d01c2b2f51cc63a Mon Sep 17 00:00:00 2001 From: sifnoc Date: Wed, 14 Aug 2024 16:04:00 +0000 Subject: [PATCH] Added assertion for node and root balance while generating MST --- .github/workflows/rust.yml | 1 + csv/entry_16_max_node_balance.csv | 17 +++++++ zk_prover/Cargo.toml | 2 +- zk_prover/README.md | 6 +++ zk_prover/examples/gen_inclusion_verifier.rs | 16 +++--- zk_prover/src/circuits/merkle_sum_tree.rs | 2 +- zk_prover/src/circuits/tests.rs | 1 + zk_prover/src/merkle_sum_tree/entry.rs | 2 +- zk_prover/src/merkle_sum_tree/mst.rs | 4 +- zk_prover/src/merkle_sum_tree/tests.rs | 24 +++++++-- .../src/merkle_sum_tree/utils/build_tree.rs | 50 ++++++++++++++++--- zk_prover/src/merkle_sum_tree/utils/mod.rs | 4 +- 12 files changed, 104 insertions(+), 25 deletions(-) create mode 100644 csv/entry_16_max_node_balance.csv diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 4cf59ad1..53386247 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -82,6 +82,7 @@ jobs: run: | cd zk_prover cargo test --release --features dev-graph -- --nocapture + cargo test --package summa-solvency --lib --features "skip-node-balance-check" -- circuits::tests::test::test_balance_not_in_range test-zk-prover-examples: runs-on: ubuntu-latest diff --git a/csv/entry_16_max_node_balance.csv b/csv/entry_16_max_node_balance.csv new file mode 100644 index 00000000..471bb800 --- /dev/null +++ b/csv/entry_16_max_node_balance.csv @@ -0,0 +1,17 @@ +username,balance_ETH_ETH,balance_USDT_ETH +dxGaEAii,18446744073709551615,18446744073709551615 +MBlfbBGI,0,0 +lAhWlEWZ,18446744073709551615,18446744073709551615 +nuZweYtO,0,0 +gbdSwiuY,18446744073709551615,18446744073709551615 +RZNneNuP,0,0 +YsscHXkp,0,0 +RkLzkDun,0,0 +HlQlnEYI,18446744073709551615,18446744073709551615 +RqkZOFYe,0,0 +NjCSRAfD,0,0 +pHniJMQY,0,0 +dOGIMzKR,18446744073709551615,18446744073709551615 +HfMDmNLp,0,0 +xPLKzCBl,0,0 +AtwIxZHo,0,2 diff --git a/zk_prover/Cargo.toml b/zk_prover/Cargo.toml index bba8596c..bfb4a764 100644 --- a/zk_prover/Cargo.toml +++ b/zk_prover/Cargo.toml @@ -7,7 +7,7 @@ edition = "2021" [features] dev-graph = ["halo2_proofs/dev-graph", "plotters"] - +skip-node-balance-check = [] [dependencies] halo2_proofs = { git = "https://github.com/summa-dev/halo2"} diff --git a/zk_prover/README.md b/zk_prover/README.md index 28e4f5d8..cf88a612 100644 --- a/zk_prover/README.md +++ b/zk_prover/README.md @@ -11,6 +11,12 @@ cargo build cargo test --release --features dev-graph ``` +For testing the overflow case that in range check, execute: + +``` +cargo test --package summa-solvency --lib --features "skip-node-balance-check" -- circuits::tests::test::test_balance_not_in_range +``` + ## Documentation The documentation for the circuits can be generated by running diff --git a/zk_prover/examples/gen_inclusion_verifier.rs b/zk_prover/examples/gen_inclusion_verifier.rs index ca1b1dc7..aafd218c 100644 --- a/zk_prover/examples/gen_inclusion_verifier.rs +++ b/zk_prover/examples/gen_inclusion_verifier.rs @@ -6,8 +6,13 @@ use num_traits::Num; use prelude::*; use halo2_solidity_verifier::{compile_solidity, BatchOpenScheme::Bdfg21, SolidityGenerator}; -use summa_solvency::circuits::utils::generate_setup_artifacts; -use summa_solvency::circuits::{merkle_sum_tree::MstInclusionCircuit, WithInstances}; +use summa_solvency::{ + circuits::{ + utils::generate_setup_artifacts, + {merkle_sum_tree::MstInclusionCircuit, WithInstances}, + }, + merkle_sum_tree::utils::calculate_max_root_balance, +}; const LEVELS: usize = 4; const N_CURRENCIES: usize = 2; @@ -52,13 +57,6 @@ fn save_solidity(name: impl AsRef, solidity: &str) { println!("Saved {path}"); } -// Calculate the maximum value that the Merkle Root can have, given N_BYTES and LEVELS -fn calculate_max_root_balance(n_bytes: usize, n_levels: usize) -> BigInt { - // The max value that can be stored in a leaf node or a sibling node, according to the constraint set in the circuit - let max_leaf_value = BigInt::from(2).pow(n_bytes as u32 * 8) - 1; - max_leaf_value * (n_levels + 1) -} - // Given a combination of `N_BYTES` and `LEVELS`, check if there is a risk of overflow in the Merkle Root fn is_there_risk_of_overflow(n_bytes: usize, n_levels: usize) -> bool { // Calculate the max root balance value diff --git a/zk_prover/src/circuits/merkle_sum_tree.rs b/zk_prover/src/circuits/merkle_sum_tree.rs index f935c307..149fcd9c 100644 --- a/zk_prover/src/circuits/merkle_sum_tree.rs +++ b/zk_prover/src/circuits/merkle_sum_tree.rs @@ -74,7 +74,7 @@ where { pub fn init_empty() -> Self { Self { - entry: Entry::zero_entry(), + entry: Entry::init_empty(), path_indices: vec![Fp::zero(); LEVELS], sibling_leaf_node_hash_preimage: [Fp::zero(); N_CURRENCIES + 1], sibling_middle_node_hash_preimages: vec![[Fp::zero(); N_CURRENCIES + 2]; LEVELS], diff --git a/zk_prover/src/circuits/tests.rs b/zk_prover/src/circuits/tests.rs index 29025d31..f76daa71 100644 --- a/zk_prover/src/circuits/tests.rs +++ b/zk_prover/src/circuits/tests.rs @@ -267,6 +267,7 @@ mod test { // Building a proof using as input a csv file with an entry that is not in range [0, 2^N_BYTES*8 - 1] should fail the range check constraint on the leaf balance #[test] + #[cfg(feature = "skip-node-balance-check")] fn test_balance_not_in_range() { let merkle_sum_tree = MerkleSumTree::::from_csv("../csv/entry_16_overflow.csv") diff --git a/zk_prover/src/merkle_sum_tree/entry.rs b/zk_prover/src/merkle_sum_tree/entry.rs index 161fb76c..071c104b 100644 --- a/zk_prover/src/merkle_sum_tree/entry.rs +++ b/zk_prover/src/merkle_sum_tree/entry.rs @@ -27,7 +27,7 @@ impl Entry { } /// Returns a zero entry where the username is 0 and the balances are all 0 - pub fn zero_entry() -> Self { + pub fn init_empty() -> Self { let empty_balances: [BigUint; N_CURRENCIES] = std::array::from_fn(|_| BigUint::from(0u32)); Entry { diff --git a/zk_prover/src/merkle_sum_tree/mst.rs b/zk_prover/src/merkle_sum_tree/mst.rs index 5dc9bb7e..9607d6f4 100644 --- a/zk_prover/src/merkle_sum_tree/mst.rs +++ b/zk_prover/src/merkle_sum_tree/mst.rs @@ -114,14 +114,14 @@ impl MerkleSumTree(&leaves, depth)?; Ok(MerkleSumTree { root, diff --git a/zk_prover/src/merkle_sum_tree/tests.rs b/zk_prover/src/merkle_sum_tree/tests.rs index a14a99e2..c0dbb962 100644 --- a/zk_prover/src/merkle_sum_tree/tests.rs +++ b/zk_prover/src/merkle_sum_tree/tests.rs @@ -1,7 +1,9 @@ #[cfg(test)] mod test { - - use crate::merkle_sum_tree::utils::big_uint_to_fp; + use crate::merkle_sum_tree::utils::{ + big_uint_to_fp, build_leaves_from_entries, build_merkle_tree_from_leaves, + parse_csv_to_entries, + }; use crate::merkle_sum_tree::{Entry, MerkleSumTree, Node, Tree}; use num_bigint::{BigUint, ToBigUint}; use rand::Rng as _; @@ -65,6 +67,20 @@ mod test { assert!(!merkle_tree.verify_proof(&proof_invalid_2)); } + #[test] + #[should_panic(expected = "Node balance is exceed limit")] + fn test_node_and_root_balance_limit() { + let (_, entries) = parse_csv_to_entries::<&str, N_CURRENCIES, N_BYTES>( + "../csv/entry_16_max_node_balance.csv", + ) + .unwrap(); + + let leaves = build_leaves_from_entries::(&entries); + let depth = (entries.len() as f64).log2().ceil() as usize; + + let _ = build_merkle_tree_from_leaves::(&leaves, depth); + } + #[test] fn test_update_mst_leaf() { let merkle_tree_1 = @@ -210,7 +226,7 @@ mod test { // The last 3 entries of the merkle tree should be zero entries for i in 13..16 { let entry = merkle_tree.entries()[i].clone(); - assert_eq!(entry, Entry::::zero_entry()); + assert_eq!(entry, Entry::::init_empty()); } // expect root hash to be different than 0 @@ -242,7 +258,7 @@ mod test { // The last 15 entries of the merkle tree should be zero entries for i in 17..32 { let entry = merkle_tree.entries()[i].clone(); - assert_eq!(entry, Entry::::zero_entry()); + assert_eq!(entry, Entry::::init_empty()); } // expect root hash to be different than 0 diff --git a/zk_prover/src/merkle_sum_tree/utils/build_tree.rs b/zk_prover/src/merkle_sum_tree/utils/build_tree.rs index 05d39008..10f185d7 100644 --- a/zk_prover/src/merkle_sum_tree/utils/build_tree.rs +++ b/zk_prover/src/merkle_sum_tree/utils/build_tree.rs @@ -1,8 +1,9 @@ -use crate::merkle_sum_tree::{Entry, Node}; +use crate::merkle_sum_tree::{utils::big_uint_to_fp, Entry, Node}; use halo2_proofs::halo2curves::bn256::Fr as Fp; +use num_bigint::BigInt; use rayon::prelude::*; -pub fn build_merkle_tree_from_leaves( +pub fn build_merkle_tree_from_leaves( leaves: &[Node], depth: usize, ) -> Result<(Node, Vec>>), Box> @@ -18,8 +19,28 @@ where tree.push(leaves.to_vec()); + // The allowed_max_root_balance should be safe if set to half of `calculate_max_root_balance`. + // To achieve this, we use a depth one level lower here. + let allowed_max_node_balance = big_uint_to_fp( + &calculate_max_root_balance(N_BYTES, depth - 1) + .to_biguint() + .unwrap(), + ); + let allowed_max_root_balance = big_uint_to_fp( + &calculate_max_root_balance(N_BYTES, depth) + .to_biguint() + .unwrap(), + ); + for level in 1..=depth { - build_middle_level(level, &mut tree) + // Determine the maximum node balance based on the current level + let max_node_balance = if level == depth { + allowed_max_root_balance + } else { + allowed_max_node_balance + }; + + build_middle_level::(level, &mut tree, max_node_balance) } let root = tree[depth][0].clone(); @@ -33,14 +54,14 @@ where [usize; N_CURRENCIES + 1]: Sized, { // Precompute the zero leaf (this will only be used if we encounter a zero entry) - let zero_leaf = Entry::::zero_entry().compute_leaf(); + let zero_leaf = Entry::::init_empty().compute_leaf(); let leaves = entries .par_iter() .map(|entry| { // If the entry is the zero entry then we return the precomputed zero leaf // Otherwise, we compute the leaf as usual - if entry == &Entry::::zero_entry() { + if entry == &Entry::::init_empty() { zero_leaf.clone() } else { entry.compute_leaf() @@ -51,9 +72,10 @@ where leaves } -fn build_middle_level( +fn build_middle_level( level: usize, tree: &mut Vec>>, + max_node_balance: Fp, ) where [usize; N_CURRENCIES + 2]: Sized, { @@ -66,6 +88,15 @@ fn build_middle_level( for (i, balance) in hash_preimage.iter_mut().enumerate().take(N_CURRENCIES) { *balance = tree[level - 1][index].balances[i] + tree[level - 1][index + 1].balances[i]; + + // This conditional is for the test case `test_balance_not_in_range` that performs exceed case while generating proof + if !cfg!(feature = "skip-node-balance-check") { + assert!( + balance.to_owned() <= max_node_balance, + "{}", + format!("Node balance is exceed limit: {:#?}", max_node_balance), + ); + } } hash_preimage[N_CURRENCIES] = tree[level - 1][index].hash; @@ -76,3 +107,10 @@ fn build_middle_level( tree.push(results); } + +// Calculate the maximum value that the Merkle Root can have, given N_BYTES and LEVELS +pub fn calculate_max_root_balance(n_bytes: usize, n_levels: usize) -> BigInt { + // The max value that can be stored in a leaf node or a sibling node, according to the constraint set in the circuit + let max_leaf_value = BigInt::from(2).pow(n_bytes as u32 * 8) - 1; + max_leaf_value * (n_levels + 1) +} diff --git a/zk_prover/src/merkle_sum_tree/utils/mod.rs b/zk_prover/src/merkle_sum_tree/utils/mod.rs index c4f0a33e..321c32a3 100644 --- a/zk_prover/src/merkle_sum_tree/utils/mod.rs +++ b/zk_prover/src/merkle_sum_tree/utils/mod.rs @@ -2,6 +2,8 @@ mod build_tree; mod csv_parser; mod operation_helpers; -pub use build_tree::{build_leaves_from_entries, build_merkle_tree_from_leaves}; +pub use build_tree::{ + build_leaves_from_entries, build_merkle_tree_from_leaves, calculate_max_root_balance, +}; pub use csv_parser::parse_csv_to_entries; pub use operation_helpers::*;