From 8f7affff3794eb8f9f046e5ca26d67a36fddd7ad Mon Sep 17 00:00:00 2001
From: parazyd <parazyd@dyne.org>
Date: Mon, 6 Mar 2023 15:34:57 +0100
Subject: [PATCH] blockchain: Implement abstraction for fetching zkas bincode
 and vks.

---
 src/blockchain/contract_store.rs | 38 ++++++++++++++++++++++++++++-
 src/consensus/validator.rs       | 41 ++++----------------------------
 2 files changed, 42 insertions(+), 37 deletions(-)

diff --git a/src/blockchain/contract_store.rs b/src/blockchain/contract_store.rs
index da4c5007a..4360f27f9 100644
--- a/src/blockchain/contract_store.rs
+++ b/src/blockchain/contract_store.rs
@@ -15,12 +15,18 @@ r* This program is distributed in the hope that it will be useful,
  * You should have received a copy of the GNU Affero General Public License
  * along with this program.  If not, see <https://www.gnu.org/licenses/>.
  */
+use std::io::Cursor;
 
 use darkfi_sdk::crypto::ContractId;
 use darkfi_serial::{deserialize, serialize};
 use log::{debug, error};
 
-use crate::{Error, Result};
+use crate::{
+    runtime::vm_runtime::SMART_CONTRACT_ZKAS_DB_NAME,
+    zk::{VerifyingKey, ZkCircuit},
+    zkas::ZkBinary,
+    Error, Result,
+};
 
 const SLED_CONTRACTS_TREE: &[u8] = b"_contracts";
 const SLED_BINCODE_TREE: &[u8] = b"_wasm_bincode";
@@ -204,4 +210,34 @@ impl ContractStateStore {
 
         Ok(())
     }
+
+    /// Abstraction function for fetching a `ZkBinary` and its respective `VerifyingKey`
+    /// from a contract's zkas sled tree.
+    pub fn get_zkas(
+        &self,
+        db: &sled::Db,
+        contract_id: &ContractId,
+        zkas_ns: &str,
+    ) -> Result<(ZkBinary, VerifyingKey)> {
+        debug!(target: "blockchain::contractstore", "Looking up \"{}:{}\" zkas circuit & vk", contract_id, zkas_ns);
+
+        let zkas_tree = self.lookup(db, contract_id, SMART_CONTRACT_ZKAS_DB_NAME)?;
+
+        let Some(zkas_bytes) = zkas_tree.get(&serialize(&zkas_ns))? else {
+            return Err(Error::ZkasBincodeNotFound)
+        };
+
+        // If anything in this function panics, that means corrupted data managed
+        // to get into this sled tree. This should not be possible.
+        let (zkbin, vkbin): (Vec<u8>, Vec<u8>) = deserialize(&zkas_bytes).unwrap();
+
+        // The first vec is the compiled zkas binary
+        let zkbin = ZkBinary::decode(&zkbin).unwrap();
+
+        // The second one is the serialized VerifyingKey for it
+        let mut vk_buf = Cursor::new(vkbin);
+        let vk = VerifyingKey::read::<Cursor<Vec<u8>>, ZkCircuit>(&mut vk_buf).unwrap();
+
+        Ok((zkbin, vk))
+    }
 }
diff --git a/src/consensus/validator.rs b/src/consensus/validator.rs
index 189242893..75736c63b 100644
--- a/src/consensus/validator.rs
+++ b/src/consensus/validator.rs
@@ -29,7 +29,7 @@ use darkfi_sdk::{
     incrementalmerkletree::{bridgetree::BridgeTree, Tree},
     pasta::{group::ff::PrimeField, pallas},
 };
-use darkfi_serial::{deserialize, serialize, Decodable, Encodable, WriteExt};
+use darkfi_serial::{serialize, Decodable, Encodable, WriteExt};
 use halo2_proofs::arithmetic::Field;
 use log::{debug, error, info, warn};
 use rand::rngs::OsRng;
@@ -45,7 +45,7 @@ use super::{
 use crate::{
     blockchain::Blockchain,
     rpc::jsonrpc::JsonNotification,
-    runtime::vm_runtime::{Runtime, SMART_CONTRACT_ZKAS_DB_NAME},
+    runtime::vm_runtime::Runtime,
     system::{Subscriber, SubscriberPtr},
     tx::Transaction,
     util::time::Timestamp,
@@ -959,21 +959,7 @@ impl ValidatorState {
                 info!(target: "consensus::validator", "Successfully executed \"metadata\" call");
 
                 // Here we'll look up verifying keys and insert them into the per-contract map.
-                // TODO: This should be abstracted
                 info!(target: "consensus::validator", "Performing VerifyingKey lookups from the sled db");
-                let zkas_tree = match self.blockchain.contracts.lookup(
-                    &self.blockchain.sled_db,
-                    &call.contract_id,
-                    SMART_CONTRACT_ZKAS_DB_NAME,
-                ) {
-                    Ok(v) => v,
-                    Err(e) => {
-                        error!(target: "consensus::validator", "Failed to lookup zkas db for contract {}: {}", call.contract_id, e);
-                        skip = true;
-                        break
-                    }
-                };
-
                 for (zkas_ns, _) in &zkp_pub {
                     let inner_vk_map =
                         verifying_keys.get_mut(&call.contract_id.to_bytes()).unwrap();
@@ -982,31 +968,14 @@ impl ValidatorState {
                         continue
                     }
 
-                    let Some(zkas_encoded_bytes) = zkas_tree.get(&serialize(&zkas_ns.as_str())).expect("get") else {
+                    let Ok((_, vk)) = self.blockchain.contracts.get_zkas(
+                        &self.blockchain.sled_db, &call.contract_id, &zkas_ns
+                    ) else {
                         error!(target: "consensus::validator", "Failed to find reference to zkas in sled");
                         skip = true;
                         break
                     };
 
-                    let (_, vk_bin): (Vec<u8>, Vec<u8>) = match deserialize(&zkas_encoded_bytes) {
-                        Ok(v) => v,
-                        Err(e) => {
-                            error!(target: "consensus::validator", "Failed to deserialize zkas data: {}", e);
-                            skip = true;
-                            break
-                        }
-                    };
-
-                    let mut vk_buf = Cursor::new(vk_bin);
-                    let vk = match VerifyingKey::read::<Cursor<Vec<u8>>, ZkCircuit>(&mut vk_buf) {
-                        Ok(v) => v,
-                        Err(e) => {
-                            error!(target: "consensus::validator", "Failed to decode VerifyingKey: {}", e);
-                            skip = true;
-                            break
-                        }
-                    };
-
                     inner_vk_map.insert(zkas_ns.to_string(), vk);
                 }