vacp2p
/
zerokit
mirror of https://github.com/vacp2p/zerokit.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Integrate pmtree into rln (#147) 

* feat: integrate pmtree with zerokit

* fix: tests
This commit is contained in:
Aaryamann Challani 2023-04-28 10:02:21 +05:30 committed by GitHub
parent 4f98fd8028
commit fd7d7d9318
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
18 changed files with 374 additions and 653 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
Cargo.lock generated
View File

@ -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",
]

3
rln/Cargo.toml
View File

@ -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"]

17
rln/src/hashers.rs
View File

@ -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

2
rln/src/lib.rs
View File

@ -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;

159
rln/src/pm_tree_adapter.rs Normal file
View File

@ -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<SledDB, PoseidonHash>,
}
pub struct PmTreeProof {
proof: pmtree::tree::MerkleProof<PoseidonHash>,
}
pub type FrOf<H> = <H as Hasher>::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::<u64>()));
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<Self> {
let default_config = PmtreeConfig::default();
PmTree::new(depth, Self::Hasher::default_leaf(), default_config)
}
fn new(depth: usize, _default_leaf: FrOf<Self::Hasher>, config: Self::Config) -> Result<Self> {
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::Hasher> {
self.tree.root()
}
fn set(&mut self, index: usize, leaf: FrOf<Self::Hasher>) -> Result<()> {
self.tree
.set(index, leaf)
.map_err(|e| Report::msg(e.to_string()))
}
fn set_range<I: IntoIterator<Item = FrOf<Self::Hasher>>>(
&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<Self::Hasher>) -> 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<Self::Proof> {
let proof = self.tree.proof(index)?;
Ok(PmTreeProof { proof })
}
fn verify(&self, leaf: &FrOf<Self::Hasher>, witness: &Self::Proof) -> Result<bool> {
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<FrOf<Self::Hasher>> {
self.proof.get_path_elements()
}
fn get_path_index(&self) -> Vec<Self::Index> {
self.proof.get_path_index()
}
fn compute_root_from(&self, leaf: &FrOf<Self::Hasher>) -> FrOf<Self::Hasher> {
self.proof.compute_root_from(leaf)
}
}

58
rln/src/poseidon_tree.rs
View File

@ -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<PoseidonHash>;
pub type MerkleProof = FullMerkleProof<PoseidonHash>;
} else if #[cfg(feature = "pmtree-ft")] {
pub type PoseidonTree = PmTree;
pub type MerkleProof = PmTreeProof;
} else {
pub type PoseidonTree = OptimalMerkleTree<PoseidonHash>;
pub type MerkleProof = OptimalMerkleProof<PoseidonHash>;
}
}
// 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)
}
}

17
rln/src/public.rs
View File

@ -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();
}

3
rln/tests/ffi.rs
View File

@ -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");
}

547
rln/tests/poseidon_tree.rs
View File

@ -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::<PoseidonHash>::default(tree_height);
FullMerkleTree::<PoseidonHash>::default(tree_height).unwrap();
gen_time_full += now.elapsed().as_nanos();
let now = Instant::now();
OptimalMerkleTree::<PoseidonHash>::default(tree_height);
OptimalMerkleTree::<PoseidonHash>::default(tree_height).unwrap();
gen_time_opt += now.elapsed().as_nanos();
}
let mut tree_full = FullMerkleTree::<PoseidonHash>::default(tree_height);
let mut tree_opt = OptimalMerkleTree::<PoseidonHash>::default(tree_height);
let mut tree_full = FullMerkleTree::<PoseidonHash>::default(tree_height).unwrap();
let mut tree_opt = OptimalMerkleTree::<PoseidonHash>::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<DBKey, Value>);
#[derive(Default)]
struct MemoryDBConfig {}
impl Database for MemoryDB {
type Config = MemoryDBConfig;
fn new(_config: Self::Config) -> Result<Self> {
Ok(MemoryDB(HashMap::new()))
}
fn load(_config: Self::Config) -> Result<Self> {
Err(Box::new(Error("Cannot load in-memory DB".to_string())))
}
fn get(&self, key: DBKey) -> Result<Option<Value>> {
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<DBKey, Value>) -> 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<Self> {
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<Self> {
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<Option<Value>> {
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<DBKey, Value>) -> 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<Fr> = (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::<PoseidonHash>::default(tree_height);
gen_time_full += now.elapsed().as_nanos();
let now = Instant::now();
OptimalMerkleTree::<PoseidonHash>::default(tree_height);
gen_time_opt += now.elapsed().as_nanos();
let now = Instant::now();
pmtree::MerkleTree::<MemoryDB, PoseidonHash>::default(tree_height).unwrap();
gen_time_pm_memory += now.elapsed().as_nanos();
let now = Instant::now();
pmtree::MerkleTree::<SledDB, PoseidonHash>::default(tree_height).unwrap();
gen_time_pm_sled += now.elapsed().as_nanos();
}
let mut tree_full = FullMerkleTree::<PoseidonHash>::default(tree_height);
let mut tree_opt = OptimalMerkleTree::<PoseidonHash>::default(tree_height);
let mut tree_pm_memory =
pmtree::MerkleTree::<MemoryDB, PoseidonHash>::default(tree_height).unwrap();
let mut tree_pm_sled =
pmtree::MerkleTree::<SledDB, PoseidonHash>::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::<MemoryDB, PoseidonHash>::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<u8> =
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::<SledDB, PoseidonHash>::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<u8> =
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(())
}
}

16
rln/tests/protocol.rs
View File

@ -10,6 +10,8 @@ mod test {
use rln::utils::str_to_fr;
use utils::{ZerokitMerkleProof, ZerokitMerkleTree};
type ConfigOf<T> = <T as ZerokitMerkleTree>::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::<PoseidonTree>::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::<PoseidonTree>::default(),
)
.unwrap();
tree.set(leaf_index, id_commitment.into()).unwrap();
let merkle_proof = tree.proof(leaf_index).expect("proof should exist");

3
utils/Cargo.toml
View File

@ -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"]

5
utils/src/lib.rs
View File

@ -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::*;

42
utils/src/merkle_tree/full_merkle_tree.rs
View File

@ -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<H: Hasher> {
pub struct FullMerkleProof<H: Hasher>(pub Vec<FullMerkleBranch<H>>);
/// Implementations
impl<H: Hasher> ZerokitMerkleTree<H> for FullMerkleTree<H> {
impl<H: Hasher> ZerokitMerkleTree for FullMerkleTree<H>
where
H: Hasher,
{
type Proof = FullMerkleProof<H>;
fn default(depth: usize) -> Self {
FullMerkleTree::<H>::new(depth, H::default_leaf())
type Hasher = H;
type Config = ();
fn default(depth: usize) -> Result<Self> {
FullMerkleTree::<H>::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<Self::Hasher>, _config: Self::Config) -> Result<Self> {
// 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<H: Hasher> ZerokitMerkleTree<H> for FullMerkleTree<H> {
let next_index = 0;
Self {
Ok(Self {
depth,
cached_nodes,
nodes,
next_index,
}
})
}
// Returns the depth of the tree
@ -97,12 +102,12 @@ impl<H: Hasher> ZerokitMerkleTree<H> for FullMerkleTree<H> {
#[must_use]
// Returns the root of the tree
fn root(&self) -> H::Fr {
fn root(&self) -> FrOf<Self::Hasher> {
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<Self::Hasher>) -> Result<()> {
self.set_range(leaf, once(hash))?;
self.next_index = max(self.next_index, leaf + 1);
Ok(())
@ -110,7 +115,11 @@ impl<H: Hasher> ZerokitMerkleTree<H> for FullMerkleTree<H> {
// Sets tree nodes, starting from start index
// Function proper of FullMerkleTree implementation
fn set_range<I: IntoIterator<Item = H::Fr>>(&mut self, start: usize, hashes: I) -> Result<()> {
fn set_range<I: IntoIterator<Item = FrOf<Self::Hasher>>>(
&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<H: Hasher> ZerokitMerkleTree<H> for FullMerkleTree<H> {
}
// Sets a leaf at the next available index
fn update_next(&mut self, leaf: H::Fr) -> Result<()> {
fn update_next(&mut self, leaf: FrOf<Self::Hasher>) -> Result<()> {
self.set(self.next_index, leaf)?;
Ok(())
}
@ -165,7 +174,7 @@ impl<H: Hasher> ZerokitMerkleTree<H> for FullMerkleTree<H> {
}
// Verifies a Merkle proof with respect to the input leaf and the tree root
fn verify(&self, hash: &H::Fr, proof: &FullMerkleProof<H>) -> Result<bool> {
fn verify(&self, hash: &FrOf<Self::Hasher>, proof: &FullMerkleProof<H>) -> Result<bool> {
Ok(proof.compute_root_from(hash) == self.root())
}
}
@ -212,8 +221,9 @@ where
}
}
impl<H: Hasher> ZerokitMerkleProof<H> for FullMerkleProof<H> {
impl<H: Hasher> ZerokitMerkleProof for FullMerkleProof<H> {
type Index = u8;
type Hasher = H;
#[must_use]
// Returns the length of a Merkle proof
@ -232,7 +242,7 @@ impl<H: Hasher> ZerokitMerkleProof<H> for FullMerkleProof<H> {
#[must_use]
/// Returns the path elements forming a Merkle proof
fn get_path_elements(&self) -> Vec<H::Fr> {
fn get_path_elements(&self) -> Vec<FrOf<Self::Hasher>> {
self.0
.iter()
.map(|x| match x {
@ -255,7 +265,7 @@ impl<H: Hasher> ZerokitMerkleProof<H> for FullMerkleProof<H> {
/// 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<Self::Hasher>) -> FrOf<Self::Hasher> {
self.0.iter().fold(*hash, |hash, branch| match branch {
FullMerkleBranch::Left(sibling) => H::hash(&[hash, *sibling]),
FullMerkleBranch::Right(sibling) => H::hash(&[*sibling, hash]),

43
utils/src/merkle_tree/merkle_tree.rs
View File

@ -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<H: Hasher>
where
H: Hasher,
{
type Proof;
pub type FrOf<H> = <H as Hasher>::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<Self>
where
Self: Sized;
fn new(depth: usize, default_leaf: FrOf<Self::Hasher>, config: Self::Config) -> Result<Self>
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<Self::Hasher>;
fn set(&mut self, index: usize, leaf: FrOf<Self::Hasher>) -> Result<()>;
fn set_range<I>(&mut self, start: usize, leaves: I) -> Result<()>
where
I: IntoIterator<Item = H::Fr>;
fn update_next(&mut self, leaf: H::Fr) -> Result<()>;
I: IntoIterator<Item = FrOf<Self::Hasher>>;
fn update_next(&mut self, leaf: FrOf<Self::Hasher>) -> Result<()>;
fn delete(&mut self, index: usize) -> Result<()>;
fn proof(&self, index: usize) -> Result<Self::Proof>;
fn verify(&self, leaf: &H::Fr, witness: &Self::Proof) -> Result<bool>;
fn verify(&self, leaf: &FrOf<Self::Hasher>, witness: &Self::Proof) -> Result<bool>;
}
pub trait ZerokitMerkleProof<H: Hasher>
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<H::Fr>;
fn get_path_elements(&self) -> Vec<FrOf<Self::Hasher>>;
fn get_path_index(&self) -> Vec<Self::Index>;
fn compute_root_from(&self, leaf: &H::Fr) -> H::Fr;
fn compute_root_from(&self, leaf: &FrOf<Self::Hasher>) -> FrOf<Self::Hasher>;
}

18
utils/src/merkle_tree/optimal_merkle_tree.rs
View File

@ -37,31 +37,33 @@ pub struct OptimalMerkleProof<H: Hasher>(pub Vec<(H::Fr, u8)>);
/// Implementations
impl<H: Hasher> ZerokitMerkleTree<H> for OptimalMerkleTree<H>
impl<H: Hasher> ZerokitMerkleTree for OptimalMerkleTree<H>
where
H: Hasher,
{
type Proof = OptimalMerkleProof<H>;
type Hasher = H;
type Config = ();
fn default(depth: usize) -> Self {
OptimalMerkleTree::<H>::new(depth, H::default_leaf())
fn default(depth: usize) -> Result<Self> {
OptimalMerkleTree::<H>::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<Self> {
let mut cached_nodes: Vec<H::Fr> = 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<H: Hasher> ZerokitMerkleProof<H> for OptimalMerkleProof<H>
impl<H: Hasher> ZerokitMerkleProof for OptimalMerkleProof<H>
where
H: Hasher,
{
type Index = u8;
type Hasher = H;
#[must_use]
// Returns the length of a Merkle proof
fn length(&self) -> usize {

4
utils/src/pm_tree/mod.rs Normal file
View File

@ -0,0 +1,4 @@
pub mod sled_adapter;
pub use self::sled_adapter::*;
pub use pmtree;
pub use sled::*;

75
utils/src/pm_tree/sled_adapter.rs Normal file
View File

@ -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<Self> {
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<Self> {
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<Option<Value>> {
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<DBKey, Value>) -> 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(())
}
}

10
utils/tests/merkle_tree.rs
View File

@ -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::<Keccak256>::new(2, [0; 32]);
let mut tree = FullMerkleTree::<Keccak256>::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::<Keccak256>::new(2, [0; 32]);
let mut tree = OptimalMerkleTree::<Keccak256>::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::<Keccak256>::new(2, [0; 32]);
let mut tree = FullMerkleTree::<Keccak256>::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::<Keccak256>::new(2, [0; 32]);
let mut tree = OptimalMerkleTree::<Keccak256>::new(2, [0; 32], ()).unwrap();
for i in 0..leaves.len() {
// We set the leaves
tree.set(i, leaves[i]).unwrap();