From c423bdea6140f3d7d20fde35f45e6b1d4b32e17c Mon Sep 17 00:00:00 2001 From: tyshko-rostyslav <122977916+tyshko-rostyslav@users.noreply.github.com> Date: Fri, 3 Mar 2023 06:03:39 +0100 Subject: [PATCH] chore(rln): update pmtree implementation (#125) * most changes * fmt * hide tests back under feature * grooming * changed `SledConfig` * requested change: rm `dbpath` --------- Co-authored-by: tyshkor --- rln/Cargo.toml | 4 +- rln/src/poseidon_tree.rs | 28 +++++ rln/tests/poseidon_tree.rs | 207 ++++++++++++++++++++++--------------- 3 files changed, 156 insertions(+), 83 deletions(-) diff --git a/rln/Cargo.toml b/rln/Cargo.toml index 7331a0d..40351bc 100644 --- a/rln/Cargo.toml +++ b/rln/Cargo.toml @@ -39,12 +39,12 @@ rand = "0.8" 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", optional = true} # serialization serde_json = "1.0.48" [dev-dependencies] -pmtree = { git = "https://github.com/Rate-Limiting-Nullifier/pmtree" } sled = "0.34.7" [features] @@ -54,4 +54,4 @@ wasm = ["wasmer/js", "wasmer/std"] fullmerkletree = ["default"] # Note: pmtree feature is still experimental -pmtree = ["default"] +pmtree-ft = ["default", "pmtree"] diff --git a/rln/src/poseidon_tree.rs b/rln/src/poseidon_tree.rs index a207a48..bbed00d 100644 --- a/rln/src/poseidon_tree.rs +++ b/rln/src/poseidon_tree.rs @@ -7,6 +7,11 @@ use crate::poseidon_hash::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::*; + // The zerokit RLN default Merkle tree implementation is the OptimalMerkleTree. // To switch to FullMerkleTree implementation, it is enough to enable the fullmerkletree feature @@ -36,3 +41,26 @@ impl utils::merkle_tree::Hasher for PoseidonHash { 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/tests/poseidon_tree.rs b/rln/tests/poseidon_tree.rs index b6211bb..5583185 100644 --- a/rln/tests/poseidon_tree.rs +++ b/rln/tests/poseidon_tree.rs @@ -82,7 +82,7 @@ mod test { // Test module for testing pmtree integration and features in zerokit // enabled only if the pmtree feature is enabled -#[cfg(feature = "pmtree")] +#[cfg(feature = "pmtree-ft")] #[cfg(test)] mod pmtree_test { @@ -91,47 +91,29 @@ mod pmtree_test { use rln::poseidon_hash::poseidon_hash; use rln::poseidon_tree::PoseidonHash; use rln::protocol::hash_to_field; - use rln::utils::{bytes_le_to_fr, fr_to_bytes_le, str_to_fr}; + use rln::utils::str_to_fr; use sled::Db as Sled; - use std::collections::HashMap; use std::fs; - use std::path::Path; + use std::{collections::HashMap, path::PathBuf}; use utils::{FullMerkleTree, OptimalMerkleTree}; - // 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) - } - } - // 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 { - fn new(_dbpath: &str) -> Result { + type Config = MemoryDBConfig; + + fn new(_config: Self::Config) -> Result { Ok(MemoryDB(HashMap::new())) } - fn load(_dbpath: &str) -> Result { - Err(Error("Cannot load in-memory DB".to_string())) + fn load(_config: Self::Config) -> Result { + Err(Box::new(Error("Cannot load in-memory DB".to_string()))) } fn get(&self, key: DBKey) -> Result> { @@ -143,36 +125,46 @@ mod pmtree_test { 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 { - fn new(dbpath: &str) -> Result { - if Path::new(dbpath).exists() { - match fs::remove_dir_all(dbpath) { + 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(Error(e.to_string())), + Err(e) => return Err(Box::new(Error(e.to_string()))), } } - let db: Sled = match sled::open(dbpath) { + let db: Sled = match config.open() { Ok(db) => db, - Err(e) => return Err(Error(e.to_string())), + Err(e) => return Err(Box::new(Error(e.to_string()))), }; Ok(SledDB(db)) } - fn load(dbpath: &str) -> Result { - let db: Sled = match sled::open(dbpath) { + fn load(config: Self::Config) -> Result { + let db: Sled = match sled::open(config.dbpath) { Ok(db) => db, - Err(e) => return Err(Error(e.to_string())), + Err(e) => return Err(Box::new(Error(e.to_string()))), }; if !db.was_recovered() { - return Err(Error("Trying to load non-existing database!".to_string())); + return Err(Box::new(Error( + "Trying to load non-existing database!".to_string(), + ))); } Ok(SledDB(db)) @@ -181,16 +173,27 @@ mod pmtree_test { fn get(&self, key: DBKey) -> Result> { match self.0.get(key) { Ok(value) => Ok(value.map(|val| val.to_vec())), - Err(e) => Err(Error(e.to_string())), + 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(Error(e.to_string())), + 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] @@ -341,6 +344,7 @@ mod pmtree_test { "0x21947ffd0bce0c385f876e7c97d6a42eec5b1fe935aab2f01c1f8a8cbcc356d2", 16 ) + .unwrap() ); let merkle_proof = tree.proof(leaf_index).expect("proof should exist"); @@ -353,83 +357,103 @@ mod pmtree_test { 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 = @@ -467,6 +491,7 @@ mod pmtree_test { "0x21947ffd0bce0c385f876e7c97d6a42eec5b1fe935aab2f01c1f8a8cbcc356d2", 16 ) + .unwrap() ); let merkle_proof = tree.proof(leaf_index).expect("proof should exist"); @@ -479,83 +504,103 @@ mod pmtree_test { 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 =