Enable verifying proofs based on the fork (#1011)

* Partially rewrite fork manager * Add SpecId to the zk output * Lint * Verify proofs with the corresponding spec ids * Disable lint * Remove unused const * Remove std * Lint * Remove Receipt output type * Remove risc0 zkvm from mockvm * Lint * Rename to last_active_spec_id * Restore fork logs * Lint * Remove unused * Remove unnecessary allow clippy --------- Co-authored-by: yaziciahmet <yaziciahmetcleargmail.com>
chainwayxyz · Aug 20, 2024 · 1bacc1e · 1bacc1e
1 parent 5dea363
commit 1bacc1e
Show file tree

Hide file tree

Showing 31 changed files with 338 additions and 313 deletions.
diff --git a/bin/citrea/src/rollup/bitcoin.rs b/bin/citrea/src/rollup/bitcoin.rs
@@ -1,3 +1,4 @@
+use std::collections::HashMap;
 use std::sync::Arc;
 
 use async_trait::async_trait;
@@ -18,6 +19,7 @@ use sov_modules_stf_blueprint::StfBlueprint;
 use sov_prover_storage_manager::ProverStorageManager;
 use sov_rollup_interface::da::DaVerifier;
 use sov_rollup_interface::services::da::SenderWithNotifier;
+use sov_rollup_interface::spec::SpecId;
 use sov_rollup_interface::zk::{Zkvm, ZkvmHost};
 use sov_state::{DefaultStorageSpec, Storage, ZkStorage};
 use sov_stf_runner::{FullNodeConfig, ProverConfig};
@@ -90,8 +92,10 @@ impl RollupBlueprint for BitcoinRollup {
     }
 
     #[instrument(level = "trace", skip(self), ret)]
-    fn get_code_commitment(&self) -> <Self::Vm as Zkvm>::CodeCommitment {
-        Digest::new(citrea_risc0::BITCOIN_DA_ID)
+    fn get_code_commitments_by_spec(&self) -> HashMap<SpecId, <Self::Vm as Zkvm>::CodeCommitment> {
+        let mut map = HashMap::new();
+        map.insert(SpecId::Genesis, Digest::new(citrea_risc0::BITCOIN_DA_ID));
+        map
     }
 
     #[instrument(level = "trace", skip_all, err)]

diff --git a/bin/citrea/src/rollup/mock.rs b/bin/citrea/src/rollup/mock.rs
@@ -1,3 +1,4 @@
+use std::collections::HashMap;
 use std::sync::Arc;
 
 use async_trait::async_trait;
@@ -13,6 +14,7 @@ use sov_modules_api::{Address, Spec};
 use sov_modules_rollup_blueprint::RollupBlueprint;
 use sov_modules_stf_blueprint::StfBlueprint;
 use sov_prover_storage_manager::ProverStorageManager;
+use sov_rollup_interface::spec::SpecId;
 use sov_rollup_interface::zk::{Zkvm, ZkvmHost};
 use sov_state::{DefaultStorageSpec, Storage, ZkStorage};
 use sov_stf_runner::{FullNodeConfig, ProverConfig};
@@ -81,8 +83,10 @@ impl RollupBlueprint for MockDemoRollup {
         Ok(rpc_methods)
     }
 
-    fn get_code_commitment(&self) -> <Self::Vm as Zkvm>::CodeCommitment {
-        Digest::new(citrea_risc0::MOCK_DA_ID)
+    fn get_code_commitments_by_spec(&self) -> HashMap<SpecId, <Self::Vm as Zkvm>::CodeCommitment> {
+        let mut map = HashMap::new();
+        map.insert(SpecId::Genesis, Digest::new(citrea_risc0::MOCK_DA_ID));
+        map
     }
 
     async fn create_da_service(

diff --git a/bin/citrea/src/rollup/mod.rs b/bin/citrea/src/rollup/mod.rs
@@ -11,7 +11,6 @@ use sov_modules_api::storage::HierarchicalStorageManager;
 use sov_modules_api::Spec;
 use sov_modules_rollup_blueprint::RollupBlueprint;
 use sov_modules_stf_blueprint::{Runtime as RuntimeTrait, StfBlueprint};
-use sov_rollup_interface::spec::SpecId;
 use sov_state::storage::NativeStorage;
 use sov_stf_runner::{FullNodeConfig, InitVariant, ProverConfig};
 use tokio::sync::broadcast;
@@ -92,10 +91,7 @@ pub trait CitreaRollupBlueprint: RollupBlueprint {
             .map(|(l2_height, _)| l2_height)
             .unwrap_or(BatchNumber(0));
 
-        let active_spec: SpecId = ledger_db.get_active_fork()?;
-
-        let mut fork_manager =
-            ForkManager::new(current_l2_height.into(), active_spec, FORKS.to_vec());
+        let mut fork_manager = ForkManager::new(FORKS.to_vec(), current_l2_height.0);
         fork_manager.register_handler(Box::new(ledger_db.clone()));
 
         let seq = CitreaSequencer::new(
@@ -180,17 +176,15 @@ pub trait CitreaRollupBlueprint: RollupBlueprint {
             },
         };
 
-        let code_commitment = self.get_code_commitment();
+        let code_commitments_by_spec = self.get_code_commitments_by_spec();
 
         let current_l2_height = ledger_db
             .get_head_soft_confirmation()
             .map_err(|e| anyhow!("Failed to get head soft confirmation: {}", e))?
             .map(|(l2_height, _)| l2_height)
             .unwrap_or(BatchNumber(0));
 
-        let active_spec: SpecId = ledger_db.get_active_fork()?;
-        let mut fork_manager =
-            ForkManager::new(current_l2_height.into(), active_spec, FORKS.to_vec());
+        let mut fork_manager = ForkManager::new(FORKS.to_vec(), current_l2_height.0);
         fork_manager.register_handler(Box::new(ledger_db.clone()));
 
         let runner = CitreaFullnode::new(
@@ -202,7 +196,7 @@ pub trait CitreaRollupBlueprint: RollupBlueprint {
             native_stf,
             storage_manager,
             init_variant,
-            code_commitment,
+            code_commitments_by_spec,
             rollup_config.sync_blocks_count,
             fork_manager,
             soft_confirmation_tx,
@@ -286,17 +280,15 @@ pub trait CitreaRollupBlueprint: RollupBlueprint {
             },
         };
 
-        let code_commitment = self.get_code_commitment();
+        let code_commitments_by_spec = self.get_code_commitments_by_spec();
 
         let current_l2_height = ledger_db
             .get_head_soft_confirmation()
             .map_err(|e| anyhow!("Failed to get head soft confirmation: {}", e))?
             .map(|(l2_height, _)| l2_height)
             .unwrap_or(BatchNumber(0));
 
-        let active_spec: SpecId = ledger_db.get_active_fork()?;
-        let mut fork_manager =
-            ForkManager::new(current_l2_height.into(), active_spec, FORKS.to_vec());
+        let mut fork_manager = ForkManager::new(FORKS.to_vec(), current_l2_height.0);
         fork_manager.register_handler(Box::new(ledger_db.clone()));
 
         let runner = CitreaProver::new(
@@ -310,7 +302,7 @@ pub trait CitreaRollupBlueprint: RollupBlueprint {
             init_variant,
             Some(prover_service),
             Some(prover_config),
-            code_commitment,
+            code_commitments_by_spec,
             rollup_config.sync_blocks_count,
             fork_manager,
             soft_confirmation_tx,

diff --git a/crates/citrea-stf/src/verifier.rs b/crates/citrea-stf/src/verifier.rs
@@ -48,7 +48,7 @@ where
         )?;
 
         println!("going into apply_soft_confirmations_from_sequencer_commitments");
-        let (final_state_root, state_diff) = self
+        let (final_state_root, state_diff, last_active_spec_id) = self
             .app
             .apply_soft_confirmations_from_sequencer_commitments(
                 data.sequencer_public_key.as_ref(),
@@ -83,6 +83,7 @@ where
             sequencer_public_key: data.sequencer_public_key,
             sequencer_da_public_key: data.sequencer_da_public_key,
             sequencer_commitments_range: data.sequencer_commitments_range,
+            last_active_spec_id,
         };
 
         zkvm.commit(&out);

diff --git a/crates/fullnode/src/runner.rs b/crates/fullnode/src/runner.rs
@@ -1,4 +1,4 @@
-use std::collections::VecDeque;
+use std::collections::{HashMap, VecDeque};
 use std::marker::PhantomData;
 use std::net::SocketAddr;
 use std::sync::Arc;
@@ -7,7 +7,7 @@ use anyhow::{anyhow, bail};
 use backoff::future::retry as retry_backoff;
 use backoff::ExponentialBackoffBuilder;
 use borsh::de::BorshDeserialize;
-use citrea_primitives::fork::{Fork, ForkManager};
+use citrea_primitives::fork::ForkManager;
 use citrea_primitives::types::SoftConfirmationHash;
 use citrea_primitives::{get_da_block_at_height, L1BlockCache, SyncError};
 use jsonrpsee::core::client::Error as JsonrpseeError;
@@ -27,6 +27,7 @@ use sov_rollup_interface::da::{
 };
 use sov_rollup_interface::rpc::SoftConfirmationStatus;
 use sov_rollup_interface::services::da::{DaService, SlotData};
+use sov_rollup_interface::spec::SpecId;
 pub use sov_rollup_interface::stf::BatchReceipt;
 use sov_rollup_interface::stf::{SoftConfirmationReceipt, StateTransitionFunction};
 use sov_rollup_interface::storage::HierarchicalStorageManager;
@@ -65,7 +66,7 @@ where
     prover_da_pub_key: Vec<u8>,
     phantom: std::marker::PhantomData<C>,
     include_tx_body: bool,
-    code_commitment: Vm::CodeCommitment,
+    code_commitments_by_spec: HashMap<SpecId, Vm::CodeCommitment>,
     accept_public_input_as_proven: bool,
     l1_block_cache: Arc<Mutex<L1BlockCache<Da>>>,
     sync_blocks_count: u64,
@@ -103,7 +104,7 @@ where
         stf: Stf,
         mut storage_manager: Sm,
         init_variant: InitVariant<Stf, Vm, Da::Spec>,
-        code_commitment: Vm::CodeCommitment,
+        code_commitments_by_spec: HashMap<SpecId, Vm::CodeCommitment>,
         sync_blocks_count: u64,
         fork_manager: ForkManager,
         soft_confirmation_tx: broadcast::Sender<u64>,
@@ -151,7 +152,7 @@ where
             prover_da_pub_key: public_keys.prover_da_pub_key,
             phantom: std::marker::PhantomData,
             include_tx_body: runner_config.include_tx_body,
-            code_commitment,
+            code_commitments_by_spec,
             accept_public_input_as_proven: runner_config
                 .accept_public_input_as_proven
                 .unwrap_or(false),
@@ -230,47 +231,37 @@ where
             l1_block.header().height()
         );
         tracing::debug!("ZK proof: {:?}", proof);
-        let state_transition = match proof.clone() {
-            Proof::Full(proof) => {
-                let code_commitment = self.code_commitment.clone();
 
-                tracing::warn!(
-                    "using code commitment: {:?}",
-                    serde_json::to_string(&code_commitment).unwrap()
-                );
+        let state_transition =
+            Vm::extract_output::<<Da as DaService>::Spec, Stf::StateRoot>(&proof)
+                .expect("Proof should be deserializable");
+        if state_transition.sequencer_da_public_key != self.sequencer_da_pub_key
+            || state_transition.sequencer_public_key != self.sequencer_pub_key
+        {
+            return Err(anyhow!(
+                "Proof verification: Sequencer public key or sequencer da public key mismatch. Skipping proof."
+            ).into());
+        }
 
-                if let Ok(proof_data) = Vm::verify_and_extract_output::<
-                    <Da as DaService>::Spec,
-                    Stf::StateRoot,
-                >(&proof, &code_commitment)
-                {
-                    if proof_data.sequencer_da_public_key != self.sequencer_da_pub_key
-                        || proof_data.sequencer_public_key != self.sequencer_pub_key
-                    {
-                        return Err(anyhow!(
-                            "Proof verification: Sequencer public key or sequencer da public key mismatch. Skipping proof."
-                        ).into());
-                    }
-                    proof_data
-                } else {
-                    return Err(anyhow!(
-                        "Proof verification: SNARK verification failed. Skipping to next proof.."
-                    )
-                    .into());
-                }
+        match &proof {
+            Proof::Full(data) => {
+                let code_commitment = self
+                    .code_commitments_by_spec
+                    .get(&state_transition.last_active_spec_id)
+                    .expect("Proof public input must contain valid spec id");
+                Vm::verify(data, code_commitment)
+                    .map_err(|err| anyhow!("Failed to verify proof: {:?}. Skipping it...", err))?;
             }
             Proof::PublicInput(_) => {
                 if !self.accept_public_input_as_proven {
                     return Err(anyhow!(
-                        "Found public input in da block number: {:?}, Skipping to next proof..",
+                        "Found public input in da block number: {}, Skipping to next proof..",
                         l1_block.header().height(),
                     )
                     .into());
                 }
-                // public input is accepted only in tests, so ok to expect
-                Vm::extract_output(&proof).expect("Proof should be deserializable")
             }
-        };
+        }
 
         let stored_state_transition = StoredStateTransition {
             initial_state_root: state_transition.initial_state_root.as_ref().to_vec(),
@@ -437,7 +428,7 @@ where
             .create_storage_on_l2_height(l2_height)?;
 
         let slot_result = self.stf.apply_soft_confirmation(
-            self.fork_manager.active_fork(),
+            self.fork_manager.active_fork().spec_id,
             self.sequencer_pub_key.as_slice(),
             // TODO(https://github.com/Sovereign-Labs/sovereign-sdk/issues/1247): incorrect pre-state root in case of re-org
             &self.state_root,

diff --git a/crates/fullnode/tests/hash_stf.rs b/crates/fullnode/tests/hash_stf.rs
@@ -7,6 +7,7 @@ use sov_modules_api::Context;
 use sov_modules_stf_blueprint::StfBlueprintTrait;
 use sov_prover_storage_manager::{new_orphan_storage, SnapshotManager};
 use sov_rollup_interface::da::{BlobReaderTrait, BlockHeaderTrait, DaSpec};
+use sov_rollup_interface::fork::Fork;
 use sov_rollup_interface::spec::SpecId;
 use sov_rollup_interface::stf::{SlotResult, StateTransitionFunction};
 use sov_rollup_interface::zk::{CumulativeStateDiff, ValidityCondition, Zkvm};
@@ -245,8 +246,8 @@ impl<Vm: Zkvm, Cond: ValidityCondition, Da: DaSpec> StateTransitionFunction<Vm,
             Vec<sov_modules_api::SignedSoftConfirmationBatch>,
         >,
         _preproven_commitment_indicies: Vec<usize>,
-        _forks: Vec<(SpecId, u64)>,
-    ) -> (Self::StateRoot, CumulativeStateDiff) {
+        _forks: Vec<Fork>,
+    ) -> (Self::StateRoot, CumulativeStateDiff, SpecId) {
         todo!()
     }
 }

diff --git a/crates/fullnode/tests/runner_initialization_tests.rs b/crates/fullnode/tests/runner_initialization_tests.rs
@@ -1,3 +1,4 @@
+use std::collections::HashMap;
 use std::sync::Arc;
 
 use citrea_fullnode::CitreaFullnode;
@@ -6,6 +7,7 @@ use sov_db::ledger_db::LedgerDB;
 use sov_mock_da::{MockAddress, MockDaConfig, MockDaService, MockDaSpec, MockValidityCond};
 use sov_mock_zkvm::{MockCodeCommitment, MockZkvm};
 use sov_prover_storage_manager::ProverStorageManager;
+use sov_rollup_interface::fork::Fork;
 use sov_rollup_interface::spec::SpecId;
 use sov_state::DefaultStorageSpec;
 use sov_stf_runner::{
@@ -55,7 +57,7 @@ fn initialize_runner(
     sov_modules_api::default_context::DefaultContext,
     LedgerDB,
 > {
-    let specs = vec![(SpecId::Genesis, 0)];
+    let forks = vec![Fork::new(SpecId::Genesis, 0)];
     let da_storage_path = storage_path.join("da").to_path_buf();
     let rollup_storage_path = storage_path.join("rollup").to_path_buf();
 
@@ -112,7 +114,10 @@ fn initialize_runner(
     // let vm = MockZkvm::new(MockValidityCond::default());
     // let verifier = MockDaVerifier::default();
 
-    let fork_manager = ForkManager::new(0, SpecId::Genesis, specs);
+    let fork_manager = ForkManager::new(forks, 0);
+
+    let mut code_commitments_by_spec = HashMap::new();
+    code_commitments_by_spec.insert(SpecId::Genesis, MockCodeCommitment([1u8; 32]));
 
     CitreaFullnode::new(
         rollup_config.runner.unwrap(),
@@ -123,7 +128,7 @@ fn initialize_runner(
         stf,
         storage_manager,
         init_variant,
-        MockCodeCommitment([1u8; 32]),
+        code_commitments_by_spec,
         10,
         fork_manager,
         broadcast::channel(1).0,