Verify proofs with the corresponding spec ids

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)]

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(

bin/citrea/src/rollup/mod.rs

@@ -176,7 +176,7 @@ 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()
@@ -196,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,
@@ -280,7 +280,7 @@ 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()
@@ -302,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,

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;
@@ -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,39 @@ 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, receipt) =
+            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(_) => {
+                let receipt = receipt.expect("Receipt must exist with full proof");
+                let code_commitment = self
+                    .code_commitments_by_spec
+                    .get(&state_transition.final_spec_id)
+                    .expect("Proof public input must contain valid spec id");
+                receipt
+                    .verify(code_commitment.clone())
+                    .map_err(|err| anyhow!("Proof verification: SNARK verification failed: {}. Skipping to next proof...", 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(),

crates/fullnode/tests/runner_initialization_tests.rs

@@ -1,3 +1,4 @@
+use std::collections::HashMap;
 use std::sync::Arc;
 
 use citrea_fullnode::CitreaFullnode;
@@ -115,6 +116,9 @@ fn initialize_runner(
 
     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(),
         rollup_config.public_keys,
@@ -124,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,

crates/prover/src/runner.rs

@@ -1,5 +1,5 @@
 use core::panic;
-use std::collections::VecDeque;
+use std::collections::{HashMap, VecDeque};
 use std::marker::PhantomData;
 use std::net::SocketAddr;
 use std::sync::Arc;
@@ -26,6 +26,7 @@ use sov_modules_stf_blueprint::StfBlueprintTrait;
 use sov_rollup_interface::da::{BlockHeaderTrait, DaData, DaSpec, SequencerCommitment};
 use sov_rollup_interface::rpc::SoftConfirmationStatus;
 use sov_rollup_interface::services::da::DaService;
+use sov_rollup_interface::spec::SpecId;
 use sov_rollup_interface::stf::{SoftConfirmationReceipt, StateTransitionFunction};
 use sov_rollup_interface::zk::{Proof, StateTransitionData, ZkvmHost};
 use sov_stf_runner::{
@@ -71,7 +72,7 @@ where
     sequencer_da_pub_key: Vec<u8>,
     phantom: std::marker::PhantomData<C>,
     prover_config: Option<ProverConfig>,
-    code_commitment: Vm::CodeCommitment,
+    code_commitments_by_spec: HashMap<SpecId, Vm::CodeCommitment>,
     l1_block_cache: Arc<Mutex<L1BlockCache<Da>>>,
     sync_blocks_count: u64,
     fork_manager: ForkManager,
@@ -111,7 +112,7 @@ where
         init_variant: InitVariant<Stf, Vm, Da::Spec>,
         prover_service: Option<Ps>,
         prover_config: Option<ProverConfig>,
-        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>,
@@ -158,7 +159,7 @@ where
             sequencer_da_pub_key: public_keys.sequencer_da_pub_key,
             phantom: std::marker::PhantomData,
             prover_config,
-            code_commitment,
+            code_commitments_by_spec,
             l1_block_cache: Arc::new(Mutex::new(L1BlockCache::new())),
             sync_blocks_count,
             fork_manager,
@@ -790,28 +791,22 @@ where
 
         // l1_height => (tx_id, proof, transition_data)
         // save proof along with tx id to db, should be queriable by slot number or slot hash
-        let transition_data: sov_modules_api::StateTransition<
-            <Da as DaService>::Spec,
-            Stf::StateRoot,
-        > = Vm::extract_output(&proof).expect("Proof should be deserializable");
+        let (transition_data, receipt) =
+            Vm::extract_output::<<Da as DaService>::Spec, Stf::StateRoot>(&proof)
+                .expect("Proof should be deserializable");
 
         match proof {
             Proof::PublicInput(_) => {
                 warn!("Proof is public input, skipping");
             }
-            Proof::Full(ref proof) => {
+            Proof::Full(_) => {
                 info!("Verifying proof!");
-                let transition_data_from_proof =
-                    Vm::verify_and_extract_output::<<Da as DaService>::Spec, Stf::StateRoot>(
-                        &proof.clone(),
-                        &self.code_commitment,
-                    )
-                    .expect("Proof should be verifiable");
-
-                info!(
-                    "transition data from proof: {:?}",
-                    transition_data_from_proof
-                );
+                let receipt = receipt.expect("Receipt must exist in full proof");
+                let code_commitment = self
+                    .code_commitments_by_spec
+                    .get(&transition_data.final_spec_id)
+                    .expect("Proof public input must contain valid spec id");
+                receipt.verify(code_commitment.clone())?;
             }
         }
 

crates/risc0-bonsai/src/host.rs

@@ -358,18 +358,24 @@ impl<'a> ZkvmHost for Risc0BonsaiHost<'a> {
 
     fn extract_output<Da: sov_rollup_interface::da::DaSpec, Root: BorshDeserialize>(
         proof: &Proof,
-    ) -> Result<sov_rollup_interface::zk::StateTransition<Da, Root>, Self::Error> {
-        let journal = match proof {
+    ) -> Result<
+        (
+            sov_rollup_interface::zk::StateTransition<Da, Root>,
+            Option<Receipt>,
+        ),
+        Self::Error,
+    > {
+        let (journal, receipt) = match proof {
             Proof::PublicInput(journal) => {
                 let journal: Journal = bincode::deserialize(journal)?;
-                journal
+                (journal, None)
             }
             Proof::Full(data) => {
                 let receipt: Receipt = bincode::deserialize(data)?;
-                receipt.journal
+                (receipt.journal.clone(), Some(receipt))
             }
         };
-        Ok(BorshDeserialize::try_from_slice(&journal.bytes)?)
+        Ok((BorshDeserialize::try_from_slice(&journal.bytes)?, receipt))
     }
 
     fn recover_proving_sessions(&self) -> Result<Vec<Proof>, anyhow::Error> {

crates/sovereign-sdk/adapters/mock-zkvm/Cargo.toml

@@ -16,6 +16,7 @@ publish = true
 anyhow = { workspace = true }
 borsh = { workspace = true }
 bincode = { workspace = true }
+risc0-zkvm = { workspace = true }
 serde = { workspace = true }
 sov-rollup-interface = { path = "../../rollup-interface" }