Просмотр исходного кода

event_graph: cleanup and refactor poc memory expensive fields

dasman 4 месяцев назад
Родитель
Сommit
29b559fc5c

+ 7 - 25
bin/darkirc/src/crypto/rln.rs

@@ -19,7 +19,10 @@
 use std::time::UNIX_EPOCH;
 
 use darkfi::{
-    event_graph::Event,
+    event_graph::{
+        rln::{closest_epoch, hash_event},
+        Event,
+    },
     zk::{
         halo2::{Field, Value},
         Proof, ProvingKey, Witness, ZkCircuit,
@@ -28,7 +31,7 @@ use darkfi::{
     Result,
 };
 use darkfi_sdk::{
-    crypto::{pasta_prelude::FromUniformBytes, poseidon_hash, smt::SmtMemoryFp},
+    crypto::{poseidon_hash, smt::SmtMemoryFp},
     pasta::pallas,
 };
 use darkfi_serial::{async_trait, SerialDecodable, SerialEncodable};
@@ -38,11 +41,6 @@ use tracing::info;
 pub const RLN_TRAPDOOR_DERIVATION_PATH: pallas::Base = pallas::Base::from_raw([4211, 0, 0, 0]);
 pub const RLN_NULLIFIER_DERIVATION_PATH: pallas::Base = pallas::Base::from_raw([4212, 0, 0, 0]);
 
-/// RLN epoch genesis in millis
-pub const RLN_GENESIS: u64 = 1_738_688_400_000;
-/// RLN epoch length in millis
-pub const RLN_EPOCH_LEN: u64 = 600_000; // 10 min
-
 pub const RLN2_REGISTER_ZKBIN: &[u8] =
     include_bytes!("../../../../src/event_graph/proof/rlnv2-diff-register.zk.bin");
 pub const RLN2_SIGNAL_ZKBIN: &[u8] =
@@ -51,21 +49,6 @@ pub const RLN2_SIGNAL_ZKBIN: &[u8] =
 /// TODO: this is arbitrary it should be based on stake
 pub const MAX_MSG_LIMIT: u64 = 100;
 
-/// Find closest epoch to given timestamp
-pub fn closest_epoch(timestamp: u64) -> u64 {
-    let time_diff = timestamp - RLN_GENESIS;
-    let epoch_idx = time_diff as f64 / RLN_EPOCH_LEN as f64;
-    let rounded = epoch_idx.round() as i64;
-    RLN_GENESIS + (rounded * RLN_EPOCH_LEN as i64) as u64
-}
-
-/// Hash message/event modulo `Fp`
-pub fn hash_event(event: &Event) -> pallas::Base {
-    let mut buf = [0u8; 64];
-    buf[..blake3::OUT_LEN].copy_from_slice(event.header.id().as_bytes());
-    pallas::Base::from_uniform_bytes(&buf)
-}
-
 #[derive(Copy, Clone, SerialEncodable, SerialDecodable)]
 pub struct RlnIdentity {
     pub nullifier: pallas::Base,
@@ -118,8 +101,7 @@ impl RlnIdentity {
         let register_zkbin = ZkBinary::decode(RLN2_REGISTER_ZKBIN, false)?;
         let register_circuit = ZkCircuit::new(witnesses, &register_zkbin);
 
-        let proof =
-            Proof::create(&register_pk, &[register_circuit], &public_inputs, &mut OsRng).unwrap();
+        let proof = Proof::create(&register_pk, &[register_circuit], &public_inputs, &mut OsRng)?;
 
         let leaf = vec![commitment];
         let leaf: Vec<_> = leaf.into_iter().map(|l| (l, l)).collect();
@@ -168,7 +150,7 @@ impl RlnIdentity {
 
         info!(target: "crypto::rln::create_signal_proof", "[RLN] Creating signal proof for event {}", event.header.id());
         let signal_zkbin = ZkBinary::decode(RLN2_SIGNAL_ZKBIN, false)?;
-        let signal_circuit = ZkCircuit::new(witnesses.clone(), &signal_zkbin);
+        let signal_circuit = ZkCircuit::new(witnesses, &signal_zkbin);
 
         let proof = Proof::create(signal_pk, &[signal_circuit], &public_inputs, &mut OsRng)?;
         Ok((proof, y, internal_nullifier, self.user_message_limit))

+ 38 - 34
bin/darkirc/src/irc/client.rs

@@ -29,7 +29,7 @@ use std::{
 use darkfi::{
     event_graph::{
         proto::EventPut,
-        rln::{process_commitment, RLNNode},
+        rln::{closest_epoch, process_commitment, Blob, RLNNode},
         Event, NULL_ID,
     },
     system::Subscription,
@@ -53,7 +53,7 @@ use super::{
     server::{IrcServer, MAX_MSG_LEN},
     NickServ, Privmsg, SERVER_NAME,
 };
-use crate::crypto::rln::{closest_epoch, RlnIdentity, RLN2_SIGNAL_ZKBIN};
+use crate::crypto::rln::{RlnIdentity, RLN2_SIGNAL_ZKBIN};
 
 const PENALTY_LIMIT: usize = 5;
 
@@ -216,34 +216,34 @@ impl Client {
                                         return Err(e)
                                     }
 
-                                    // If we have a RLN identity, now we'll build a ZK proof.
-                                    // Also I really want GOTO in Rust... Fags.
-                                    if let Some(ref mut rln_identity) = *self.server.rln_identity.write().await {
-                                        // If the current epoch is different, we can reset the message counter
-                                        if rln_identity.last_epoch != closest_epoch(event.header.timestamp) {
-                                            rln_identity.last_epoch = closest_epoch(event.header.timestamp);
-                                            rln_identity.message_id = 0;
-                                        }
-
-                                        rln_identity.message_id += 1;
-
-                                        let (proof, y, internal_nullifier, user_msg_limit) = match self.create_rln_signal_proof(&rln_identity, &event).await {
-                                            Ok(v) => v,
-                                            Err(e) => {
-                                                // TODO: Send a message to the IRC client telling that sending went wrong
-                                                error!("[IRC CLIENT] Failed creating RLN signal proof: {e}");
-                                                // Just use an empty "proof"
-                                                (Proof::new(vec![]), pallas::Base::from(0), pallas::Base::from(0), 0)
+                                    let blob = {
+                                        // If we have a RLN identity, now we'll build a ZK proof.
+                                        // Also I really want GOTO in Rust... Fags.
+                                        if let Some(ref mut rln_identity) = *self.server.rln_identity.write().await {
+                                            if rln_identity.last_epoch != closest_epoch(event.header.timestamp) {
+                                                rln_identity.last_epoch = closest_epoch(event.header.timestamp);
+                                                rln_identity.message_id = 0;
                                             }
-                                        };
 
-                                        let blob = serialize_async(&(proof, y, internal_nullifier, user_msg_limit)).await;
+                                            rln_identity.message_id += 1;
+
+                                            let (proof, y, internal_nullifier, user_msg_limit) = match self.create_rln_signal_proof(&rln_identity, &event).await {
+                                                Ok(v) => v,
+                                                Err(e) => {
+                                                    // TODO: Send a message to the IRC client telling that sending went wrong
+                                                    error!("[IRC CLIENT] Failed creating RLN signal proof: {e}");
+                                                    return Err(e)
+                                                }
+                                            };
+
+                                            let blob = Blob{ proof, y, internal_nullifier, user_msg_limit };
+                                            serialize_async(&blob).await
+                                        } else {
+                                            vec![]
+                                        }
+                                    };
 
-                                        self.server.darkirc.p2p.broadcast(&EventPut(event, blob)).await;
-                                    } else {
-                                        // Broadcast it
-                                        self.server.darkirc.p2p.broadcast(&EventPut(event, vec![])).await;
-                                    }
+                                    self.server.darkirc.p2p.broadcast(&EventPut(event, blob)).await;
                                 }
                             }
                         }
@@ -596,14 +596,18 @@ impl Client {
         // Retrieve the ZK proving key from the db
         let signal_zkbin = ZkBinary::decode(RLN2_SIGNAL_ZKBIN, false)?;
         let signal_circuit = ZkCircuit::new(empty_witnesses(&signal_zkbin)?, &signal_zkbin);
-        let Some(proving_key) = self.server.server_store.get("rlnv2-diff-signal-pk")? else {
-            return Err(Error::DatabaseError(
-                "RLN signal proving key not found in server store".to_string(),
-            ))
+        let signal_pk = {
+            let Some(proving_key) = self.server.server_store.get("rlnv2-diff-signal-pk")? else {
+                return Err(Error::DatabaseError(
+                    "RLN signal proving key not found in server store".to_string(),
+                ))
+            };
+            let mut reader = Cursor::new(&*proving_key);
+            let pk = ProvingKey::read(&mut reader, signal_circuit)?;
+            drop(proving_key);
+            pk
         };
-        let mut reader = Cursor::new(proving_key);
-        let proving_key = ProvingKey::read(&mut reader, signal_circuit)?;
 
-        rln_identity.create_signal_proof(event, &identity_tree, &proving_key)
+        rln_identity.create_signal_proof(event, &identity_tree, &signal_pk)
     }
 }

+ 5 - 2
bin/darkirc/src/irc/services/nickserv.rs

@@ -19,7 +19,10 @@
 use std::{io::Cursor, str::SplitAsciiWhitespace, sync::Arc, time::UNIX_EPOCH};
 
 use darkfi::{
-    event_graph::{rln::RLNNode, Event},
+    event_graph::{
+        rln::{closest_epoch, RLNNode},
+        Event,
+    },
     zk::{empty_witnesses, ProvingKey, ZkCircuit},
     zkas::ZkBinary,
     Error, Result,
@@ -30,7 +33,7 @@ use smol::lock::RwLock;
 
 use super::super::{client::ReplyType, rpl::*};
 use crate::{
-    crypto::rln::{closest_epoch, RlnIdentity, RLN2_REGISTER_ZKBIN},
+    crypto::rln::{RlnIdentity, RLN2_REGISTER_ZKBIN},
     IrcServer,
 };
 

+ 9 - 2
bin/darkirc/src/main.rs

@@ -95,10 +95,10 @@ struct Args {
     irc_tls_secret: Option<String>,
 
     /// How many DAGs to sync.
-    #[structopt(short, long, default_value = "8")]
+    #[structopt(long, default_value = "8")]
     dags_count: usize,
 
-    #[structopt(short, long, default_value = "~/.local/share/darkfi/darkirc_db")]
+    #[structopt(long, default_value = "~/.local/share/darkfi/darkirc_db")]
     /// Datastore (DB) path
     datastore: String,
 
@@ -193,6 +193,13 @@ impl DarkIrc {
     }
 }
 
+// #[global_allocator]
+// static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
+
+// #[allow(non_upper_case_globals)]
+// #[export_name = "malloc_conf"]
+// pub static malloc_conf: &[u8] = b"dirty_decay_ms:1000,muzzy_decay_ms:1000\0";
+
 async_daemonize!(realmain);
 async fn realmain(args: Args, ex: Arc<Executor<'static>>) -> Result<()> {
     if args.fast_mode {

+ 2 - 2
bin/darkirc/src/settings.rs

@@ -23,12 +23,12 @@ use std::{
 };
 
 use crypto_box::PublicKey;
-use darkfi::{Error::ParseFailed, Result};
+use darkfi::{event_graph::rln::closest_epoch, Error::ParseFailed, Result};
 use darkfi_sdk::{crypto::pasta_prelude::PrimeField, pasta::pallas};
 use tracing::info;
 
 use crate::{
-    crypto::rln::{closest_epoch, RlnIdentity},
+    crypto::rln::RlnIdentity,
     irc::{IrcChannel, IrcContact},
 };
 

+ 10 - 19
src/event_graph/mod.rs

@@ -50,7 +50,6 @@ use crate::{
     },
     net::{channel::Channel, P2pPtr},
     system::{msleep, Publisher, PublisherPtr, StoppableTask, StoppableTaskPtr, Subscription},
-    zk::{ProvingKey, VerifyingKey},
     Error, Result,
 };
 
@@ -371,14 +370,8 @@ pub struct EventGraph {
     deg_publisher: PublisherPtr<DegEvent>,
     /// Run in fast mode where if set we sync only headers.
     fast_mode: bool,
-    /// Registering verify key
-    register_vk: VerifyingKey,
-    /// Signaling verify key
-    signal_vk: VerifyingKey,
-    /// Slashing proving key
-    slash_pk: ProvingKey,
-    /// Slashing verify key
-    slash_vk: VerifyingKey,
+    /// sled DB
+    sled_db: sled::Db,
     /// RLN identity storage
     pub rln_identity_tree: RwLock<SmtMemoryFp>,
 }
@@ -405,12 +398,13 @@ impl EventGraph {
         hours_rotation: u64,
         ex: Arc<Executor<'_>>,
     ) -> Result<EventGraphPtr> {
-        let register_vk = build_register_vk(&sled_db)?;
-        let signal_vk = build_signal_vk(&sled_db)?;
-        let slash_pk = build_slash_pk(&sled_db)?;
-        let slash_vk = build_slash_vk(&sled_db)?;
+        let _register_vk = build_register_vk(&sled_db)?;
+        let _signal_vk = build_signal_vk(&sled_db)?;
+        let _slash_pk = build_slash_pk(&sled_db)?;
+        let _slash_vk = build_slash_vk(&sled_db)?;
 
         let hasher = PoseidonFp::new();
+        // let store = AccountStorage::new(&sled_db, "name".to_owned());
         let store = MemoryStorageFp::new();
         let identity_tree = SmtMemoryFp::new(store, hasher.clone(), &EMPTY_NODES_FP);
 
@@ -429,7 +423,7 @@ impl EventGraph {
         .new(sled_db.clone(), hours_rotation)
         .await;
 
-        let static_dag = Self::static_new(sled_db).await?;
+        let static_dag = Self::static_new(&sled_db).await?;
 
         let self_ = Arc::new(Self {
             p2p,
@@ -447,10 +441,7 @@ impl EventGraph {
             synced: RwLock::new(false),
             deg_enabled: RwLock::new(false),
             deg_publisher: Publisher::new(),
-            register_vk,
-            signal_vk,
-            slash_pk,
-            slash_vk,
+            sled_db,
             rln_identity_tree: RwLock::new(identity_tree),
         });
 
@@ -1256,7 +1247,7 @@ impl EventGraph {
         Ok(result)
     }
 
-    pub async fn static_new(sled_db: sled::Db) -> Result<sled::Tree> {
+    pub async fn static_new(sled_db: &sled::Db) -> Result<sled::Tree> {
         let static_dag = sled_db.open_tree("static-dag")?;
 
         let genesis = generate_genesis(0);

+ 39 - 23
src/event_graph/proto.rs

@@ -43,6 +43,7 @@ use super::{
     Event, EventGraphPtr, LayerUTips, NULL_ID, NULL_PARENTS,
 };
 use crate::{
+    event_graph::rln::{read_register_vk, read_signal_vk, read_slash_pk, read_slash_vk, Blob},
     impl_p2p_message,
     net::{
         metering::{MeteringConfiguration, DEFAULT_METERING_CONFIGURATION},
@@ -309,12 +310,7 @@ impl ProtocolEventGraph {
                 if blob.is_empty() {
                     break
                 }
-                let (proof, y, internal_nullifier, user_msg_limit): (
-                    Proof,
-                    pallas::Base,
-                    pallas::Base,
-                    u64,
-                ) = match deserialize_async_partial(&blob).await {
+                let rcvd_blob: Blob = match deserialize_async_partial(&blob).await {
                     Ok((v, _)) => v,
                     Err(e) => {
                         error!(target: "event_graph::protocol::handle_event_put()","[EVENTGRAPH] Failed deserializing event ephemeral data: {}", e);
@@ -333,30 +329,42 @@ impl ProtocolEventGraph {
                 let external_nullifier = poseidon_hash([epoch, rln_app_identifier]);
                 let x = hash_event(&event);
                 let identity_root = self.event_graph.rln_identity_tree.read().await.root();
-                let public_inputs =
-                    vec![identity_root, external_nullifier, x, y, internal_nullifier];
-
-                if metadata.is_duplicate(&external_nullifier, &internal_nullifier, &x, &y) {
+                let public_inputs = vec![
+                    identity_root,
+                    external_nullifier,
+                    x,
+                    rcvd_blob.y,
+                    rcvd_blob.internal_nullifier,
+                ];
+
+                if metadata.is_duplicate(
+                    &external_nullifier,
+                    &rcvd_blob.internal_nullifier,
+                    &x,
+                    &rcvd_blob.y,
+                ) {
                     error!(target: "event_graph::protocol::handle_event_put()", "[RLN] Duplicate Message!");
                     verification_failed = true;
                     break
                 }
 
-                if metadata.is_reused(&external_nullifier, &internal_nullifier) {
+                if metadata.is_reused(&external_nullifier, &rcvd_blob.internal_nullifier) {
                     info!(target: "event_graph::protocol::handle_event_put()", "[RLN] Metadata is reused.. slashing..");
-                    let shares = metadata.get_shares(&external_nullifier, &internal_nullifier);
+                    let shares =
+                        metadata.get_shares(&external_nullifier, &rcvd_blob.internal_nullifier);
                     let secret = sss_recover(&shares);
 
                     // Broadcast slashing event
-                    let slash_pk = &self.event_graph.slash_pk;
+                    let slash_pk = read_slash_pk(&self.event_graph.sled_db)?;
+                    // let slash_pk = &self.event_graph.slash_pk;
                     let mut identity_tree = self.event_graph.rln_identity_tree.write().await;
 
                     info!("[RLN] Creating slashing proof");
                     let (proof, identity_root) = match create_slash_proof(
                         secret,
-                        user_msg_limit,
+                        rcvd_blob.user_msg_limit,
                         &mut identity_tree,
-                        slash_pk,
+                        &slash_pk,
                     ) {
                         Ok(v) => v,
                         Err(e) => {
@@ -368,10 +376,12 @@ impl ProtocolEventGraph {
                     drop(identity_tree);
 
                     let blob =
-                        serialize_async(&(proof, secret, user_msg_limit, identity_root)).await;
+                        serialize_async(&(proof, secret, rcvd_blob.user_msg_limit, identity_root))
+                            .await;
 
                     let evgr = &self.event_graph;
-                    let identity_secret_hash = poseidon_hash([secret, user_msg_limit.into()]);
+                    let identity_secret_hash =
+                        poseidon_hash([secret, rcvd_blob.user_msg_limit.into()]);
                     let identity_commitment = poseidon_hash([identity_secret_hash]);
                     let rln_commitment = RLNNode::Slashing(identity_commitment);
                     let st_event =
@@ -384,11 +394,16 @@ impl ProtocolEventGraph {
                 }
 
                 // At this point we can safely add the shares
-                metadata.add_share(external_nullifier, internal_nullifier, x, y)?;
+                metadata.add_share(
+                    external_nullifier,
+                    rcvd_blob.internal_nullifier,
+                    x,
+                    rcvd_blob.y,
+                )?;
 
                 info!(target: "event_graph::protocol::handle_event_put()", "[RLN] Verifying incoming Event RLN proof");
-                verification_failed =
-                    proof.verify(&self.event_graph.signal_vk, &public_inputs).is_err();
+                let signal_vk = read_signal_vk(&self.event_graph.sled_db)?;
+                verification_failed = rcvd_blob.proof.verify(&signal_vk, &public_inputs).is_err();
 
                 break
             }
@@ -707,7 +722,8 @@ impl ProtocolEventGraph {
                     info!("registering account: {:?}", commitment);
                     let public_inputs = vec![commitment, user_msg_limit.into()];
 
-                    if proof.verify(&self.event_graph.register_vk, &public_inputs).is_err() {
+                    let register_vk = read_register_vk(&self.event_graph.sled_db)?;
+                    if proof.verify(&register_vk, &public_inputs).is_err() {
                         error!(target: "event_graph::protocol::handle_static_put()", "[RLN] Incoming Event RLN Registration proof verification failed");
                         continue
                     }
@@ -728,8 +744,8 @@ impl ProtocolEventGraph {
 
                     let public_inputs =
                         vec![secret, pallas::Base::from(user_msg_limit), identity_root];
-
-                    if proof.verify(&self.event_graph.slash_vk, &public_inputs).is_err() {
+                    let slash_vk = read_slash_vk(&self.event_graph.sled_db)?;
+                    if proof.verify(&slash_vk, &public_inputs).is_err() {
                         error!(target: "event_graph::protocol::handle_static_put()", "[RLN] Incoming Event RLN Slashing proof verification failed");
                         continue
                     }

+ 154 - 60
src/event_graph/rln.rs

@@ -34,7 +34,7 @@ use crate::{
     event_graph::Event,
     zk::{empty_witnesses, Proof, ProvingKey, VerifyingKey, Witness, ZkCircuit},
     zkas::ZkBinary,
-    Result,
+    Error, Result,
 };
 
 pub const RLN2_REGISTER_ZKBIN: &[u8] = include_bytes!("proof/rlnv2-diff-register.zk.bin");
@@ -46,6 +46,64 @@ pub const RLN_GENESIS: u64 = 1_738_688_400_000;
 /// RLN epoch length in millis
 pub const RLN_EPOCH_LEN: u64 = 600_000; // 10 min
 
+#[derive(SerialEncodable, SerialDecodable)]
+pub struct Blob {
+    pub proof: Proof,
+    pub y: pallas::Base,
+    pub internal_nullifier: pallas::Base,
+    pub user_msg_limit: u64,
+}
+
+// pub type SmtAccntFp = SparseMerkleTree<
+//     'static,
+//     SMT_FP_DEPTH,
+//     { SMT_FP_DEPTH + 1 },
+//     pallas::Base,
+//     PoseidonFp,
+//     AccountStorage,
+// >;
+
+// #[derive(Clone)]
+// pub struct AccountStorage {
+//     pub tree: sled::Tree,
+// }
+
+// impl AccountStorage {
+//     pub fn new(sled_db: &sled::Db, name: String) -> Self {
+//         Self { tree: sled_db.open_tree(name).unwrap() }
+//     }
+// }
+
+// impl StorageAdapter for AccountStorage {
+//     type Value = pallas::Base;
+
+//     fn put(&mut self, key: BigUint, value: pallas::Base) -> ContractResult {
+//         self.tree.insert(key.to_bytes_le(), &value.to_repr()).unwrap();
+//         Ok(())
+//     }
+
+//     fn get(&self, key: &BigUint) -> Option<pallas::Base> {
+//         let value = match self.tree.get(&key.to_bytes_le()) {
+//             Ok(v) => v,
+//             Err(e) => {
+//                 error!("SledStorage::get(): Fetching key {:?} from Accounts tree: {}", key, e,);
+//                 return None
+//             }
+//         };
+
+//         let value = value?;
+//         let mut repr = [0; 32];
+//         repr.copy_from_slice(&value);
+
+//         pallas::Base::from_repr(repr).into()
+//     }
+
+//     fn del(&mut self, key: &BigUint) -> ContractResult {
+//         self.tree.remove(key.to_bytes_le()).unwrap();
+//         Ok(())
+//     }
+// }
+
 /// Hash message/event modulo `Fp`
 pub fn hash_event(event: &Event) -> pallas::Base {
     let mut buf = [0u8; 64];
@@ -222,87 +280,123 @@ pub fn sss_recover(shares: &[(pallas::Base, pallas::Base)]) -> pallas::Base {
 }
 
 /// Helper function to read or build register verifying key
-pub(super) fn build_register_vk(sled_db: &sled::Db) -> Result<VerifyingKey> {
+pub(super) fn build_register_vk(sled_db: &sled::Db) -> Result<()> {
+    // sanity check
+    if sled_db.get("rlnv2-diff-register-vk")?.is_some() {
+        return Ok(())
+    }
     let register_zkbin = ZkBinary::decode(RLN2_REGISTER_ZKBIN, false).unwrap();
     let register_empty_circuit =
         ZkCircuit::new(empty_witnesses(&register_zkbin).unwrap(), &register_zkbin);
 
-    match sled_db.get("rlnv2-diff-register-vk")? {
-        Some(vk) => {
-            let mut reader = Cursor::new(vk);
-            Ok(VerifyingKey::read(&mut reader, register_empty_circuit)?)
-        }
-        None => {
-            info!(target: "irc::server", "[RLN] Creating RlnV2_Diff_Register VerifyingKey");
-            let verifyingkey = VerifyingKey::build(register_zkbin.k, &register_empty_circuit);
-            let mut buf = vec![];
-            verifyingkey.write(&mut buf)?;
-            sled_db.insert("rlnv2-diff-register-vk", buf)?;
-            Ok(verifyingkey)
-        }
+    info!(target: "irc::server", "[RLN] Creating RlnV2_Diff_Register VerifyingKey");
+    let verifyingkey = VerifyingKey::build(register_zkbin.k, &register_empty_circuit);
+    let mut buf = vec![];
+    verifyingkey.write(&mut buf)?;
+    sled_db.insert("rlnv2-diff-register-vk", buf)?;
+    Ok(())
+}
+
+/// Helper function to read register verifying key
+pub(super) fn read_register_vk(sled_db: &sled::Db) -> Result<VerifyingKey> {
+    if let Some(vk) = sled_db.get("rlnv2-diff-register-vk")? {
+        let register_zkbin = ZkBinary::decode(RLN2_REGISTER_ZKBIN, false).unwrap();
+        let register_empty_circuit =
+            ZkCircuit::new(empty_witnesses(&register_zkbin).unwrap(), &register_zkbin);
+        let mut reader = Cursor::new(vk);
+        Ok(VerifyingKey::read(&mut reader, register_empty_circuit)?)
+    } else {
+        Err(Error::Custom("Error reading register verifying key".to_owned()))
     }
 }
 
-/// Helper function to read or build signal verifying key
-pub(super) fn build_signal_vk(sled_db: &sled::Db) -> Result<VerifyingKey> {
+/// Helper function to build signal verifying key
+pub(super) fn build_signal_vk(sled_db: &sled::Db) -> Result<()> {
+    // sanity check
+    if sled_db.get("rlnv2-diff-signal-vk")?.is_some() {
+        return Ok(())
+    }
     let signal_zkbin = ZkBinary::decode(RLN2_SIGNAL_ZKBIN, false).unwrap();
     let signal_empty_circuit =
         ZkCircuit::new(empty_witnesses(&signal_zkbin).unwrap(), &signal_zkbin);
 
-    match sled_db.get("rlnv2-diff-signal-vk")? {
-        Some(vk) => {
-            let mut reader = Cursor::new(vk);
-            Ok(VerifyingKey::read(&mut reader, signal_empty_circuit)?)
-        }
-        None => {
-            info!(target: "irc::server", "[RLN] Creating RlnV2_Diff_Signal VerifyingKey");
-            let verifyingkey = VerifyingKey::build(signal_zkbin.k, &signal_empty_circuit);
-            let mut buf = vec![];
-            verifyingkey.write(&mut buf)?;
-            sled_db.insert("rlnv2-diff-signal-vk", buf)?;
-            Ok(verifyingkey)
-        }
+    info!(target: "irc::server", "[RLN] Creating RlnV2_Diff_Signal VerifyingKey");
+    let verifyingkey = VerifyingKey::build(signal_zkbin.k, &signal_empty_circuit);
+    let mut buf = vec![];
+    verifyingkey.write(&mut buf)?;
+    sled_db.insert("rlnv2-diff-signal-vk", buf)?;
+    Ok(())
+}
+
+/// Helper function to read signal verifying key
+pub(super) fn read_signal_vk(sled_db: &sled::Db) -> Result<VerifyingKey> {
+    if let Some(vk) = sled_db.get("rlnv2-diff-signal-vk")? {
+        let signal_zkbin = ZkBinary::decode(RLN2_SIGNAL_ZKBIN, false).unwrap();
+        let signal_empty_circuit =
+            ZkCircuit::new(empty_witnesses(&signal_zkbin).unwrap(), &signal_zkbin);
+        let mut reader = Cursor::new(vk);
+        Ok(VerifyingKey::read(&mut reader, signal_empty_circuit)?)
+    } else {
+        Err(Error::Custom("Error Reading signal verifying key".to_owned()))
     }
 }
 
-/// Helper function to read or build slash proving key
-pub(super) fn build_slash_pk(sled_db: &sled::Db) -> Result<ProvingKey> {
+/// Helper function to build slash proving key
+pub(super) fn build_slash_pk(sled_db: &sled::Db) -> Result<()> {
+    // sanity check
+    if sled_db.get("rlnv2-diff-slash-pk")?.is_some() {
+        return Ok(())
+    }
     let slash_zkbin = ZkBinary::decode(RLN2_SLASH_ZKBIN, false).unwrap();
     let slash_empty_circuit = ZkCircuit::new(empty_witnesses(&slash_zkbin).unwrap(), &slash_zkbin);
 
-    match sled_db.get("rlnv2-diff-slash-pk")? {
-        Some(pk) => {
-            let mut reader = Cursor::new(pk);
-            Ok(ProvingKey::read(&mut reader, slash_empty_circuit)?)
-        }
-        None => {
-            info!(target: "irc::server", "[RLN] Creating RlnV2_Diff_Slash ProvingKey");
-            let provingkey = ProvingKey::build(slash_zkbin.k, &slash_empty_circuit);
-            let mut buf = vec![];
-            provingkey.write(&mut buf)?;
-            sled_db.insert("rlnv2-diff-slash-pk", buf)?;
-            Ok(provingkey)
-        }
+    info!(target: "irc::server", "[RLN] Creating RlnV2_Diff_Slash ProvingKey");
+    let verifyingkey = VerifyingKey::build(slash_zkbin.k, &slash_empty_circuit);
+    let mut buf = vec![];
+    verifyingkey.write(&mut buf)?;
+    sled_db.insert("rlnv2-diff-slash-pk", buf)?;
+    Ok(())
+}
+
+/// Helper function to read slash proving key
+pub(super) fn read_slash_pk(sled_db: &sled::Db) -> Result<ProvingKey> {
+    if let Some(vk) = sled_db.get("rlnv2-diff-slash-pk")? {
+        let slash_zkbin = ZkBinary::decode(RLN2_SLASH_ZKBIN, false).unwrap();
+        let slash_empty_circuit =
+            ZkCircuit::new(empty_witnesses(&slash_zkbin).unwrap(), &slash_zkbin);
+        let mut reader = Cursor::new(vk);
+        Ok(ProvingKey::read(&mut reader, slash_empty_circuit)?)
+    } else {
+        Err(Error::Custom("Error Reading slash proving key".to_owned()))
     }
 }
 
-/// Helper function to read or build slash verifying key
-pub(super) fn build_slash_vk(sled_db: &sled::Db) -> Result<VerifyingKey> {
+/// Helper function to build slash verifying key
+pub(super) fn build_slash_vk(sled_db: &sled::Db) -> Result<()> {
+    // sanity check
+    if sled_db.get("rlnv2-diff-slash-vk")?.is_some() {
+        return Ok(())
+    }
     let slash_zkbin = ZkBinary::decode(RLN2_SLASH_ZKBIN, false).unwrap();
     let slash_empty_circuit = ZkCircuit::new(empty_witnesses(&slash_zkbin).unwrap(), &slash_zkbin);
 
-    match sled_db.get("rlnv2-diff-slash-vk")? {
-        Some(vk) => {
-            let mut reader = Cursor::new(vk);
-            Ok(VerifyingKey::read(&mut reader, slash_empty_circuit)?)
-        }
-        None => {
-            info!(target: "irc::server", "[RLN] Creating RlnV2_Diff_Slash VerifyingKey");
-            let verifyingkey = VerifyingKey::build(slash_zkbin.k, &slash_empty_circuit);
-            let mut buf = vec![];
-            verifyingkey.write(&mut buf)?;
-            sled_db.insert("rlnv2-diff-slash-vk", buf)?;
-            Ok(verifyingkey)
-        }
+    info!(target: "irc::server", "[RLN] Creating RlnV2_Diff_Slash VerifyingKey");
+    let verifyingkey = VerifyingKey::build(slash_zkbin.k, &slash_empty_circuit);
+    let mut buf = vec![];
+    verifyingkey.write(&mut buf)?;
+    sled_db.insert("rlnv2-diff-slash-vk", buf)?;
+    Ok(())
+}
+
+/// Helper function to read slash proving key
+pub(super) fn read_slash_vk(sled_db: &sled::Db) -> Result<VerifyingKey> {
+    if let Some(vk) = sled_db.get("rlnv2-diff-slash-pk")? {
+        let slash_zkbin = ZkBinary::decode(RLN2_SLASH_ZKBIN, false).unwrap();
+        let slash_empty_circuit =
+            ZkCircuit::new(empty_witnesses(&slash_zkbin).unwrap(), &slash_zkbin);
+        let mut reader = Cursor::new(vk);
+        Ok(VerifyingKey::read(&mut reader, slash_empty_circuit)?)
+    } else {
+        Err(Error::Custom("Error Reading slash verifying key".to_owned()))
     }
 }