feat(prover): adding keystore object to handle reading and writing of…

… prover keys (matter-labs#1132)

## What ❔

* Adding keystore object as one interface to handle reading and writing
of the verification and setup keys
* As a side effect, also removing dependency on the old test_harness
version from key generator

## Why ❔

* This will allow us to specify different (non default) directories for
verification and setup keys in the future.

prover/proof_fri_compressor/Cargo.toml

@@ -33,3 +33,6 @@ ctrlc = { version = "3.1", features = ["termination"] }
 async-trait = "0.1"
 bincode = "1.0"
 reqwest = { version = "0.11", features = ["blocking"] }
+serde_json = "1.0"
+serde = { version = "1.0", features = ["derive"] }
+

prover/proof_fri_compressor/src/compressor.rs

@@ -8,7 +8,10 @@ use zkevm_test_harness_1_3_3::{
     abstract_zksync_circuit::concrete_circuits::{
         ZkSyncCircuit, ZkSyncProof, ZkSyncVerificationKey,
     },
-    bellman::{bn256::Bn256, plonk::better_better_cs::proof::Proof},
+    bellman::{
+        bn256::Bn256,
+        plonk::better_better_cs::{proof::Proof, setup::VerificationKey as SnarkVerificationKey},
+    },
     witness::oracle::VmWitnessOracle,
 };
 use zksync_dal::ConnectionPool;
@@ -26,7 +29,7 @@ use zksync_prover_fri_types::{
 use zksync_prover_interface::outputs::L1BatchProofForL1;
 use zksync_queued_job_processor::JobProcessor;
 use zksync_types::L1BatchNumber;
-use zksync_vk_setup_data_server_fri::{get_recursive_layer_vk_for_circuit_type, get_snark_vk};
+use zksync_vk_setup_data_server_fri::keystore::Keystore;
 
 use crate::metrics::METRICS;
 
@@ -60,10 +63,12 @@ impl ProofCompressor {
         compression_mode: u8,
         verify_wrapper_proof: bool,
     ) -> anyhow::Result<Proof<Bn256, ZkSyncCircuit<Bn256, VmWitnessOracle<Bn256>>>> {
-        let scheduler_vk = get_recursive_layer_vk_for_circuit_type(
-            ZkSyncRecursionLayerStorageType::SchedulerCircuit as u8,
-        )
-        .context("get_recursiver_layer_vk_for_circuit_type()")?;
+        let keystore = Keystore::default();
+        let scheduler_vk = keystore
+            .load_recursive_layer_verification_key(
+                ZkSyncRecursionLayerStorageType::SchedulerCircuit as u8,
+            )
+            .context("get_recursiver_layer_vk_for_circuit_type()")?;
         let config = WrapperConfig::new(compression_mode);
 
         let (wrapper_proof, _) = wrap_proof(proof, scheduler_vk, config);
@@ -76,7 +81,15 @@ impl ProofCompressor {
             bincode::deserialize(&serialized)
                 .expect("Failed to deserialize proof with ZkSyncCircuit");
         if verify_wrapper_proof {
-            let existing_vk = get_snark_vk().context("get_snark_vk()")?;
+            // We're fetching the key as String and deserializing it here
+            // as we don't want to include the old version of prover in the main libraries.
+            let existing_vk_serialized = keystore
+                .load_snark_verification_key()
+                .context("get_snark_vk()")?;
+            let existing_vk = serde_json::from_str::<
+                SnarkVerificationKey<Bn256, ZkSyncCircuit<Bn256, VmWitnessOracle<Bn256>>>,
+            >(&existing_vk_serialized)?;
+
             let vk = ZkSyncVerificationKey::from_verification_key_and_numeric_type(0, existing_vk);
             let scheduler_proof = ZkSyncProof::from_proof_and_numeric_type(0, proof.clone());
             match vk.verify_proof(&scheduler_proof) {

prover/prover_fri/src/gpu_prover_job_processor.rs

@@ -27,9 +27,7 @@ pub mod gpu_prover {
     };
     use zksync_queued_job_processor::{async_trait, JobProcessor};
     use zksync_types::basic_fri_types::CircuitIdRoundTuple;
-    use zksync_vk_setup_data_server_fri::{
-        get_setup_data_for_circuit_type, GoldilocksGpuProverSetupData,
-    };
+    use zksync_vk_setup_data_server_fri::{keystore::Keystore, GoldilocksGpuProverSetupData};
 
     use crate::{
         metrics::METRICS,
@@ -103,9 +101,10 @@ pub mod gpu_prover {
                     .clone(),
                 SetupLoadMode::FromDisk => {
                     let started_at = Instant::now();
-                    let artifact: GoldilocksGpuProverSetupData =
-                        get_setup_data_for_circuit_type(key.clone())
-                            .context("get_setup_data_for_circuit_type()")?;
+                    let keystore = Keystore::default();
+                    let artifact: GoldilocksGpuProverSetupData = keystore
+                        .load_gpu_setup_data_for_circuit_type(key.clone())
+                        .context("load_gpu_setup_data_for_circuit_type()")?;
 
                     METRICS.gpu_setup_data_load_time[&key.circuit_id.to_string()]
                         .observe(started_at.elapsed());
@@ -326,10 +325,12 @@ pub mod gpu_prover {
                     &config.specialized_group_id,
                     prover_setup_metadata_list
                 );
+                let keystore = Keystore::default();
                 for prover_setup_metadata in prover_setup_metadata_list {
                     let key = setup_metadata_to_setup_data_key(&prover_setup_metadata);
-                    let setup_data = get_setup_data_for_circuit_type(key.clone())
-                        .context("get_setup_data_for_circuit_type()")?;
+                    let setup_data = keystore
+                        .load_gpu_setup_data_for_circuit_type(key.clone())
+                        .context("load_gpu_setup_data_for_circuit_type()")?;
                     cache.insert(key, Arc::new(setup_data));
                 }
                 SetupLoadMode::FromMemory(cache)

prover/prover_fri/src/main.rs

@@ -208,8 +208,6 @@ async fn get_prover_tasks(
     pool: ConnectionPool,
     circuit_ids_for_round_to_be_proven: Vec<CircuitIdRoundTuple>,
 ) -> anyhow::Result<Vec<JoinHandle<anyhow::Result<()>>>> {
-    use std::sync::Arc;
-
     use gpu_prover_job_processor::gpu_prover;
     use socket_listener::gpu_socket_listener;
     use tokio::sync::Mutex;

prover/prover_fri/src/prover_job_processor.rs

@@ -25,9 +25,7 @@ use zksync_prover_fri_types::{
 use zksync_prover_fri_utils::fetch_next_circuit;
 use zksync_queued_job_processor::{async_trait, JobProcessor};
 use zksync_types::{basic_fri_types::CircuitIdRoundTuple, protocol_version::L1VerifierConfig};
-use zksync_vk_setup_data_server_fri::{
-    get_cpu_setup_data_for_circuit_type, GoldilocksProverSetupData,
-};
+use zksync_vk_setup_data_server_fri::{keystore::Keystore, GoldilocksProverSetupData};
 
 use crate::{
     metrics::{CircuitLabels, Layer, METRICS},
@@ -88,9 +86,10 @@ impl Prover {
                 .clone(),
             SetupLoadMode::FromDisk => {
                 let started_at = Instant::now();
-                let artifact: GoldilocksProverSetupData =
-                    get_cpu_setup_data_for_circuit_type(key.clone())
-                        .context("get_cpu_setup_data_for_circuit_type()")?;
+                let keystore = Keystore::default();
+                let artifact: GoldilocksProverSetupData = keystore
+                    .load_cpu_setup_data_for_circuit_type(key.clone())
+                    .context("get_cpu_setup_data_for_circuit_type()")?;
                 METRICS.gpu_setup_data_load_time[&key.circuit_id.to_string()]
                     .observe(started_at.elapsed());
 
@@ -299,9 +298,11 @@ pub fn load_setup_data_cache(config: &FriProverConfig) -> anyhow::Result<SetupLo
                 &config.specialized_group_id,
                 prover_setup_metadata_list
             );
+            let keystore = Keystore::default();
             for prover_setup_metadata in prover_setup_metadata_list {
                 let key = setup_metadata_to_setup_data_key(&prover_setup_metadata);
-                let setup_data = get_cpu_setup_data_for_circuit_type(key.clone())
+                let setup_data = keystore
+                    .load_cpu_setup_data_for_circuit_type(key.clone())
                     .context("get_cpu_setup_data_for_circuit_type()")?;
                 cache.insert(key, Arc::new(setup_data));
             }

prover/prover_fri/src/socket_listener.rs

@@ -18,9 +18,7 @@ pub mod gpu_socket_listener {
     use zksync_object_store::bincode;
     use zksync_prover_fri_types::{CircuitWrapper, ProverServiceDataKey, WitnessVectorArtifacts};
     use zksync_types::basic_fri_types::AggregationRound;
-    use zksync_vk_setup_data_server_fri::{
-        get_finalization_hints, get_round_for_recursive_circuit_type,
-    };
+    use zksync_vk_setup_data_server_fri::keystore::Keystore;
 
     use crate::{
         metrics::METRICS,
@@ -178,23 +176,24 @@ pub mod gpu_socket_listener {
         circuit_id: u8,
     ) -> anyhow::Result<ProvingAssembly> {
         let started_at = Instant::now();
+        let keystore = Keystore::default();
         let cs = match circuit_wrapper {
             CircuitWrapper::Base(base_circuit) => {
                 let key = ProverServiceDataKey::new(
                     base_circuit.numeric_circuit_type(),
                     AggregationRound::BasicCircuits,
                 );
-                let finalization_hint =
-                    get_finalization_hints(key).context("get_finalization_hints()")?;
+                let finalization_hint = keystore
+                    .load_finalization_hints(key)
+                    .context("get_finalization_hints()")?;
                 init_base_layer_cs_for_repeated_proving(base_circuit, &finalization_hint)
             }
             CircuitWrapper::Recursive(recursive_circuit) => {
-                let key = ProverServiceDataKey::new(
-                    recursive_circuit.numeric_circuit_type(),
-                    get_round_for_recursive_circuit_type(recursive_circuit.numeric_circuit_type()),
-                );
-                let finalization_hint =
-                    get_finalization_hints(key).context("get_finalization_hints()")?;
+                let key =
+                    ProverServiceDataKey::new_recursive(recursive_circuit.numeric_circuit_type());
+                let finalization_hint = keystore
+                    .load_finalization_hints(key)
+                    .context("get_finalization_hints()")?;
                 init_recursive_layer_cs_for_repeated_proving(recursive_circuit, &finalization_hint)
             }
         };

prover/prover_fri/tests/basic_test.rs

@@ -10,7 +10,9 @@ use zksync_prover_fri_types::{
     keys::FriCircuitKey, CircuitWrapper, ProverJob, ProverServiceDataKey,
 };
 use zksync_types::{basic_fri_types::AggregationRound, L1BatchNumber};
-use zksync_vk_setup_data_server_fri::generate_cpu_base_layer_setup_data;
+use zksync_vk_setup_data_server_fri::{
+    keystore::Keystore, setup_data_generator::generate_cpu_base_layer_setup_data,
+};
 
 fn compare_serialized<T: Serialize>(expected: &T, actual: &T) {
     let serialized_expected = bincode::serialize(expected).unwrap();
@@ -54,8 +56,9 @@ async fn prover_and_assert_base_layer(
         CircuitWrapper::Base(base) => base.clone(),
         CircuitWrapper::Recursive(_) => anyhow::bail!("Expected base layer circuit"),
     };
+    let keystore = Keystore::default();
     let setup_data = Arc::new(
-        generate_cpu_base_layer_setup_data(circuit)
+        generate_cpu_base_layer_setup_data(&keystore, circuit)
             .context("generate_cpu_base_layers_setup_data()")?,
     );
     let setup_key = ProverServiceDataKey::new(circuit_id, aggregation_round);

prover/prover_fri_types/src/lib.rs

@@ -6,7 +6,9 @@ use circuit_definitions::{
     boojum::{cs::implementations::witness::WitnessVec, field::goldilocks::GoldilocksField},
     circuit_definitions::{
         base_layer::{ZkSyncBaseLayerCircuit, ZkSyncBaseLayerProof},
-        recursion_layer::{ZkSyncRecursionLayerProof, ZkSyncRecursiveLayerCircuit},
+        recursion_layer::{
+            ZkSyncRecursionLayerProof, ZkSyncRecursionLayerStorageType, ZkSyncRecursiveLayerCircuit,
+        },
     },
     zkevm_circuits::scheduler::block_header::BlockAuxilaryOutputWitness,
     ZkSyncDefaultRoundFunction,
@@ -112,10 +114,38 @@ pub struct ProverServiceDataKey {
     pub round: AggregationRound,
 }
 
+fn get_round_for_recursive_circuit_type(circuit_type: u8) -> AggregationRound {
+    match circuit_type {
+        circuit_type if circuit_type == ZkSyncRecursionLayerStorageType::SchedulerCircuit as u8 => {
+            AggregationRound::Scheduler
+        }
+        circuit_type if circuit_type == ZkSyncRecursionLayerStorageType::NodeLayerCircuit as u8 => {
+            AggregationRound::NodeAggregation
+        }
+        _ => AggregationRound::LeafAggregation,
+    }
+}
+
 impl ProverServiceDataKey {
     pub fn new(circuit_id: u8, round: AggregationRound) -> Self {
         Self { circuit_id, round }
     }
+
+    /// Creates a new data key for recursive type - with auto selection of the aggregation round.
+    pub fn new_recursive(circuit_id: u8) -> Self {
+        Self {
+            circuit_id,
+            round: get_round_for_recursive_circuit_type(circuit_id),
+        }
+    }
+
+    /// Data key for snark wrapper.
+    pub fn snark() -> Self {
+        Self {
+            circuit_id: 1,
+            round: AggregationRound::Scheduler,
+        }
+    }
 }
 
 #[derive(serde::Serialize, serde::Deserialize)]

prover/vk_setup_data_generator_server_fri/Cargo.toml

@@ -18,7 +18,6 @@ vlog = { path = "../../core/lib/vlog" }
 zksync_types = { path = "../../core/lib/types" }
 zksync_prover_fri_types = { path = "../prover_fri_types" }
 
-zkevm_test_harness_1_3_3 = { git = "https://github.com/matter-labs/era-zkevm_test_harness.git", branch = "v1.3.3", package = "zkevm_test_harness" }
 zkevm_test_harness = { git = "https://github.com/matter-labs/era-zkevm_test_harness.git", branch = "v1.4.1" }
 circuit_definitions = { git = "https://github.com/matter-labs/era-zkevm_test_harness.git", branch = "v1.4.1", features = [
     "log_tracing",

prover/vk_setup_data_generator_server_fri/src/commitment_generator.rs

@@ -1,12 +1,13 @@
 use anyhow::Context;
 use zksync_vk_setup_data_server_fri::{
     commitment_utils::generate_commitments,
+    keystore::Keystore,
     vk_commitment_helper::{get_toml_formatted_value, read_contract_toml, write_contract_toml},
 };
 
-pub fn read_and_update_contract_toml(dryrun: bool) -> anyhow::Result<()> {
+pub fn read_and_update_contract_toml(keystore: &Keystore, dryrun: bool) -> anyhow::Result<()> {
     let mut contract_doc = read_contract_toml().context("read_contract_toml()")?;
-    let vk_commitments = generate_commitments().context("generate_commitments()")?;
+    let vk_commitments = generate_commitments(keystore).context("generate_commitments()")?;
 
     contract_doc["contracts"]["FRI_RECURSION_LEAF_LEVEL_VK_HASH"] =
         get_toml_formatted_value(vk_commitments.leaf);
@@ -29,6 +30,6 @@ mod test {
 
     #[test]
     fn test_read_and_update_contract_toml() {
-        read_and_update_contract_toml(true).unwrap();
+        read_and_update_contract_toml(&Keystore::default(), true).unwrap();
     }
 }

prover/vk_setup_data_generator_server_fri/src/commitment_utils.rs

@@ -15,11 +15,11 @@ use zksync_types::{
     H256,
 };
 
-use crate::{get_recursive_layer_vk_for_circuit_type, utils::get_leaf_vk_params};
+use crate::{keystore::Keystore, utils::get_leaf_vk_params};
 
 lazy_static! {
     // TODO: do not initialize a static const with data read in runtime.
-    static ref COMMITMENTS: Lazy<L1VerifierConfig> = Lazy::new(|| { circuit_commitments().unwrap() });
+    static ref COMMITMENTS: Lazy<L1VerifierConfig> = Lazy::new(|| { circuit_commitments(&Keystore::default()).unwrap() });
 }
 
 #[derive(Debug)]
@@ -29,8 +29,8 @@ pub struct VkCommitments {
     pub scheduler: String,
 }
 
-fn circuit_commitments() -> anyhow::Result<L1VerifierConfig> {
-    let commitments = generate_commitments().context("generate_commitments()")?;
+fn circuit_commitments(keystore: &Keystore) -> anyhow::Result<L1VerifierConfig> {
+    let commitments = generate_commitments(keystore).context("generate_commitments()")?;
     let snark_wrapper_vk = std::env::var("CONTRACTS_SNARK_WRAPPER_VK_HASH")
         .context("SNARK wrapper VK not found in the config")?;
     Ok(L1VerifierConfig {
@@ -55,8 +55,8 @@ fn circuit_commitments() -> anyhow::Result<L1VerifierConfig> {
     })
 }
 
-pub fn generate_commitments() -> anyhow::Result<VkCommitments> {
-    let leaf_vk_params = get_leaf_vk_params().context("get_leaf_vk_params()")?;
+pub fn generate_commitments(keystore: &Keystore) -> anyhow::Result<VkCommitments> {
+    let leaf_vk_params = get_leaf_vk_params(keystore).context("get_leaf_vk_params()")?;
     let leaf_layer_params = leaf_vk_params
         .iter()
         .map(|el| el.1.clone())
@@ -65,16 +65,18 @@ pub fn generate_commitments() -> anyhow::Result<VkCommitments> {
         .unwrap();
     let leaf_vk_commitment = compute_leaf_vks_and_params_commitment(leaf_layer_params);
 
-    let node_vk = get_recursive_layer_vk_for_circuit_type(
-        ZkSyncRecursionLayerStorageType::NodeLayerCircuit as u8,
-    )
-    .context("get_recursive_layer_vk_for_circuit_type(NodeLayerCircuit)")?;
+    let node_vk = keystore
+        .load_recursive_layer_verification_key(
+            ZkSyncRecursionLayerStorageType::NodeLayerCircuit as u8,
+        )
+        .context("get_recursive_layer_vk_for_circuit_type(NodeLayerCircuit)")?;
     let node_vk_commitment = compute_node_vk_commitment(node_vk.clone());
 
-    let scheduler_vk = get_recursive_layer_vk_for_circuit_type(
-        ZkSyncRecursionLayerStorageType::SchedulerCircuit as u8,
-    )
-    .context("get_recursive_layer_vk_for_circuit_type(SchedulerCircuit)")?;
+    let scheduler_vk = keystore
+        .load_recursive_layer_verification_key(
+            ZkSyncRecursionLayerStorageType::SchedulerCircuit as u8,
+        )
+        .context("get_recursive_layer_vk_for_circuit_type(SchedulerCircuit)")?;
     let scheduler_vk_commitment = compute_node_vk_commitment(scheduler_vk.clone());
 
     let hex_concatenator = |hex_array: [GoldilocksField; 4]| {