feat: ability to verify e2e STARK proof + CLI command

Author: stephenh-axiom-xyz

crates/cli/src/commands/verify.rs

@@ -3,14 +3,20 @@ use std::path::{Path, PathBuf};
 use clap::Parser;
 use eyre::Result;
 use openvm_sdk::{
-    fs::{read_app_proof_from_file, read_app_vk_from_file},
+    fs::{
+        decode_from_file, read_agg_stark_pk_from_file, read_app_proof_from_file,
+        read_app_vk_from_file,
+    },
     Sdk,
 };
 
 use super::KeygenCargoArgs;
 #[cfg(feature = "evm-verify")]
 use crate::default::default_evm_halo2_verifier_path;
-use crate::util::{get_app_vk_path, get_files_with_ext, get_manifest_path_and_dir, get_target_dir};
+use crate::{
+    default::default_agg_stark_pk_path,
+    util::{get_app_vk_path, get_files_with_ext, get_manifest_path_and_dir, get_target_dir},
+};
 
 #[derive(Parser)]
 #[command(name = "verify", about = "Verify a proof")]
@@ -41,6 +47,15 @@ enum VerifySubCommand {
         #[command(flatten)]
         cargo_args: KeygenCargoArgs,
     },
+    Stark {
+        #[arg(
+            long,
+            action,
+            help = "Path to STARK proof, by default will search the working directory for a file with extension .stark.proof",
+            help_heading = "OpenVM Options"
+        )]
+        proof: Option<PathBuf>,
+    },
     #[cfg(feature = "evm-verify")]
     Evm {
         #[arg(
@@ -85,6 +100,28 @@ impl VerifyCmd {
                 let app_proof = read_app_proof_from_file(proof_path)?;
                 sdk.verify_app_proof(&app_vk, &app_proof)?;
             }
+            VerifySubCommand::Stark { proof } => {
+                let agg_stark_pk = read_agg_stark_pk_from_file(default_agg_stark_pk_path())
+                    .map_err(|e| {
+                        eyre::eyre!(
+                        "Failed to read aggregation STARK proving key: {}\nPlease run 'cargo openvm setup' first",
+                        e
+                    )
+                    })?;
+                let proof_path = if let Some(proof) = proof {
+                    proof.clone()
+                } else {
+                    let files = get_files_with_ext(Path::new("."), "stark.proof")?;
+                    if files.len() > 1 {
+                        return Err(eyre::eyre!("multiple .stark.proof files found, please specify the path using option --proof"));
+                    } else if files.is_empty() {
+                        return Err(eyre::eyre!("no .stark.proof file found, please specify the path using option --proof"));
+                    }
+                    files[0].clone()
+                };
+                let stark_proof = decode_from_file(proof_path)?;
+                sdk.verify_e2e_stark_proof(agg_stark_pk, &stark_proof)?;
+            }
             #[cfg(feature = "evm-verify")]
             VerifySubCommand::Evm { proof } => {
                 use openvm_sdk::fs::{

crates/sdk/src/lib.rs

@@ -1,4 +1,4 @@
-use std::{fs::read, marker::PhantomData, path::Path, sync::Arc};
+use std::{borrow::Borrow, fs::read, marker::PhantomData, path::Path, sync::Arc};
 
 #[cfg(feature = "evm-verify")]
 use alloy_sol_types::sol;
@@ -11,16 +11,19 @@ use openvm_build::{
 };
 use openvm_circuit::{
     arch::{
-        hasher::poseidon2::vm_poseidon2_hasher, instructions::exe::VmExe, verify_segments,
-        ContinuationVmProof, ExecutionError, InitFileGenerator, VerifiedExecutionPayload, VmConfig,
-        VmExecutor, VmVerificationError,
+        hasher::{poseidon2::vm_poseidon2_hasher, Hasher},
+        instructions::exe::VmExe,
+        verify_segments, ContinuationVmProof, ExecutionError, InitFileGenerator,
+        VerifiedExecutionPayload, VmConfig, VmExecutor, PROGRAM_CACHED_TRACE_INDEX,
+        PUBLIC_VALUES_AIR_ID,
     },
     system::{
         memory::{tree::public_values::extract_public_values, CHUNK},
-        program::trace::VmCommittedExe,
+        program::trace::{compute_exe_commit, VmCommittedExe},
     },
 };
 use openvm_continuations::verifier::{
+    common::types::VmVerifierPvs,
     internal::types::VmStarkProof,
     root::{types::RootVmVerifierInput, RootVmVerifierConfig},
 };
@@ -33,7 +36,7 @@ use openvm_stark_backend::proof::Proof;
 use openvm_stark_sdk::{
     config::{baby_bear_poseidon2::BabyBearPoseidon2Engine, FriParameters},
     engine::StarkFriEngine,
-    openvm_stark_backend::{verifier::VerificationError, Chip},
+    openvm_stark_backend::Chip,
 };
 use openvm_transpiler::{
     elf::Elf,
@@ -228,7 +231,7 @@ impl<E: StarkFriEngine<SC>> GenericSdk<E> {
         &self,
         app_vk: &AppVerifyingKey,
         proof: &ContinuationVmProof<SC>,
-    ) -> Result<VerifiedContinuationVmPayload, VmVerificationError> {
+    ) -> Result<VerifiedContinuationVmPayload> {
         let engine = E::new(app_vk.fri_params);
         let VerifiedExecutionPayload {
             exe_commit,
@@ -250,9 +253,10 @@ impl<E: StarkFriEngine<SC>> GenericSdk<E> {
         &self,
         app_vk: &AppVerifyingKey,
         proof: &Proof<SC>,
-    ) -> Result<(), VerificationError> {
+    ) -> Result<()> {
         let e = E::new(app_vk.fri_params);
-        e.verify(&app_vk.app_vm_vk, proof)
+        e.verify(&app_vk.app_vm_vk, proof)?;
+        Ok(())
     }
 
     pub fn agg_keygen(
@@ -322,6 +326,51 @@ impl<E: StarkFriEngine<SC>> GenericSdk<E> {
         Ok(proof)
     }
 
+    pub fn verify_e2e_stark_proof(
+        &self,
+        agg_stark_pk: AggStarkProvingKey,
+        proof: &VmStarkProof<SC>,
+    ) -> Result<[F; CHUNK]> {
+        let program_commit =
+            proof.proof.commitments.main_trace[PROGRAM_CACHED_TRACE_INDEX].as_ref();
+        let internal_commit: &[_; CHUNK] = &agg_stark_pk
+            .internal_committed_exe
+            .get_program_commit()
+            .into();
+
+        let vm_pk = if program_commit == internal_commit {
+            &agg_stark_pk.internal_vm_pk
+        } else {
+            &agg_stark_pk.leaf_vm_pk
+        };
+        let e = E::new(vm_pk.fri_params);
+        e.verify(&vm_pk.vm_pk.get_vk(), &proof.proof)?;
+
+        let public_values_air_proof_data = proof
+            .proof
+            .per_air
+            .iter()
+            .find(|p| p.air_id == PUBLIC_VALUES_AIR_ID)
+            .unwrap();
+        let pvs: &VmVerifierPvs<_> =
+            public_values_air_proof_data.public_values[..VmVerifierPvs::<u8>::width()].borrow();
+
+        let hasher = vm_poseidon2_hasher();
+        let public_values_root = hasher.merkle_root(&proof.user_public_values);
+        if public_values_root != pvs.public_values_commit {
+            return Err(eyre::eyre!(
+                "Invalid public values root: expected {:?}, got {:?}",
+                pvs.public_values_commit,
+                public_values_root
+            ));
+        }
+
+        let start_pc = pvs.connector.initial_pc;
+        let initial_memory_root = &pvs.memory.initial_root;
+        let exe_commit = compute_exe_commit(&hasher, program_commit, initial_memory_root, start_pc);
+        Ok(exe_commit)
+    }
+
     #[cfg(feature = "evm-prove")]
     pub fn generate_evm_proof<VC: VmConfig<F>>(
         &self,