From 2c4b3991269ccc9466470c28cfdc031b60a9fc81 Mon Sep 17 00:00:00 2001 From: Aaryamann Challani <43716372+rymnc@users.noreply.github.com> Date: Fri, 5 May 2023 15:15:33 +0530 Subject: [PATCH] feat(rln): ability to pass rln tree config in ffi (#150) * feat(rln): pass config in rln ffi * fix: for rln-wasm * fix: ffi tests * fix: clippy * fix: test cursor --- Cargo.lock | 2 + rln/Cargo.toml | 1 + rln/src/pm_tree_adapter.rs | 34 +++++++++++- rln/src/public.rs | 58 ++++++++++++++------ rln/tests/ffi.rs | 34 ++++++++---- rln/tests/public.rs | 4 +- utils/Cargo.toml | 1 + utils/src/merkle_tree/full_merkle_tree.rs | 16 +++++- utils/src/merkle_tree/merkle_tree.rs | 4 +- utils/src/merkle_tree/optimal_merkle_tree.rs | 16 +++++- utils/tests/merkle_tree.rs | 19 +++++-- 11 files changed, 150 insertions(+), 39 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 6649951..3d08f7e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2395,6 +2395,7 @@ dependencies = [ "once_cell", "rand", "rand_chacha", + "serde", "serde_json", "sled", "thiserror", @@ -3117,6 +3118,7 @@ dependencies = [ "num-bigint", "num-traits", "pmtree", + "serde", "sled", "tiny-keccak", ] diff --git a/rln/Cargo.toml b/rln/Cargo.toml index 2771644..a6dba75 100644 --- a/rln/Cargo.toml +++ b/rln/Cargo.toml @@ -41,6 +41,7 @@ utils = { path = "../utils/", default-features = false } # serialization serde_json = "1.0.48" +serde = { version = "1.0.130", features = ["derive"] } include_dir = "=0.7.3" diff --git a/rln/src/pm_tree_adapter.rs b/rln/src/pm_tree_adapter.rs index 9aaab57..660b86d 100644 --- a/rln/src/pm_tree_adapter.rs +++ b/rln/src/pm_tree_adapter.rs @@ -2,7 +2,10 @@ 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 serde_json::Value; use std::fmt::Debug; +use std::path::PathBuf; +use std::str::FromStr; use utils::*; pub struct PmTree { @@ -37,10 +40,39 @@ impl pmtree::Hasher for PoseidonHash { } } +fn get_tmp_path() -> std::path::PathBuf { + std::env::temp_dir().join(format!("pmtree-{}", rand::random::())) +} + +fn get_tmp() -> bool { + true +} + pub struct PmtreeConfig(pm_tree::Config); + +impl FromStr for PmtreeConfig { + type Err = Report; + + fn from_str(s: &str) -> Result { + let config: Value = serde_json::from_str(s)?; + + let temporary = config["temporary"].as_bool(); + let path = config["path"].as_str(); + let path = match path { + Some(path) => Some(PathBuf::from(path)), + None => None, + }; + + let config = pm_tree::Config::new() + .temporary(temporary.unwrap_or(get_tmp())) + .path(path.unwrap_or(get_tmp_path())); + Ok(PmtreeConfig(config)) + } +} + impl Default for PmtreeConfig { fn default() -> Self { - let tmp_path = std::env::temp_dir().join(format!("pmtree-{}", rand::random::())); + let tmp_path = get_tmp_path(); PmtreeConfig(pm_tree::Config::new().temporary(true).path(tmp_path)) } } diff --git a/rln/src/public.rs b/rln/src/public.rs index 6b64d69..c715210 100644 --- a/rln/src/public.rs +++ b/rln/src/public.rs @@ -12,8 +12,10 @@ use ark_serialize::{CanonicalDeserialize, CanonicalSerialize, Read, Write}; use cfg_if::cfg_if; use color_eyre::{Report, Result}; use num_bigint::BigInt; +use serde_json::{json, Value}; use std::io::Cursor; -use utils::{ZerokitMerkleProof, ZerokitMerkleTree}; +use std::str::FromStr; +use utils::{Hasher, ZerokitMerkleProof, ZerokitMerkleTree}; cfg_if! { if #[cfg(not(target_arch = "wasm32"))] { @@ -62,7 +64,7 @@ impl RLN<'_> { /// use std::io::Cursor; /// /// let tree_height = 20; - /// let resources = Cursor::new("tree_height_20"); + /// let resources = Cursor::new(json!({"resources_folder": "tree_height_20"}); /// /// // We create a new RLN instance /// let mut rln = RLN::new(tree_height, resources); @@ -73,15 +75,30 @@ impl RLN<'_> { let mut input: Vec = Vec::new(); input_data.read_to_end(&mut input)?; - let resources_folder = String::from_utf8(input)?; + let rln_config: Value = serde_json::from_str(&String::from_utf8(input)?)?; + let resources_folder = rln_config["resources_folder"] + .as_str() + .unwrap_or(TEST_RESOURCES_FOLDER); + let tree_config_opt = rln_config["tree_config"].as_str(); - let witness_calculator = circom_from_folder(&resources_folder)?; + let witness_calculator = circom_from_folder(resources_folder)?; - let proving_key = zkey_from_folder(&resources_folder)?; - let verification_key = vk_from_folder(&resources_folder)?; + let proving_key = zkey_from_folder(resources_folder)?; + let verification_key = vk_from_folder(resources_folder)?; + + let tree_config: ::Config = match tree_config_opt { + Some(tree_config_str) => { + ::Config::from_str(tree_config_str)? + } + None => ::Config::default(), + }; // We compute a default empty tree - let tree = PoseidonTree::default(tree_height)?; + let tree = PoseidonTree::new( + tree_height, + ::Hasher::default_leaf(), + tree_config, + )?; Ok(RLN { witness_calculator, @@ -990,7 +1007,7 @@ impl RLN<'_> { impl Default for RLN<'_> { fn default() -> Self { let tree_height = TEST_TREE_HEIGHT; - let buffer = Cursor::new(TEST_RESOURCES_FOLDER); + let buffer = Cursor::new(json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string()); Self::new(tree_height, buffer).unwrap() } } @@ -1066,6 +1083,7 @@ mod test { use super::*; use ark_std::{rand::thread_rng, UniformRand}; use rand::Rng; + use serde_json::json; use utils::ZerokitMerkleTree; #[test] @@ -1082,7 +1100,8 @@ mod test { } // We create a new tree - let input_buffer = Cursor::new(TEST_RESOURCES_FOLDER); + let input_buffer = + Cursor::new(json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string()); let mut rln = RLN::new(tree_height, input_buffer).unwrap(); // We first add leaves one by one specifying the index @@ -1177,7 +1196,8 @@ mod test { let set_index = rng.gen_range(0..no_of_leaves) as usize; // We create a new tree - let input_buffer = Cursor::new(TEST_RESOURCES_FOLDER); + let input_buffer = + Cursor::new(json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string()); let mut rln = RLN::new(tree_height, input_buffer).unwrap(); // We add leaves in a batch into the tree @@ -1248,7 +1268,8 @@ mod test { let bad_index = (1 << tree_height) - rng.gen_range(0..no_of_leaves) as usize; // We create a new tree - let input_buffer = Cursor::new(TEST_RESOURCES_FOLDER); + let input_buffer = + Cursor::new(json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string()); let mut rln = RLN::new(tree_height, input_buffer).unwrap(); // Get root of empty tree @@ -1277,7 +1298,8 @@ mod test { fn test_groth16_proof() { let tree_height = TEST_TREE_HEIGHT; - let input_buffer = Cursor::new(TEST_RESOURCES_FOLDER); + let input_buffer = + Cursor::new(json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string()); let mut rln = RLN::new(tree_height, input_buffer).unwrap(); // Note: we only test Groth16 proof generation, so we ignore setting the tree in the RLN object @@ -1321,7 +1343,8 @@ mod test { } // We create a new RLN instance - let input_buffer = Cursor::new(TEST_RESOURCES_FOLDER); + let input_buffer = + Cursor::new(json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string()); let mut rln = RLN::new(tree_height, input_buffer).unwrap(); // We add leaves in a batch into the tree @@ -1385,7 +1408,8 @@ mod test { } // We create a new RLN instance - let input_buffer = Cursor::new(TEST_RESOURCES_FOLDER); + let input_buffer = + Cursor::new(json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string()); let mut rln = RLN::new(tree_height, input_buffer).unwrap(); // We add leaves in a batch into the tree @@ -1481,7 +1505,8 @@ mod test { } // We create a new RLN instance - let input_buffer = Cursor::new(TEST_RESOURCES_FOLDER); + let input_buffer = + Cursor::new(json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string()); let mut rln = RLN::new(tree_height, input_buffer).unwrap(); // We add leaves in a batch into the tree @@ -1567,7 +1592,8 @@ mod test { let tree_height = TEST_TREE_HEIGHT; // We create a new RLN instance - let input_buffer = Cursor::new(TEST_RESOURCES_FOLDER); + let input_buffer = + Cursor::new(json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string()); let mut rln = RLN::new(tree_height, input_buffer).unwrap(); // Generate identity pair diff --git a/rln/tests/ffi.rs b/rln/tests/ffi.rs index 40acea6..32f6746 100644 --- a/rln/tests/ffi.rs +++ b/rln/tests/ffi.rs @@ -8,6 +8,7 @@ mod test { use rln::protocol::*; use rln::public::RLN; use rln::utils::*; + use serde_json::json; use std::fs::File; use std::io::Read; use std::mem::MaybeUninit; @@ -28,7 +29,8 @@ mod test { // We create a RLN instance let mut rln_pointer = MaybeUninit::<*mut RLN>::uninit(); - let input_buffer = &Buffer::from(TEST_RESOURCES_FOLDER.as_bytes()); + let input_config = json!({ "resource_folder": TEST_RESOURCES_FOLDER }).to_string(); + let input_buffer = &Buffer::from(input_config.as_bytes()); let success = new(tree_height, input_buffer, rln_pointer.as_mut_ptr()); assert!(success, "RLN object creation failed"); let rln_pointer = unsafe { &mut *rln_pointer.assume_init() }; @@ -135,7 +137,8 @@ mod test { // We create a RLN instance let mut rln_pointer = MaybeUninit::<*mut RLN>::uninit(); - let input_buffer = &Buffer::from(TEST_RESOURCES_FOLDER.as_bytes()); + let input_config = json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string(); + let input_buffer = &Buffer::from(input_config.as_bytes()); let success = new(tree_height, input_buffer, rln_pointer.as_mut_ptr()); assert!(success, "RLN object creation failed"); let rln_pointer = unsafe { &mut *rln_pointer.assume_init() }; @@ -232,7 +235,8 @@ mod test { // We create a RLN instance let mut rln_pointer = MaybeUninit::<*mut RLN>::uninit(); - let input_buffer = &Buffer::from(TEST_RESOURCES_FOLDER.as_bytes()); + let input_config = json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string(); + let input_buffer = &Buffer::from(input_config.as_bytes()); let success = new(tree_height, input_buffer, rln_pointer.as_mut_ptr()); assert!(success, "RLN object creation failed"); let rln_pointer = unsafe { &mut *rln_pointer.assume_init() }; @@ -272,7 +276,8 @@ mod test { // We create a RLN instance let mut rln_pointer = MaybeUninit::<*mut RLN>::uninit(); - let input_buffer = &Buffer::from(TEST_RESOURCES_FOLDER.as_bytes()); + let input_config = json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string(); + let input_buffer = &Buffer::from(input_config.as_bytes()); let success = new(tree_height, input_buffer, rln_pointer.as_mut_ptr()); assert!(success, "RLN object creation failed"); let rln_pointer = unsafe { &mut *rln_pointer.assume_init() }; @@ -440,7 +445,8 @@ mod test { // We create a RLN instance let mut rln_pointer = MaybeUninit::<*mut RLN>::uninit(); - let input_buffer = &Buffer::from(TEST_RESOURCES_FOLDER.as_bytes()); + let input_config = json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string(); + let input_buffer = &Buffer::from(input_config.as_bytes()); let success = new(tree_height, input_buffer, rln_pointer.as_mut_ptr()); assert!(success, "RLN object creation failed"); let rln_pointer = unsafe { &mut *rln_pointer.assume_init() }; @@ -502,7 +508,8 @@ mod test { // We create a RLN instance using a resource folder path let mut rln_pointer = MaybeUninit::<*mut RLN>::uninit(); - let input_buffer = &Buffer::from(TEST_RESOURCES_FOLDER.as_bytes()); + let input_config = json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string(); + let input_buffer = &Buffer::from(input_config.as_bytes()); let success = new(tree_height, input_buffer, rln_pointer.as_mut_ptr()); assert!(success, "RLN object creation failed"); let rln_pointer = unsafe { &mut *rln_pointer.assume_init() }; @@ -582,7 +589,8 @@ mod test { // We create a RLN instance let mut rln_pointer = MaybeUninit::<*mut RLN>::uninit(); - let input_buffer = &Buffer::from(TEST_RESOURCES_FOLDER.as_bytes()); + let input_config = json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string(); + let input_buffer = &Buffer::from(input_config.as_bytes()); let success = new(tree_height, input_buffer, rln_pointer.as_mut_ptr()); assert!(success, "RLN object creation failed"); let rln_pointer = unsafe { &mut *rln_pointer.assume_init() }; @@ -666,7 +674,8 @@ mod test { // We create a RLN instance let mut rln_pointer = MaybeUninit::<*mut RLN>::uninit(); - let input_buffer = &Buffer::from(TEST_RESOURCES_FOLDER.as_bytes()); + let input_config = json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string(); + let input_buffer = &Buffer::from(input_config.as_bytes()); let success = new(tree_height, input_buffer, rln_pointer.as_mut_ptr()); assert!(success, "RLN object creation failed"); let rln_pointer = unsafe { &mut *rln_pointer.assume_init() }; @@ -784,7 +793,8 @@ mod test { // We create a RLN instance let mut rln_pointer = MaybeUninit::<*mut RLN>::uninit(); - let input_buffer = &Buffer::from(TEST_RESOURCES_FOLDER.as_bytes()); + let input_config = json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string(); + let input_buffer = &Buffer::from(input_config.as_bytes()); let success = new(tree_height, input_buffer, rln_pointer.as_mut_ptr()); assert!(success, "RLN object creation failed"); let rln_pointer = unsafe { &mut *rln_pointer.assume_init() }; @@ -943,7 +953,8 @@ mod test { // We create a RLN instance let mut rln_pointer = MaybeUninit::<*mut RLN>::uninit(); - let input_buffer = &Buffer::from(TEST_RESOURCES_FOLDER.as_bytes()); + let input_config = json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string(); + let input_buffer = &Buffer::from(input_config.as_bytes()); let success = new(tree_height, input_buffer, rln_pointer.as_mut_ptr()); assert!(success, "RLN object creation failed"); let rln_pointer = unsafe { &mut *rln_pointer.assume_init() }; @@ -982,7 +993,8 @@ mod test { let tree_height = TEST_TREE_HEIGHT; // We create a RLN instance let mut rln_pointer = MaybeUninit::<*mut RLN>::uninit(); - let input_buffer = &Buffer::from(TEST_RESOURCES_FOLDER.as_bytes()); + let input_config = json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string(); + let input_buffer = &Buffer::from(input_config.as_bytes()); let success = new(tree_height, input_buffer, rln_pointer.as_mut_ptr()); assert!(success, "RLN object creation failed"); let rln_pointer = unsafe { &mut *rln_pointer.assume_init() }; diff --git a/rln/tests/public.rs b/rln/tests/public.rs index d1bfd67..b6ba276 100644 --- a/rln/tests/public.rs +++ b/rln/tests/public.rs @@ -7,6 +7,7 @@ mod test { use rln::protocol::{compute_tree_root, deserialize_identity_tuple}; use rln::public::{hash as public_hash, poseidon_hash as public_poseidon_hash, RLN}; use rln::utils::*; + use serde_json::json; use std::io::Cursor; #[test] @@ -15,7 +16,8 @@ mod test { let tree_height = TEST_TREE_HEIGHT; let leaf_index = 3; - let input_buffer = Cursor::new(TEST_RESOURCES_FOLDER); + let input_buffer = + Cursor::new(json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string()); let mut rln = RLN::new(tree_height, input_buffer).unwrap(); // generate identity diff --git a/utils/Cargo.toml b/utils/Cargo.toml index 74ff44e..462a02a 100644 --- a/utils/Cargo.toml +++ b/utils/Cargo.toml @@ -10,6 +10,7 @@ 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" +serde = "1.0.44" [dev-dependencies] ark-bn254 = "=0.4.0" diff --git a/utils/src/merkle_tree/full_merkle_tree.rs b/utils/src/merkle_tree/full_merkle_tree.rs index 19fc295..5a248ff 100644 --- a/utils/src/merkle_tree/full_merkle_tree.rs +++ b/utils/src/merkle_tree/full_merkle_tree.rs @@ -4,6 +4,7 @@ use std::{ cmp::max, fmt::Debug, iter::{once, repeat, successors}, + str::FromStr, }; //////////////////////////////////////////////////////////// @@ -44,6 +45,17 @@ pub enum FullMerkleBranch { #[derive(Clone, PartialEq, Eq)] pub struct FullMerkleProof(pub Vec>); +#[derive(Default)] +pub struct FullMerkleConfig(()); + +impl FromStr for FullMerkleConfig { + type Err = Report; + + fn from_str(_s: &str) -> Result { + Ok(FullMerkleConfig::default()) + } +} + /// Implementations impl ZerokitMerkleTree for FullMerkleTree where @@ -51,10 +63,10 @@ where { type Proof = FullMerkleProof; type Hasher = H; - type Config = (); + type Config = FullMerkleConfig; fn default(depth: usize) -> Result { - FullMerkleTree::::new(depth, Self::Hasher::default_leaf(), ()) + FullMerkleTree::::new(depth, Self::Hasher::default_leaf(), Self::Config::default()) } /// Creates a new `MerkleTree` diff --git a/utils/src/merkle_tree/merkle_tree.rs b/utils/src/merkle_tree/merkle_tree.rs index e831dda..d2a112f 100644 --- a/utils/src/merkle_tree/merkle_tree.rs +++ b/utils/src/merkle_tree/merkle_tree.rs @@ -13,6 +13,8 @@ //! * Disk based storage backend (using mmaped files should be easy) //! * Implement serialization for tree and Merkle proof +use std::str::FromStr; + use color_eyre::Result; /// In the Hasher trait we define the node type, the default leaf @@ -35,7 +37,7 @@ pub type FrOf = ::Fr; pub trait ZerokitMerkleTree { type Proof: ZerokitMerkleProof; type Hasher: Hasher; - type Config: Default; + type Config: Default + FromStr; fn default(depth: usize) -> Result where diff --git a/utils/src/merkle_tree/optimal_merkle_tree.rs b/utils/src/merkle_tree/optimal_merkle_tree.rs index 71ab1b7..989da97 100644 --- a/utils/src/merkle_tree/optimal_merkle_tree.rs +++ b/utils/src/merkle_tree/optimal_merkle_tree.rs @@ -1,6 +1,7 @@ use crate::merkle_tree::{Hasher, ZerokitMerkleProof, ZerokitMerkleTree}; use color_eyre::{Report, Result}; use std::collections::HashMap; +use std::str::FromStr; use std::{cmp::max, fmt::Debug}; //////////////////////////////////////////////////////////// @@ -35,6 +36,17 @@ where #[derive(Clone, PartialEq, Eq)] pub struct OptimalMerkleProof(pub Vec<(H::Fr, u8)>); +#[derive(Default)] +pub struct OptimalMerkleConfig(()); + +impl FromStr for OptimalMerkleConfig { + type Err = Report; + + fn from_str(_s: &str) -> Result { + Ok(OptimalMerkleConfig::default()) + } +} + /// Implementations impl ZerokitMerkleTree for OptimalMerkleTree @@ -43,10 +55,10 @@ where { type Proof = OptimalMerkleProof; type Hasher = H; - type Config = (); + type Config = OptimalMerkleConfig; fn default(depth: usize) -> Result { - OptimalMerkleTree::::new(depth, H::default_leaf(), ()) + OptimalMerkleTree::::new(depth, H::default_leaf(), Self::Config::default()) } /// Creates a new `MerkleTree` diff --git a/utils/tests/merkle_tree.rs b/utils/tests/merkle_tree.rs index 68ba819..1c99046 100644 --- a/utils/tests/merkle_tree.rs +++ b/utils/tests/merkle_tree.rs @@ -3,7 +3,10 @@ mod test { use hex_literal::hex; use tiny_keccak::{Hasher as _, Keccak}; - use utils::{FullMerkleTree, Hasher, OptimalMerkleTree, ZerokitMerkleProof, ZerokitMerkleTree}; + use utils::{ + FullMerkleConfig, FullMerkleTree, Hasher, OptimalMerkleConfig, OptimalMerkleTree, + ZerokitMerkleProof, ZerokitMerkleTree, + }; #[derive(Clone, Copy, Eq, PartialEq)] struct Keccak256; @@ -44,14 +47,17 @@ mod test { hex!("a9bb8c3f1f12e9aa903a50c47f314b57610a3ab32f2d463293f58836def38d36"), ]; - let mut tree = FullMerkleTree::::new(2, [0; 32], ()).unwrap(); + let mut tree = + FullMerkleTree::::new(2, [0; 32], FullMerkleConfig::default()).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], ()).unwrap(); + let mut tree = + OptimalMerkleTree::::new(2, [0; 32], OptimalMerkleConfig::default()) + .unwrap(); assert_eq!(tree.root(), default_tree_root); for i in 0..leaves.len() { tree.set(i, leaves[i]).unwrap(); @@ -69,7 +75,8 @@ mod test { ]; // We thest the FullMerkleTree implementation - let mut tree = FullMerkleTree::::new(2, [0; 32], ()).unwrap(); + let mut tree = + FullMerkleTree::::new(2, [0; 32], FullMerkleConfig::default()).unwrap(); for i in 0..leaves.len() { // We set the leaves tree.set(i, leaves[i]).unwrap(); @@ -93,7 +100,9 @@ mod test { } // We test the OptimalMerkleTree implementation - let mut tree = OptimalMerkleTree::::new(2, [0; 32], ()).unwrap(); + let mut tree = + OptimalMerkleTree::::new(2, [0; 32], OptimalMerkleConfig::default()) + .unwrap(); for i in 0..leaves.len() { // We set the leaves tree.set(i, leaves[i]).unwrap();