Merge pull request #63 from okx/feature/add-merkle-proof-sdk

Feature/add merkle proof sdk

.gitignore

@@ -7,4 +7,5 @@ config/local.toml
 config/prod.toml
 test-data/user-data*/
 my_permanent_leveldb/
-*level_db*/
+*level_db*/
+*.json

README.md

@@ -35,14 +35,27 @@ output_proof_path="global_proof.json"
 
 cargo run --release --package zk-por-cli --bin zk-por-cli prove --cfg-path ${cfg_dir_path} --output-path ${output_proof_path}
 ```
+
+- get-merkle-proof
+```
+cargo run --release --package zk-por-cli --bin zk-por-cli get-merkle-proof --account-path account.json --output-path merkle_proof.json --cfg-path config
+```
+
 - verify
 ```
 global_root_path="global_proof.json"
 
 # optional. If not provided, will skip verifying the inclusion
-arg_inclusion_proof_path="--inclusion-proof-path inclusion_proof.json"
+inclusion_proof_path="merkle_proof.json"
 
-cargo run --features zk-por-core/verifier --release --package zk-por-cli --bin zk-por-cli verify --global-proof-path ${global_root_path} ${arg_inclusion_proof_path}
+cargo run --features zk-por-core/verifier --release --package zk-por-cli --bin zk-por-cli verify --global-proof-path ${global_root_path} --inclusion-proof-path ${inclusion_proof_path} --root 11288199779358641579,2344540219612146741,6809171731163302525,17936043556479519168
+```
+
+## cli
+```
+./target/release/zk-por-cli --help
+./target/release/zk-por-cli prove --cfg-path ${cfg_dir_path} --output-path ${output_proof_path}
+./target/release/zk-por-cli get-merkle-proof --account-path account.json --output-path merkle_proof.json --cfg-path config
 ```
 
 ## code coverage

crates/zk-por-cli/src/lib.rs

@@ -1,3 +1,4 @@
 mod constant;
+pub mod merkle_proof;
 pub mod prover;
 pub mod verifier;

crates/zk-por-cli/src/main.rs

@@ -1,8 +1,8 @@
 use std::{path::PathBuf, str::FromStr};
 
 use clap::{Parser, Subcommand};
-use zk_por_cli::{prover::prove, verifier::verify};
-use zk_por_core::error::PoRError;
+use zk_por_cli::{merkle_proof::get_merkle_proof, prover::prove, verifier::verify};
+use zk_por_core::{config::ProverConfig, error::PoRError};
 
 #[derive(Parser)]
 #[command(version, about, long_about = None)]
@@ -25,13 +25,19 @@ pub enum ZkPorCommitCommands {
     },
     GetMerkleProof {
         #[arg(short, long)]
-        user_id: String,
+        cfg_path: String, // path to config file
+        #[arg(short, long)]
+        account_path: String,
+        #[arg(short, long)]
+        output_path: String, // path to output file
     },
     Verify {
         #[arg(short, long)]
         global_proof_path: String,
         #[arg(short, long)]
         inclusion_proof_path: Option<String>,
+        #[arg(short, long)]
+        root: Option<String>,
     },
 }
 
@@ -45,16 +51,17 @@ impl Execute for ZkPorCommitCommands {
                 let output_file = PathBuf::from_str(&output_path).unwrap();
                 prove(prover_cfg, output_file)
             }
-            ZkPorCommitCommands::GetMerkleProof { user_id } => {
-                // TODO: implement this
-                _ = user_id;
-                Ok(())
+            ZkPorCommitCommands::GetMerkleProof { cfg_path, account_path, output_path } => {
+                let cfg = zk_por_core::config::ProverConfig::load(&cfg_path)
+                    .map_err(|e| PoRError::ConfigError(e))?;
+                let prover_cfg: ProverConfig = cfg.try_deserialize().unwrap();
+                get_merkle_proof(account_path.to_string(), prover_cfg, output_path.to_string())
             }
-            ZkPorCommitCommands::Verify { global_proof_path, inclusion_proof_path } => {
+            ZkPorCommitCommands::Verify { global_proof_path, inclusion_proof_path, root } => {
                 let global_proof_path = PathBuf::from_str(&global_proof_path).unwrap();
                 let inclusion_proof_path =
                     inclusion_proof_path.as_ref().map(|p| PathBuf::from_str(&p).unwrap());
-                verify(global_proof_path, inclusion_proof_path)
+                verify(global_proof_path, inclusion_proof_path, root.clone())
             }
         }
     }

crates/zk-por-cli/src/merkle_proof.rs

@@ -0,0 +1,67 @@
+use std::{fs::File, io::Write, str::FromStr};
+
+use serde_json::json;
+use zk_por_core::{
+    account::Account,
+    config::ProverConfig,
+    database::{DataBase, DbOption},
+    error::PoRError,
+    global::GlobalConfig,
+    merkle_proof::MerkleProof,
+    parser::{AccountParser, FileAccountReader, FileManager, FilesCfg},
+};
+
+use crate::constant::RECURSION_BRANCHOUT_NUM;
+
+pub fn get_merkle_proof(
+    account_path: String,
+    cfg: ProverConfig,
+    output_path: String,
+) -> Result<(), PoRError> {
+    let database = DataBase::new(DbOption {
+        user_map_dir: cfg.db.level_db_user_path.to_string(),
+        gmst_dir: cfg.db.level_db_gmst_path.to_string(),
+    });
+
+    let batch_size = cfg.prover.batch_size as usize;
+    let token_num = cfg.prover.num_of_tokens as usize;
+
+    // the path to dump the final generated proof
+    let file_manager = FileManager {};
+    let account_parser = FileAccountReader::new(
+        FilesCfg {
+            dir: std::path::PathBuf::from_str(&cfg.prover.user_data_path).unwrap(),
+            batch_size: cfg.prover.batch_size,
+            num_of_tokens: cfg.prover.num_of_tokens,
+        },
+        &file_manager,
+    );
+    account_parser.log_state();
+    // let mut account_parser: Box<dyn AccountParser> = Box::new(parser);
+
+    let batch_num = account_parser.total_num_of_users().div_ceil(batch_size);
+
+    let global_cfg = GlobalConfig {
+        num_of_tokens: token_num,
+        num_of_batches: batch_num,
+        batch_size: batch_size,
+        recursion_branchout_num: RECURSION_BRANCHOUT_NUM,
+    };
+
+    // the account json format is a map of tokens to token values, wheras the format in the merkle proof is given by a vec of token values.
+    let account = Account::new_from_file_path(account_path);
+
+    if account.is_err() {
+        return Err(account.unwrap_err());
+    }
+
+    let merkle_proof = MerkleProof::new_from_account(&account.unwrap(), &database, &global_cfg)
+        .expect("Una ble to generate merkle proof");
+
+    let mut file = File::create(output_path.clone())
+        .expect(format!("fail to create proof file at {:#?}", output_path).as_str());
+    file.write_all(json!(merkle_proof).to_string().as_bytes())
+        .expect("fail to write proof to file");
+
+    Ok(())
+}

crates/zk-por-cli/src/prover.rs

@@ -110,7 +110,7 @@ pub fn prove(cfg: ProverConfig, proof_output_path: PathBuf) -> Result<(), PoRErr
             accounts.resize(account_num + pad_num, Account::get_empty_account(token_num));
         }
 
-        assert_eq!(account_num % batch_size, 0);
+        assert_eq!(accounts.len() % batch_size, 0);
 
         tracing::debug!(
             "parse {} times, with number of accounts {}, number of batches {}",
@@ -307,7 +307,10 @@ pub fn prove(cfg: ProverConfig, proof_output_path: PathBuf) -> Result<(), PoRErr
     // persist gmst to database
 
     let global_mst = GLOBAL_MST.get().unwrap();
+
     let _g = global_mst.read().expect("unable to get a lock");
+    let root_hash = _g.get_root().expect("no root");
+    tracing::info!("root hash is {:?}", root_hash);
     let start = std::time::Instant::now();
     _g.persist(&mut database);
     tracing::info!("persist gmst to db in {:?}", start.elapsed());

crates/zk-por-cli/src/verifier.rs

@@ -7,12 +7,15 @@ use zk_por_core::{
     circuit_config::{get_recursive_circuit_configs, STANDARD_CONFIG},
     circuit_registry::registry::CircuitRegistry,
     error::PoRError,
+    merkle_proof::MerkleProof,
+    util::get_hash_from_hash_string,
     Proof,
 };
 
 pub fn verify(
     global_proof_path: PathBuf,
     merkle_inclusion_path: Option<PathBuf>,
+    root: Option<String>,
 ) -> Result<(), PoRError> {
     let proof_file = File::open(&global_proof_path).unwrap();
     let reader = std::io::BufReader::new(proof_file);
@@ -63,10 +66,28 @@ pub fn verify(
     }
     println!("successfully verify the global proof for round {}", round_num);
 
-    // TODO: verify the inclusion proof
     if let Some(merkle_inclusion_path) = merkle_inclusion_path {
-        _ = merkle_inclusion_path;
-        println!("successfully verify the inclusion proof for user for round {}", round_num);
+        let merkle_path = File::open(&merkle_inclusion_path).unwrap();
+        let reader = std::io::BufReader::new(merkle_path);
+
+        // Parse the JSON as Proof
+        let proof: MerkleProof = from_reader(reader).unwrap();
+
+        if root.is_none() {
+            return Err(PoRError::InvalidParameter(
+                "Require root for merkle proof verification".to_string(),
+            ));
+        }
+
+        let res = proof.verify_merkle_proof(get_hash_from_hash_string(root.unwrap()));
+
+        if res.is_err() {
+            let res_err = res.unwrap_err();
+            return Err(res_err);
+        } else {
+            println!("successfully verify the inclusion proof for user for round {}", round_num);
+            return Ok(());
+        }
     }
 
     Ok(())

crates/zk-por-core/src/account.rs

@@ -3,15 +3,20 @@ use plonky2::{
     plonk::config::Hasher,
 };
 use plonky2_field::types::Field;
+use serde::{Deserialize, Serialize};
+use serde_json::Value;
+use std::{collections::BTreeMap, fs::File, io::BufReader};
 
 use crate::{
     database::{DataBase, UserId},
+    error::PoRError,
+    parser::parse_account_state,
     types::F,
 };
 use rand::Rng;
 
 /// A struct representing a users account. It represents their equity and debt as a Vector of goldilocks field elements.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct Account {
     pub id: String, // 256 bit hex string
     pub equity: Vec<F>,
@@ -34,6 +39,14 @@ impl Account {
         hash
     }
 
+    pub fn get_empty_account_with_user_id(user_id: String, num_of_tokens: usize) -> Account {
+        Self {
+            id: user_id,
+            equity: vec![F::default(); num_of_tokens],
+            debt: vec![F::default(); num_of_tokens],
+        }
+    }
+
     pub fn get_empty_account(num_of_tokens: usize) -> Account {
         Self {
             id: "0".repeat(64),
@@ -59,6 +72,27 @@ impl Account {
             .map(|seg| F::from_canonical_u64(u64::from_str_radix(seg, 16).unwrap()))
             .collect::<Vec<F>>()
     }
+
+    /// Get a new account from a file path
+    pub fn new_from_file_path(path: String) -> Result<Self, PoRError> {
+        let file_res = File::open(path);
+
+        if file_res.is_err() {
+            return Err(PoRError::InvalidParameter("Invalid account json file path".to_string()));
+        }
+
+        let reader = BufReader::new(file_res.unwrap());
+        // Deserialize the json data to a struct
+        let account_map_res: Result<BTreeMap<String, Value>, _> = serde_json::from_reader(reader);
+
+        if account_map_res.is_err() {
+            return Err(PoRError::InvalidParameter("Invalid account json".to_string()));
+        }
+
+        let account = parse_account_state(&account_map_res.unwrap());
+
+        Ok(account)
+    }
 }
 
 pub fn persist_account_id_to_gmst_pos(
@@ -70,12 +104,9 @@ pub fn persist_account_id_to_gmst_pos(
         .iter()
         .enumerate()
         .map(|(i, acct)| {
-            let hex_decode = hex::decode(&acct.id).unwrap();
-            assert_eq!(hex_decode.len(), 32);
-            let mut array = [0u8; 32];
-            array.copy_from_slice(&hex_decode);
-
-            (UserId(array), (i + start_idx) as u32)
+            let user_id = UserId::from_hex_string(acct.id.to_string()).unwrap();
+            // tracing::debug!("persist account {:?} with index: {:?}", acct.id, i + start_idx);
+            (user_id, (i + start_idx) as u32)
         })
         .collect::<Vec<(UserId, u32)>>();
     db.add_batch_users(user_batch);
@@ -105,6 +136,10 @@ pub fn gen_accounts_with_random_data(num_accounts: usize, num_assets: usize) ->
 }
 
 pub fn gen_empty_accounts(batch_size: usize, num_assets: usize) -> Vec<Account> {
-    let accounts = vec![Account::get_empty_account(num_assets); batch_size];
+    let accounts =
+        vec![
+            Account::get_empty_account_with_user_id(UserId::rand().to_string(), num_assets);
+            batch_size
+        ];
     accounts
 }

crates/zk-por-core/src/config.rs

@@ -40,6 +40,22 @@ pub struct ConfigDb {
     pub level_db_gmst_path: String,
 }
 
+impl ConfigDb {
+    pub fn load(dir: &str) -> Result<Config, ConfigError> {
+        let env = std::env::var("ENV").unwrap_or("default".into());
+        Config::builder()
+            // .add_source(File::with_name(&format!("{}/default", dir)))
+            .add_source(File::with_name(&format!("{}/{}", dir, env)).required(false))
+            .add_source(File::with_name(&format!("{}/local", dir)).required(false))
+            .add_source(config::Environment::with_prefix("ZKPOR"))
+            .build()
+    }
+    pub fn try_new() -> Result<Self, ConfigError> {
+        let config = Self::load("config")?;
+        config.try_deserialize()
+    }
+}
+
 #[derive(Debug, Clone, Deserialize)]
 pub struct ProverConfig {
     pub log: ConfigLog,

crates/zk-por-core/src/database.rs

@@ -1,10 +1,11 @@
 use std::str::FromStr;
 
+use hex::ToHex;
 use plonky2::{hash::hash_types::HashOut, plonk::config::GenericHashOut};
 use rand::Rng;
 use zk_por_db::LevelDb;
 
-use crate::types::F;
+use crate::{error::PoRError, types::F};
 
 #[derive(Debug, Clone, Copy)]
 pub struct UserId(pub [u8; 32]);
@@ -16,6 +17,29 @@ impl UserId {
         rng.fill(&mut bytes);
         Self(bytes)
     }
+
+    pub fn to_string(&self) -> String {
+        self.0.encode_hex()
+    }
+
+    pub fn from_hex_string(hex_str: String) -> Result<Self, PoRError> {
+        if hex_str.len() != 64 {
+            tracing::error!("User Id: {:?} is not a valid id, length is not 256 bits", hex_str);
+            return Err(PoRError::InvalidParameter(hex_str));
+        }
+
+        let decode_res = hex::decode(hex_str.clone());
+
+        if decode_res.is_err() {
+            tracing::error!("User Id: {:?} is not a valid id", hex_str);
+            return Err(PoRError::InvalidParameter(hex_str));
+        }
+
+        let mut arr = [0u8; 32];
+        arr.copy_from_slice(&decode_res.unwrap());
+
+        Ok(UserId { 0: arr })
+    }
 }
 
 impl db_key::Key for UserId {