diff --git a/Cargo.lock b/Cargo.lock index e036c48..5622c26 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2110,7 +2110,7 @@ dependencies = [ [[package]] name = "pmtree" version = "1.0.0" -source = "git+https://github.com/Rate-Limiting-Nullifier/pmtree?rev=f6d1a1fecad72cd39e6808e78085091d541dc882#f6d1a1fecad72cd39e6808e78085091d541dc882" +source = "git+https://github.com/Rate-Limiting-Nullifier/pmtree?rev=b3a02216cece3e9c24e1754ea381bf784fd1df48#b3a02216cece3e9c24e1754ea381bf784fd1df48" dependencies = [ "rayon", ] @@ -2393,7 +2393,6 @@ dependencies = [ "num-bigint", "num-traits", "once_cell", - "pmtree", "rand", "rand_chacha", "serde_json", @@ -3117,6 +3116,8 @@ dependencies = [ "hex-literal", "num-bigint", "num-traits", + "pmtree", + "sled", "tiny-keccak", ] diff --git a/rln/Cargo.toml b/rln/Cargo.toml index c89daa9..2d92d77 100644 --- a/rln/Cargo.toml +++ b/rln/Cargo.toml @@ -38,7 +38,6 @@ rand = "=0.8.5" rand_chacha = "=0.3.1" tiny-keccak = { version = "=2.0.2", features = ["keccak"] } utils = { path = "../utils/", default-features = false } -pmtree = { git = "https://github.com/Rate-Limiting-Nullifier/pmtree", rev = "f6d1a1fecad72cd39e6808e78085091d541dc882", optional = true} # serialization serde_json = "1.0.48" @@ -55,4 +54,4 @@ wasm = ["wasmer/js", "wasmer/std"] fullmerkletree = ["default"] # Note: pmtree feature is still experimental -pmtree-ft = ["default", "pmtree"] +pmtree-ft = ["default", "utils/pmtree-ft"] diff --git a/rln/src/hashers.rs b/rln/src/hashers.rs index f691f8c..31eeeab 100644 --- a/rln/src/hashers.rs +++ b/rln/src/hashers.rs @@ -28,6 +28,23 @@ pub fn poseidon_hash(input: &[Fr]) -> Fr { .expect("hash with fixed input size can't fail") } +// The zerokit RLN Merkle tree Hasher +#[derive(Clone, Copy, PartialEq, Eq)] +pub struct PoseidonHash; + +// The default Hasher trait used by Merkle tree implementation in utils +impl utils::merkle_tree::Hasher for PoseidonHash { + type Fr = Fr; + + fn default_leaf() -> Self::Fr { + Self::Fr::from(0) + } + + fn hash(inputs: &[Self::Fr]) -> Self::Fr { + poseidon_hash(inputs) + } +} + // Hashes arbitrary signal to the underlying prime field pub fn hash_to_field(signal: &[u8]) -> Fr { // We hash the input signal using Keccak256 diff --git a/rln/src/lib.rs b/rln/src/lib.rs index 5fc9c5d..f524d61 100644 --- a/rln/src/lib.rs +++ b/rln/src/lib.rs @@ -2,6 +2,8 @@ pub mod circuit; pub mod hashers; +#[cfg(feature = "pmtree-ft")] +pub mod pm_tree_adapter; pub mod poseidon_tree; pub mod protocol; pub mod public; diff --git a/rln/src/pm_tree_adapter.rs b/rln/src/pm_tree_adapter.rs new file mode 100644 index 0000000..9aaab57 --- /dev/null +++ b/rln/src/pm_tree_adapter.rs @@ -0,0 +1,159 @@ +use crate::circuit::Fr; +use crate::hashers::{poseidon_hash, PoseidonHash}; +use crate::utils::{bytes_le_to_fr, fr_to_bytes_le}; +use color_eyre::{Report, Result}; +use std::fmt::Debug; +use utils::*; + +pub struct PmTree { + tree: pmtree::MerkleTree, +} + +pub struct PmTreeProof { + proof: pmtree::tree::MerkleProof, +} + +pub type FrOf = ::Fr; + +// The pmtree Hasher trait used by pmtree Merkle tree +impl pmtree::Hasher for PoseidonHash { + type Fr = Fr; + + fn default_leaf() -> Self::Fr { + Fr::from(0) + } + + fn serialize(value: Self::Fr) -> pmtree::Value { + fr_to_bytes_le(&value) + } + + fn deserialize(value: pmtree::Value) -> Self::Fr { + let (fr, _) = bytes_le_to_fr(&value); + fr + } + + fn hash(inputs: &[Self::Fr]) -> Self::Fr { + poseidon_hash(inputs) + } +} + +pub struct PmtreeConfig(pm_tree::Config); +impl Default for PmtreeConfig { + fn default() -> Self { + let tmp_path = std::env::temp_dir().join(format!("pmtree-{}", rand::random::())); + PmtreeConfig(pm_tree::Config::new().temporary(true).path(tmp_path)) + } +} +impl Debug for PmtreeConfig { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + self.0.fmt(f) + } +} + +impl Clone for PmtreeConfig { + fn clone(&self) -> Self { + PmtreeConfig(self.0.clone()) + } +} + +impl ZerokitMerkleTree for PmTree { + type Proof = PmTreeProof; + type Hasher = PoseidonHash; + type Config = PmtreeConfig; + + fn default(depth: usize) -> Result { + let default_config = PmtreeConfig::default(); + PmTree::new(depth, Self::Hasher::default_leaf(), default_config) + } + + fn new(depth: usize, _default_leaf: FrOf, config: Self::Config) -> Result { + let tree_loaded = pmtree::MerkleTree::load(config.clone().0); + let tree = match tree_loaded { + Ok(tree) => tree, + Err(_) => pmtree::MerkleTree::new(depth, config.0)?, + }; + + Ok(PmTree { tree }) + } + + fn depth(&self) -> usize { + self.tree.depth() + } + + fn capacity(&self) -> usize { + self.tree.capacity() + } + + fn leaves_set(&mut self) -> usize { + self.tree.leaves_set() + } + + fn root(&self) -> FrOf { + self.tree.root() + } + + fn set(&mut self, index: usize, leaf: FrOf) -> Result<()> { + self.tree + .set(index, leaf) + .map_err(|e| Report::msg(e.to_string())) + } + + fn set_range>>( + &mut self, + start: usize, + values: I, + ) -> Result<()> { + self.tree + .set_range(start, values) + .map_err(|e| Report::msg(e.to_string())) + } + + fn update_next(&mut self, leaf: FrOf) -> Result<()> { + self.tree + .update_next(leaf) + .map_err(|e| Report::msg(e.to_string())) + } + + fn delete(&mut self, index: usize) -> Result<()> { + self.tree + .delete(index) + .map_err(|e| Report::msg(e.to_string())) + } + + fn proof(&self, index: usize) -> Result { + let proof = self.tree.proof(index)?; + Ok(PmTreeProof { proof }) + } + + fn verify(&self, leaf: &FrOf, witness: &Self::Proof) -> Result { + if self.tree.verify(leaf, &witness.proof) { + Ok(true) + } else { + Err(Report::msg("verify failed")) + } + } +} + +impl ZerokitMerkleProof for PmTreeProof { + type Index = u8; + type Hasher = PoseidonHash; + + fn length(&self) -> usize { + self.proof.length() + } + + fn leaf_index(&self) -> usize { + self.proof.leaf_index() + } + + fn get_path_elements(&self) -> Vec> { + self.proof.get_path_elements() + } + + fn get_path_index(&self) -> Vec { + self.proof.get_path_index() + } + fn compute_root_from(&self, leaf: &FrOf) -> FrOf { + self.proof.compute_root_from(leaf) + } +} diff --git a/rln/src/poseidon_tree.rs b/rln/src/poseidon_tree.rs index e85ccf9..e86e4a5 100644 --- a/rln/src/poseidon_tree.rs +++ b/rln/src/poseidon_tree.rs @@ -2,15 +2,16 @@ // Implementation inspired by https://github.com/worldcoin/semaphore-rs/blob/d462a4372f1fd9c27610f2acfe4841fab1d396aa/src/poseidon_tree.rs (no differences) -use crate::circuit::Fr; -use crate::hashers::poseidon_hash; use cfg_if::cfg_if; -use utils::merkle_tree::*; -#[cfg(feature = "pmtree-ft")] -use crate::utils::{bytes_le_to_fr, fr_to_bytes_le}; -#[cfg(feature = "pmtree-ft")] -use pmtree::*; +cfg_if! { + if #[cfg(feature = "pmtree-ft")] { + use crate::pm_tree_adapter::*; + } else { + use crate::hashers::{PoseidonHash}; + use utils::merkle_tree::*; + } +} // The zerokit RLN default Merkle tree implementation is the OptimalMerkleTree. // To switch to FullMerkleTree implementation, it is enough to enable the fullmerkletree feature @@ -19,48 +20,11 @@ cfg_if! { if #[cfg(feature = "fullmerkletree")] { pub type PoseidonTree = FullMerkleTree; pub type MerkleProof = FullMerkleProof; + } else if #[cfg(feature = "pmtree-ft")] { + pub type PoseidonTree = PmTree; + pub type MerkleProof = PmTreeProof; } else { pub type PoseidonTree = OptimalMerkleTree; pub type MerkleProof = OptimalMerkleProof; } } - -// The zerokit RLN Merkle tree Hasher -#[derive(Clone, Copy, PartialEq, Eq)] -pub struct PoseidonHash; - -// The default Hasher trait used by Merkle tree implementation in utils -impl utils::merkle_tree::Hasher for PoseidonHash { - type Fr = Fr; - - fn default_leaf() -> Self::Fr { - Self::Fr::from(0) - } - - fn hash(inputs: &[Self::Fr]) -> Self::Fr { - poseidon_hash(inputs) - } -} - -#[cfg(feature = "pmtree-ft")] -// The pmtree Hasher trait used by pmtree Merkle tree -impl pmtree::Hasher for PoseidonHash { - type Fr = Fr; - - fn default_leaf() -> Self::Fr { - Fr::from(0) - } - - fn serialize(value: Self::Fr) -> Value { - fr_to_bytes_le(&value) - } - - fn deserialize(value: Value) -> Self::Fr { - let (fr, _) = bytes_le_to_fr(&value); - fr - } - - fn hash(inputs: &[Self::Fr]) -> Self::Fr { - poseidon_hash(inputs) - } -} diff --git a/rln/src/public.rs b/rln/src/public.rs index 1eba48a..03f54b8 100644 --- a/rln/src/public.rs +++ b/rln/src/public.rs @@ -10,11 +10,10 @@ use ark_groth16::{ProvingKey, VerifyingKey}; use ark_relations::r1cs::ConstraintMatrices; use ark_serialize::{CanonicalDeserialize, CanonicalSerialize, Read, Write}; use cfg_if::cfg_if; -use color_eyre::Result; +use color_eyre::{Report, Result}; use num_bigint::BigInt; use std::io::Cursor; use utils::{ZerokitMerkleProof, ZerokitMerkleTree}; -// use rkyv::Deserialize; cfg_if! { if #[cfg(not(target_arch = "wasm32"))] { @@ -82,7 +81,7 @@ impl RLN<'_> { let verification_key = vk_from_folder(&resources_folder)?; // We compute a default empty tree - let tree = PoseidonTree::default(tree_height); + let tree = PoseidonTree::default(tree_height)?; Ok(RLN { witness_calculator, @@ -140,7 +139,7 @@ impl RLN<'_> { let verification_key = vk_from_raw(&vk_vec, &zkey_vec)?; // We compute a default empty tree - let tree = PoseidonTree::default(tree_height); + let tree = PoseidonTree::default(tree_height)?; Ok(RLN { #[cfg(not(target_arch = "wasm32"))] @@ -164,7 +163,7 @@ impl RLN<'_> { /// - `tree_height`: the height of the Merkle tree. pub fn set_tree(&mut self, tree_height: usize) -> Result<()> { // We compute a default empty tree of desired height - self.tree = PoseidonTree::default(tree_height); + self.tree = PoseidonTree::default(tree_height)?; Ok(()) } @@ -239,7 +238,10 @@ impl RLN<'_> { let (leaves, _) = bytes_le_to_vec_fr(&leaves_byte)?; // We set the leaves - self.tree.set_range(index, leaves) + self.tree + .set_range(index, leaves) + .map_err(|_| Report::msg("Could not set leaves"))?; + Ok(()) } /// Resets the tree state to default and sets multiple leaves starting from index 0. @@ -1061,7 +1063,6 @@ mod test { use ark_std::{rand::thread_rng, UniformRand}; use rand::Rng; use utils::ZerokitMerkleTree; - // use rkyv::Deserialize; #[test] // We test merkle batch Merkle tree additions @@ -1132,7 +1133,7 @@ mod test { // We now delete all leaves set and check if the root corresponds to the empty tree root // delete calls over indexes higher than no_of_leaves are ignored and will not increase self.tree.next_index - for i in 0..2 * no_of_leaves { + for i in 0..no_of_leaves { rln.delete_leaf(i).unwrap(); } diff --git a/rln/tests/ffi.rs b/rln/tests/ffi.rs index b9dbf9f..40acea6 100644 --- a/rln/tests/ffi.rs +++ b/rln/tests/ffi.rs @@ -96,8 +96,7 @@ mod test { // We now delete all leaves set and check if the root corresponds to the empty tree root // delete calls over indexes higher than no_of_leaves are ignored and will not increase self.tree.next_index - let delete_range = 2 * no_of_leaves; - for i in 0..delete_range { + for i in 0..no_of_leaves { let success = delete_leaf(rln_pointer, i); assert!(success, "delete leaf call failed"); } diff --git a/rln/tests/poseidon_tree.rs b/rln/tests/poseidon_tree.rs index 7072b39..f4b972f 100644 --- a/rln/tests/poseidon_tree.rs +++ b/rln/tests/poseidon_tree.rs @@ -5,7 +5,7 @@ #[cfg(test)] mod test { use rln::circuit::*; - use rln::poseidon_tree::*; + use rln::hashers::PoseidonHash; use utils::{FullMerkleTree, OptimalMerkleTree, ZerokitMerkleProof, ZerokitMerkleTree}; #[test] @@ -25,16 +25,16 @@ mod test { for _ in 0..sample_size.try_into().unwrap() { let now = Instant::now(); - FullMerkleTree::::default(tree_height); + FullMerkleTree::::default(tree_height).unwrap(); gen_time_full += now.elapsed().as_nanos(); let now = Instant::now(); - OptimalMerkleTree::::default(tree_height); + OptimalMerkleTree::::default(tree_height).unwrap(); gen_time_opt += now.elapsed().as_nanos(); } - let mut tree_full = FullMerkleTree::::default(tree_height); - let mut tree_opt = OptimalMerkleTree::::default(tree_height); + let mut tree_full = FullMerkleTree::::default(tree_height).unwrap(); + let mut tree_opt = OptimalMerkleTree::::default(tree_height).unwrap(); for i in 0..sample_size.try_into().unwrap() { let now = Instant::now(); @@ -78,540 +78,3 @@ mod test { ); } } - -// Test module for testing pmtree integration and features in zerokit -// enabled only if the pmtree feature is enabled - -#[cfg(feature = "pmtree-ft")] -#[cfg(test)] -mod pmtree_test { - - use pmtree::*; - use rln::circuit::Fr; - use rln::hashers::{hash_to_field, poseidon_hash}; - use rln::poseidon_tree::PoseidonHash; - use rln::protocol::hash_to_field; - use rln::utils::str_to_fr; - use sled::Db as Sled; - use std::fs; - use std::{collections::HashMap, path::PathBuf}; - use utils::{FullMerkleTree, OptimalMerkleTree}; - - // pmtree supports in-memory and on-disk databases (Database trait) for storing the Merkle tree state - - // We implement Database for hashmaps, an in-memory database - struct MemoryDB(HashMap); - - #[derive(Default)] - struct MemoryDBConfig {} - - impl Database for MemoryDB { - type Config = MemoryDBConfig; - - fn new(_config: Self::Config) -> Result { - Ok(MemoryDB(HashMap::new())) - } - - fn load(_config: Self::Config) -> Result { - Err(Box::new(Error("Cannot load in-memory DB".to_string()))) - } - - fn get(&self, key: DBKey) -> Result> { - Ok(self.0.get(&key).cloned()) - } - - fn put(&mut self, key: DBKey, value: Value) -> Result<()> { - self.0.insert(key, value); - - Ok(()) - } - - fn put_batch(&mut self, subtree: HashMap) -> Result<()> { - self.0.extend(subtree); - Ok(()) - } - } - - // We implement Database for sled DB, an on-disk database - struct SledDB(Sled); - - impl Database for SledDB { - type Config = sled::Config; - - fn new(config: Self::Config) -> Result { - let dbpath = config.path; - if config.dbpath.exists() { - match fs::remove_dir_all(&config.dbpath) { - Ok(x) => x, - Err(e) => return Err(Box::new(Error(e.to_string()))), - } - } - - let db: Sled = match config.open() { - Ok(db) => db, - Err(e) => return Err(Box::new(Error(e.to_string()))), - }; - - Ok(SledDB(db)) - } - - fn load(config: Self::Config) -> Result { - let db: Sled = match sled::open(config.dbpath) { - Ok(db) => db, - Err(e) => return Err(Box::new(Error(e.to_string()))), - }; - - if !db.was_recovered() { - return Err(Box::new(Error( - "Trying to load non-existing database!".to_string(), - ))); - } - - Ok(SledDB(db)) - } - - fn get(&self, key: DBKey) -> Result> { - match self.0.get(key) { - Ok(value) => Ok(value.map(|val| val.to_vec())), - Err(e) => Err(Box::new(Error(e.to_string()))), - } - } - - fn put(&mut self, key: DBKey, value: Value) -> Result<()> { - match self.0.insert(key, value) { - Ok(_) => Ok(()), - Err(e) => Err(Box::new(Error(e.to_string()))), - } - } - - fn put_batch(&mut self, subtree: HashMap) -> Result<()> { - let mut batch = sled::Batch::default(); - - for (key, value) in subtree { - batch.insert(&key, value); - } - - self.0.apply_batch(batch)?; - Ok(()) - } - } - - #[test] - /// A basic performance comparison between the two supported Merkle Tree implementations and in-memory/on-disk pmtree implementations - fn test_zerokit_and_pmtree_merkle_implementations_performances() { - use std::time::{Duration, Instant}; - - let tree_height = 20; - let sample_size = 100; - - let leaves: Vec = (0..sample_size).map(|s| Fr::from(s)).collect(); - - let mut gen_time_full: u128 = 0; - let mut upd_time_full: u128 = 0; - let mut gen_time_opt: u128 = 0; - let mut upd_time_opt: u128 = 0; - let mut gen_time_pm_memory: u128 = 0; - let mut upd_time_pm_memory: u128 = 0; - let mut gen_time_pm_sled: u128 = 0; - let mut upd_time_pm_sled: u128 = 0; - - for _ in 0..sample_size.try_into().unwrap() { - let now = Instant::now(); - FullMerkleTree::::default(tree_height); - gen_time_full += now.elapsed().as_nanos(); - - let now = Instant::now(); - OptimalMerkleTree::::default(tree_height); - gen_time_opt += now.elapsed().as_nanos(); - - let now = Instant::now(); - pmtree::MerkleTree::::default(tree_height).unwrap(); - gen_time_pm_memory += now.elapsed().as_nanos(); - - let now = Instant::now(); - pmtree::MerkleTree::::default(tree_height).unwrap(); - gen_time_pm_sled += now.elapsed().as_nanos(); - } - - let mut tree_full = FullMerkleTree::::default(tree_height); - let mut tree_opt = OptimalMerkleTree::::default(tree_height); - let mut tree_pm_memory = - pmtree::MerkleTree::::default(tree_height).unwrap(); - let mut tree_pm_sled = - pmtree::MerkleTree::::default(tree_height).unwrap(); - - for i in 0..sample_size.try_into().unwrap() { - let now = Instant::now(); - tree_full.set(i, leaves[i]).unwrap(); - upd_time_full += now.elapsed().as_nanos(); - let proof = tree_full.proof(i).expect("index should be set"); - assert_eq!(proof.leaf_index(), i); - - let now = Instant::now(); - tree_opt.set(i, leaves[i]).unwrap(); - upd_time_opt += now.elapsed().as_nanos(); - let proof = tree_opt.proof(i).expect("index should be set"); - assert_eq!(proof.leaf_index(), i); - - let now = Instant::now(); - tree_pm_memory.set(i, leaves[i]).unwrap(); - upd_time_pm_memory += now.elapsed().as_nanos(); - let proof = tree_pm_memory.proof(i).expect("index should be set"); - assert_eq!(proof.leaf_index(), i); - - let now = Instant::now(); - tree_pm_sled.set(i, leaves[i]).unwrap(); - upd_time_pm_sled += now.elapsed().as_nanos(); - let proof = tree_pm_sled.proof(i).expect("index should be set"); - assert_eq!(proof.leaf_index(), i); - } - - // We check all roots are the same - let tree_full_root = tree_full.root(); - let tree_opt_root = tree_opt.root(); - let tree_pm_memory_root = tree_pm_memory.root(); - let tree_pm_sled_root = tree_pm_sled.root(); - - assert_eq!(tree_full_root, tree_opt_root); - assert_eq!(tree_opt_root, tree_pm_memory_root); - assert_eq!(tree_pm_memory_root, tree_pm_sled_root); - - println!(" Average tree generation time:"); - println!( - " - Full Merkle Tree: {:?}", - Duration::from_nanos((gen_time_full / sample_size).try_into().unwrap()) - ); - println!( - " - Optimal Merkle Tree: {:?}", - Duration::from_nanos((gen_time_opt / sample_size).try_into().unwrap()) - ); - - println!( - " - Pmtree-HashMap Merkle Tree: {:?}", - Duration::from_nanos((gen_time_pm_memory / sample_size).try_into().unwrap()) - ); - - println!( - " - Pmtree-Sled Merkle Tree: {:?}", - Duration::from_nanos((gen_time_pm_sled / sample_size).try_into().unwrap()) - ); - - println!(" Average update_next execution time:"); - println!( - " - Full Merkle Tree: {:?}", - Duration::from_nanos((upd_time_full / sample_size).try_into().unwrap()) - ); - - println!( - " - Optimal Merkle Tree: {:?}", - Duration::from_nanos((upd_time_opt / sample_size).try_into().unwrap()) - ); - - println!( - " - Pmtree-HashMap Merkle Tree: {:?}", - Duration::from_nanos((upd_time_pm_memory / sample_size).try_into().unwrap()) - ); - - println!( - " - Pmtree-Sled Merkle Tree: {:?}", - Duration::from_nanos((upd_time_pm_sled / sample_size).try_into().unwrap()) - ); - } - - // The following two tests contain values that come from public::test_merkle_proof test - // We check that pmtree and zerokit Merkle tree implementations match. - - #[test] - fn test_pmtree_hashmap() -> Result<()> { - let tree_height = 20; - - let mut tree = pmtree::MerkleTree::::default(tree_height).unwrap(); - - let leaf_index = 3; - - let identity_secret = hash_to_field(b"test-merkle-proof"); - let id_commitment = poseidon_hash(&[identity_secret]); - - // let default_leaf = Fr::from(0); - tree.set(leaf_index, id_commitment).unwrap(); - - // We check correct computation of the root - let root = tree.root(); - - assert_eq!( - root, - str_to_fr( - "0x21947ffd0bce0c385f876e7c97d6a42eec5b1fe935aab2f01c1f8a8cbcc356d2", - 16 - ) - .unwrap() - ); - - let merkle_proof = tree.proof(leaf_index).expect("proof should exist"); - let path_elements = merkle_proof.get_path_elements(); - let identity_path_index = merkle_proof.get_path_index(); - - // We check correct computation of the path and indexes - // These values refers to tree height = 20 - let expected_path_elements = vec![ - str_to_fr( - "0x0000000000000000000000000000000000000000000000000000000000000000", - 16, - ) - .unwrap(), - str_to_fr( - "0x2098f5fb9e239eab3ceac3f27b81e481dc3124d55ffed523a839ee8446b64864", - 16, - ) - .unwrap(), - str_to_fr( - "0x1069673dcdb12263df301a6ff584a7ec261a44cb9dc68df067a4774460b1f1e1", - 16, - ) - .unwrap(), - str_to_fr( - "0x18f43331537ee2af2e3d758d50f72106467c6eea50371dd528d57eb2b856d238", - 16, - ) - .unwrap(), - str_to_fr( - "0x07f9d837cb17b0d36320ffe93ba52345f1b728571a568265caac97559dbc952a", - 16, - ) - .unwrap(), - str_to_fr( - "0x2b94cf5e8746b3f5c9631f4c5df32907a699c58c94b2ad4d7b5cec1639183f55", - 16, - ) - .unwrap(), - str_to_fr( - "0x2dee93c5a666459646ea7d22cca9e1bcfed71e6951b953611d11dda32ea09d78", - 16, - ) - .unwrap(), - str_to_fr( - "0x078295e5a22b84e982cf601eb639597b8b0515a88cb5ac7fa8a4aabe3c87349d", - 16, - ) - .unwrap(), - str_to_fr( - "0x2fa5e5f18f6027a6501bec864564472a616b2e274a41211a444cbe3a99f3cc61", - 16, - ) - .unwrap(), - str_to_fr( - "0x0e884376d0d8fd21ecb780389e941f66e45e7acce3e228ab3e2156a614fcd747", - 16, - ) - .unwrap(), - str_to_fr( - "0x1b7201da72494f1e28717ad1a52eb469f95892f957713533de6175e5da190af2", - 16, - ) - .unwrap(), - str_to_fr( - "0x1f8d8822725e36385200c0b201249819a6e6e1e4650808b5bebc6bface7d7636", - 16, - ) - .unwrap(), - str_to_fr( - "0x2c5d82f66c914bafb9701589ba8cfcfb6162b0a12acf88a8d0879a0471b5f85a", - 16, - ) - .unwrap(), - str_to_fr( - "0x14c54148a0940bb820957f5adf3fa1134ef5c4aaa113f4646458f270e0bfbfd0", - 16, - ) - .unwrap(), - str_to_fr( - "0x190d33b12f986f961e10c0ee44d8b9af11be25588cad89d416118e4bf4ebe80c", - 16, - ) - .unwrap(), - str_to_fr( - "0x22f98aa9ce704152ac17354914ad73ed1167ae6596af510aa5b3649325e06c92", - 16, - ) - .unwrap(), - str_to_fr( - "0x2a7c7c9b6ce5880b9f6f228d72bf6a575a526f29c66ecceef8b753d38bba7323", - 16, - ) - .unwrap(), - str_to_fr( - "0x2e8186e558698ec1c67af9c14d463ffc470043c9c2988b954d75dd643f36b992", - 16, - ) - .unwrap(), - str_to_fr( - "0x0f57c5571e9a4eab49e2c8cf050dae948aef6ead647392273546249d1c1ff10f", - 16, - ) - .unwrap(), - str_to_fr( - "0x1830ee67b5fb554ad5f63d4388800e1cfe78e310697d46e43c9ce36134f72cca", - 16, - ) - .unwrap(), - ]; - - let expected_identity_path_index: Vec = - vec![1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]; - - assert_eq!(path_elements, expected_path_elements); - assert_eq!(identity_path_index, expected_identity_path_index); - - // We check correct verification of the proof - assert!(tree.verify(&id_commitment, &merkle_proof)); - - Ok(()) - } - - #[test] - fn test_pmtree_sled() -> Result<()> { - let tree_height = 20; - - let mut tree = pmtree::MerkleTree::::default(tree_height).unwrap(); - - let leaf_index = 3; - - let identity_secret = hash_to_field(b"test-merkle-proof"); - let id_commitment = poseidon_hash(&[identity_secret]); - - // let default_leaf = Fr::from(0); - tree.set(leaf_index, id_commitment).unwrap(); - - // We check correct computation of the root - let root = tree.root(); - - assert_eq!( - root, - str_to_fr( - "0x21947ffd0bce0c385f876e7c97d6a42eec5b1fe935aab2f01c1f8a8cbcc356d2", - 16 - ) - .unwrap() - ); - - let merkle_proof = tree.proof(leaf_index).expect("proof should exist"); - let path_elements = merkle_proof.get_path_elements(); - let identity_path_index = merkle_proof.get_path_index(); - - // We check correct computation of the path and indexes - // These values refers to tree height = 20 - let expected_path_elements = vec![ - str_to_fr( - "0x0000000000000000000000000000000000000000000000000000000000000000", - 16, - ) - .unwrap(), - str_to_fr( - "0x2098f5fb9e239eab3ceac3f27b81e481dc3124d55ffed523a839ee8446b64864", - 16, - ) - .unwrap(), - str_to_fr( - "0x1069673dcdb12263df301a6ff584a7ec261a44cb9dc68df067a4774460b1f1e1", - 16, - ) - .unwrap(), - str_to_fr( - "0x18f43331537ee2af2e3d758d50f72106467c6eea50371dd528d57eb2b856d238", - 16, - ) - .unwrap(), - str_to_fr( - "0x07f9d837cb17b0d36320ffe93ba52345f1b728571a568265caac97559dbc952a", - 16, - ) - .unwrap(), - str_to_fr( - "0x2b94cf5e8746b3f5c9631f4c5df32907a699c58c94b2ad4d7b5cec1639183f55", - 16, - ) - .unwrap(), - str_to_fr( - "0x2dee93c5a666459646ea7d22cca9e1bcfed71e6951b953611d11dda32ea09d78", - 16, - ) - .unwrap(), - str_to_fr( - "0x078295e5a22b84e982cf601eb639597b8b0515a88cb5ac7fa8a4aabe3c87349d", - 16, - ) - .unwrap(), - str_to_fr( - "0x2fa5e5f18f6027a6501bec864564472a616b2e274a41211a444cbe3a99f3cc61", - 16, - ) - .unwrap(), - str_to_fr( - "0x0e884376d0d8fd21ecb780389e941f66e45e7acce3e228ab3e2156a614fcd747", - 16, - ) - .unwrap(), - str_to_fr( - "0x1b7201da72494f1e28717ad1a52eb469f95892f957713533de6175e5da190af2", - 16, - ) - .unwrap(), - str_to_fr( - "0x1f8d8822725e36385200c0b201249819a6e6e1e4650808b5bebc6bface7d7636", - 16, - ) - .unwrap(), - str_to_fr( - "0x2c5d82f66c914bafb9701589ba8cfcfb6162b0a12acf88a8d0879a0471b5f85a", - 16, - ) - .unwrap(), - str_to_fr( - "0x14c54148a0940bb820957f5adf3fa1134ef5c4aaa113f4646458f270e0bfbfd0", - 16, - ) - .unwrap(), - str_to_fr( - "0x190d33b12f986f961e10c0ee44d8b9af11be25588cad89d416118e4bf4ebe80c", - 16, - ) - .unwrap(), - str_to_fr( - "0x22f98aa9ce704152ac17354914ad73ed1167ae6596af510aa5b3649325e06c92", - 16, - ) - .unwrap(), - str_to_fr( - "0x2a7c7c9b6ce5880b9f6f228d72bf6a575a526f29c66ecceef8b753d38bba7323", - 16, - ) - .unwrap(), - str_to_fr( - "0x2e8186e558698ec1c67af9c14d463ffc470043c9c2988b954d75dd643f36b992", - 16, - ) - .unwrap(), - str_to_fr( - "0x0f57c5571e9a4eab49e2c8cf050dae948aef6ead647392273546249d1c1ff10f", - 16, - ) - .unwrap(), - str_to_fr( - "0x1830ee67b5fb554ad5f63d4388800e1cfe78e310697d46e43c9ce36134f72cca", - 16, - ) - .unwrap(), - ]; - - let expected_identity_path_index: Vec = - vec![1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]; - - assert_eq!(path_elements, expected_path_elements); - assert_eq!(identity_path_index, expected_identity_path_index); - - // We check correct verification of the proof - assert!(tree.verify(&id_commitment, &merkle_proof)); - - Ok(()) - } -} diff --git a/rln/tests/protocol.rs b/rln/tests/protocol.rs index 8a6c4f7..c63c9cb 100644 --- a/rln/tests/protocol.rs +++ b/rln/tests/protocol.rs @@ -10,6 +10,8 @@ mod test { use rln::utils::str_to_fr; use utils::{ZerokitMerkleProof, ZerokitMerkleTree}; + type ConfigOf = ::Config; + // Input generated with https://github.com/oskarth/zk-kit/commit/b6a872f7160c7c14e10a0ea40acab99cbb23c9a8 const WITNESS_JSON_15: &str = r#" { @@ -172,7 +174,12 @@ mod test { // generate merkle tree let default_leaf = Fr::from(0); - let mut tree = PoseidonTree::new(tree_height, default_leaf); + let mut tree = PoseidonTree::new( + tree_height, + default_leaf, + ConfigOf::::default(), + ) + .unwrap(); tree.set(leaf_index, id_commitment.into()).unwrap(); // We check correct computation of the root @@ -382,7 +389,12 @@ mod test { //// generate merkle tree let default_leaf = Fr::from(0); - let mut tree = PoseidonTree::new(tree_height, default_leaf); + let mut tree = PoseidonTree::new( + tree_height, + default_leaf, + ConfigOf::::default(), + ) + .unwrap(); tree.set(leaf_index, id_commitment.into()).unwrap(); let merkle_proof = tree.proof(leaf_index).expect("proof should exist"); diff --git a/utils/Cargo.toml b/utils/Cargo.toml index 8560037..74ff44e 100644 --- a/utils/Cargo.toml +++ b/utils/Cargo.toml @@ -8,6 +8,8 @@ license = "MIT OR Apache-2.0" ark-ff = { version = "=0.4.1", default-features = false, features = ["asm"] } num-bigint = { version = "=0.4.3", default-features = false, features = ["rand"] } color-eyre = "=0.6.2" +pmtree = { git = "https://github.com/Rate-Limiting-Nullifier/pmtree", rev = "b3a02216cece3e9c24e1754ea381bf784fd1df48", optional = true} +sled = "=0.34.7" [dev-dependencies] ark-bn254 = "=0.4.0" @@ -18,3 +20,4 @@ tiny-keccak = { version = "2.0.2", features = ["keccak"] } [features] default = ["parallel"] parallel = ["ark-ff/parallel"] +pmtree-ft = ["pmtree"] diff --git a/utils/src/lib.rs b/utils/src/lib.rs index 94e8694..597f453 100644 --- a/utils/src/lib.rs +++ b/utils/src/lib.rs @@ -3,3 +3,8 @@ pub use self::poseidon::*; pub mod merkle_tree; pub use self::merkle_tree::*; + +#[cfg(feature = "pmtree-ft")] +pub mod pm_tree; +#[cfg(feature = "pmtree-ft")] +pub use self::pm_tree::*; diff --git a/utils/src/merkle_tree/full_merkle_tree.rs b/utils/src/merkle_tree/full_merkle_tree.rs index 45399c7..19fc295 100644 --- a/utils/src/merkle_tree/full_merkle_tree.rs +++ b/utils/src/merkle_tree/full_merkle_tree.rs @@ -1,4 +1,4 @@ -use crate::merkle_tree::{Hasher, ZerokitMerkleProof, ZerokitMerkleTree}; +use crate::merkle_tree::{FrOf, Hasher, ZerokitMerkleProof, ZerokitMerkleTree}; use color_eyre::{Report, Result}; use std::{ cmp::max, @@ -45,16 +45,21 @@ pub enum FullMerkleBranch { pub struct FullMerkleProof(pub Vec>); /// Implementations - -impl ZerokitMerkleTree for FullMerkleTree { +impl ZerokitMerkleTree for FullMerkleTree +where + H: Hasher, +{ type Proof = FullMerkleProof; - fn default(depth: usize) -> Self { - FullMerkleTree::::new(depth, H::default_leaf()) + type Hasher = H; + type Config = (); + + fn default(depth: usize) -> Result { + FullMerkleTree::::new(depth, Self::Hasher::default_leaf(), ()) } /// Creates a new `MerkleTree` /// depth - the height of the tree made only of hash nodes. 2^depth is the maximum number of leaves hash nodes - fn new(depth: usize, initial_leaf: H::Fr) -> Self { + fn new(depth: usize, initial_leaf: FrOf, _config: Self::Config) -> Result { // Compute cache node values, leaf to root let cached_nodes = successors(Some(initial_leaf), |prev| Some(H::hash(&[*prev, *prev]))) .take(depth + 1) @@ -72,12 +77,12 @@ impl ZerokitMerkleTree for FullMerkleTree { let next_index = 0; - Self { + Ok(Self { depth, cached_nodes, nodes, next_index, - } + }) } // Returns the depth of the tree @@ -97,12 +102,12 @@ impl ZerokitMerkleTree for FullMerkleTree { #[must_use] // Returns the root of the tree - fn root(&self) -> H::Fr { + fn root(&self) -> FrOf { self.nodes[0] } // Sets a leaf at the specified tree index - fn set(&mut self, leaf: usize, hash: H::Fr) -> Result<()> { + fn set(&mut self, leaf: usize, hash: FrOf) -> Result<()> { self.set_range(leaf, once(hash))?; self.next_index = max(self.next_index, leaf + 1); Ok(()) @@ -110,7 +115,11 @@ impl ZerokitMerkleTree for FullMerkleTree { // Sets tree nodes, starting from start index // Function proper of FullMerkleTree implementation - fn set_range>(&mut self, start: usize, hashes: I) -> Result<()> { + fn set_range>>( + &mut self, + start: usize, + hashes: I, + ) -> Result<()> { let index = self.capacity() + start - 1; let mut count = 0; // first count number of hashes, and check that they fit in the tree @@ -131,7 +140,7 @@ impl ZerokitMerkleTree for FullMerkleTree { } // Sets a leaf at the next available index - fn update_next(&mut self, leaf: H::Fr) -> Result<()> { + fn update_next(&mut self, leaf: FrOf) -> Result<()> { self.set(self.next_index, leaf)?; Ok(()) } @@ -165,7 +174,7 @@ impl ZerokitMerkleTree for FullMerkleTree { } // Verifies a Merkle proof with respect to the input leaf and the tree root - fn verify(&self, hash: &H::Fr, proof: &FullMerkleProof) -> Result { + fn verify(&self, hash: &FrOf, proof: &FullMerkleProof) -> Result { Ok(proof.compute_root_from(hash) == self.root()) } } @@ -212,8 +221,9 @@ where } } -impl ZerokitMerkleProof for FullMerkleProof { +impl ZerokitMerkleProof for FullMerkleProof { type Index = u8; + type Hasher = H; #[must_use] // Returns the length of a Merkle proof @@ -232,7 +242,7 @@ impl ZerokitMerkleProof for FullMerkleProof { #[must_use] /// Returns the path elements forming a Merkle proof - fn get_path_elements(&self) -> Vec { + fn get_path_elements(&self) -> Vec> { self.0 .iter() .map(|x| match x { @@ -255,7 +265,7 @@ impl ZerokitMerkleProof for FullMerkleProof { /// Computes the Merkle root corresponding by iteratively hashing a Merkle proof with a given input leaf #[must_use] - fn compute_root_from(&self, hash: &H::Fr) -> H::Fr { + fn compute_root_from(&self, hash: &FrOf) -> FrOf { self.0.iter().fold(*hash, |hash, branch| match branch { FullMerkleBranch::Left(sibling) => H::hash(&[hash, *sibling]), FullMerkleBranch::Right(sibling) => H::hash(&[*sibling, hash]), diff --git a/utils/src/merkle_tree/merkle_tree.rs b/utils/src/merkle_tree/merkle_tree.rs index 26b93a3..e831dda 100644 --- a/utils/src/merkle_tree/merkle_tree.rs +++ b/utils/src/merkle_tree/merkle_tree.rs @@ -28,39 +28,42 @@ pub trait Hasher { fn hash(input: &[Self::Fr]) -> Self::Fr; } -/// In the ZerokitMerkleTree trait we define the methods that are required to be implemented by a Merkle tree -/// Including, OptimalMerkleTree, FullMerkleTree, Pmtree -pub trait ZerokitMerkleTree -where - H: Hasher, -{ - type Proof; +pub type FrOf = ::Fr; - fn default(depth: usize) -> Self; - fn new(depth: usize, default_leaf: H::Fr) -> Self; +/// In the ZerokitMerkleTree trait we define the methods that are required to be implemented by a Merkle tree +/// Including, OptimalMerkleTree, FullMerkleTree +pub trait ZerokitMerkleTree { + type Proof: ZerokitMerkleProof; + type Hasher: Hasher; + type Config: Default; + + fn default(depth: usize) -> Result + where + Self: Sized; + fn new(depth: usize, default_leaf: FrOf, config: Self::Config) -> Result + where + Self: Sized; fn depth(&self) -> usize; fn capacity(&self) -> usize; fn leaves_set(&mut self) -> usize; - fn root(&self) -> H::Fr; - fn set(&mut self, index: usize, leaf: H::Fr) -> Result<()>; + fn root(&self) -> FrOf; + fn set(&mut self, index: usize, leaf: FrOf) -> Result<()>; fn set_range(&mut self, start: usize, leaves: I) -> Result<()> where - I: IntoIterator; - fn update_next(&mut self, leaf: H::Fr) -> Result<()>; + I: IntoIterator>; + fn update_next(&mut self, leaf: FrOf) -> Result<()>; fn delete(&mut self, index: usize) -> Result<()>; fn proof(&self, index: usize) -> Result; - fn verify(&self, leaf: &H::Fr, witness: &Self::Proof) -> Result; + fn verify(&self, leaf: &FrOf, witness: &Self::Proof) -> Result; } -pub trait ZerokitMerkleProof -where - H: Hasher, -{ +pub trait ZerokitMerkleProof { type Index; + type Hasher: Hasher; fn length(&self) -> usize; fn leaf_index(&self) -> usize; - fn get_path_elements(&self) -> Vec; + fn get_path_elements(&self) -> Vec>; fn get_path_index(&self) -> Vec; - fn compute_root_from(&self, leaf: &H::Fr) -> H::Fr; + fn compute_root_from(&self, leaf: &FrOf) -> FrOf; } diff --git a/utils/src/merkle_tree/optimal_merkle_tree.rs b/utils/src/merkle_tree/optimal_merkle_tree.rs index acca647..71ab1b7 100644 --- a/utils/src/merkle_tree/optimal_merkle_tree.rs +++ b/utils/src/merkle_tree/optimal_merkle_tree.rs @@ -37,31 +37,33 @@ pub struct OptimalMerkleProof(pub Vec<(H::Fr, u8)>); /// Implementations -impl ZerokitMerkleTree for OptimalMerkleTree +impl ZerokitMerkleTree for OptimalMerkleTree where H: Hasher, { type Proof = OptimalMerkleProof; + type Hasher = H; + type Config = (); - fn default(depth: usize) -> Self { - OptimalMerkleTree::::new(depth, H::default_leaf()) + fn default(depth: usize) -> Result { + OptimalMerkleTree::::new(depth, H::default_leaf(), ()) } /// Creates a new `MerkleTree` /// depth - the height of the tree made only of hash nodes. 2^depth is the maximum number of leaves hash nodes - fn new(depth: usize, default_leaf: H::Fr) -> Self { + fn new(depth: usize, default_leaf: H::Fr, _config: Self::Config) -> Result { let mut cached_nodes: Vec = Vec::with_capacity(depth + 1); cached_nodes.push(default_leaf); for i in 0..depth { cached_nodes.push(H::hash(&[cached_nodes[i]; 2])); } cached_nodes.reverse(); - OptimalMerkleTree { + Ok(OptimalMerkleTree { cached_nodes: cached_nodes.clone(), depth, nodes: HashMap::new(), next_index: 0, - } + }) } // Returns the depth of the tree @@ -205,11 +207,13 @@ where } } -impl ZerokitMerkleProof for OptimalMerkleProof +impl ZerokitMerkleProof for OptimalMerkleProof where H: Hasher, { type Index = u8; + type Hasher = H; + #[must_use] // Returns the length of a Merkle proof fn length(&self) -> usize { diff --git a/utils/src/pm_tree/mod.rs b/utils/src/pm_tree/mod.rs new file mode 100644 index 0000000..b7aa25e --- /dev/null +++ b/utils/src/pm_tree/mod.rs @@ -0,0 +1,4 @@ +pub mod sled_adapter; +pub use self::sled_adapter::*; +pub use pmtree; +pub use sled::*; diff --git a/utils/src/pm_tree/sled_adapter.rs b/utils/src/pm_tree/sled_adapter.rs new file mode 100644 index 0000000..ba155da --- /dev/null +++ b/utils/src/pm_tree/sled_adapter.rs @@ -0,0 +1,75 @@ +use pmtree::*; + +use sled::Db as Sled; +use std::collections::HashMap; + +pub struct SledDB(Sled); + +impl Database for SledDB { + type Config = sled::Config; + + fn new(config: Self::Config) -> PmtreeResult { + let db: Sled = match config.open() { + Ok(db) => db, + Err(e) => { + return Err(PmtreeErrorKind::DatabaseError( + DatabaseErrorKind::CustomError(format!( + "Cannot create database: {} {:#?}", + e, config + )), + )) + } + }; + + Ok(SledDB(db)) + } + + fn load(config: Self::Config) -> PmtreeResult { + let db: Sled = match sled::open(&config.path) { + Ok(db) => db, + Err(e) => { + return Err(PmtreeErrorKind::DatabaseError( + DatabaseErrorKind::CustomError(format!("Cannot load database: {}", e)), + )) + } + }; + + if !db.was_recovered() { + return Err(PmtreeErrorKind::DatabaseError( + DatabaseErrorKind::CustomError(format!( + "Database was not recovered: {}", + config.path.display() + )), + )); + } + + Ok(SledDB(db)) + } + + fn get(&self, key: DBKey) -> PmtreeResult> { + match self.0.get(key) { + Ok(value) => Ok(value.map(|val| val.to_vec())), + Err(_e) => Err(PmtreeErrorKind::TreeError(TreeErrorKind::InvalidKey)), + } + } + + fn put(&mut self, key: DBKey, value: Value) -> PmtreeResult<()> { + match self.0.insert(key, value) { + Ok(_) => Ok(()), + Err(_e) => Err(PmtreeErrorKind::TreeError(TreeErrorKind::InvalidKey)), + } + } + + fn put_batch(&mut self, subtree: HashMap) -> PmtreeResult<()> { + let mut batch = sled::Batch::default(); + + for (key, value) in subtree { + batch.insert(&key, value); + } + + self.0 + .apply_batch(batch) + .map_err(|_| PmtreeErrorKind::TreeError(TreeErrorKind::InvalidKey))?; + Ok(()) + } +} diff --git a/utils/tests/merkle_tree.rs b/utils/tests/merkle_tree.rs index e21f193..68ba819 100644 --- a/utils/tests/merkle_tree.rs +++ b/utils/tests/merkle_tree.rs @@ -4,7 +4,7 @@ mod test { use hex_literal::hex; use tiny_keccak::{Hasher as _, Keccak}; use utils::{FullMerkleTree, Hasher, OptimalMerkleTree, ZerokitMerkleProof, ZerokitMerkleTree}; - + #[derive(Clone, Copy, Eq, PartialEq)] struct Keccak256; impl Hasher for Keccak256 { @@ -44,14 +44,14 @@ mod test { hex!("a9bb8c3f1f12e9aa903a50c47f314b57610a3ab32f2d463293f58836def38d36"), ]; - let mut tree = FullMerkleTree::::new(2, [0; 32]); + let mut tree = FullMerkleTree::::new(2, [0; 32], ()).unwrap(); assert_eq!(tree.root(), default_tree_root); for i in 0..leaves.len() { tree.set(i, leaves[i]).unwrap(); assert_eq!(tree.root(), roots[i]); } - let mut tree = OptimalMerkleTree::::new(2, [0; 32]); + let mut tree = OptimalMerkleTree::::new(2, [0; 32], ()).unwrap(); assert_eq!(tree.root(), default_tree_root); for i in 0..leaves.len() { tree.set(i, leaves[i]).unwrap(); @@ -69,7 +69,7 @@ mod test { ]; // We thest the FullMerkleTree implementation - let mut tree = FullMerkleTree::::new(2, [0; 32]); + let mut tree = FullMerkleTree::::new(2, [0; 32], ()).unwrap(); for i in 0..leaves.len() { // We set the leaves tree.set(i, leaves[i]).unwrap(); @@ -93,7 +93,7 @@ mod test { } // We test the OptimalMerkleTree implementation - let mut tree = OptimalMerkleTree::::new(2, [0; 32]); + let mut tree = OptimalMerkleTree::::new(2, [0; 32], ()).unwrap(); for i in 0..leaves.len() { // We set the leaves tree.set(i, leaves[i]).unwrap();