implement pre-execution before proving

Signed-off-by: Jun Kimura <jun.kimura@datachain.jp>

Author: bluele

Cargo.lock

@@ -265,6 +265,11 @@ pub struct ZKDCAPRemoteAttestation {
     pub bonsai_api_url: Option<String>,
     #[clap(long = "bonsai_api_key", help = "Bonsai API key")]
     pub bonsai_api_key: Option<String>,
+    #[clap(
+        long = "disable_pre_execution",
+        help = "Disable pre-execution before proving"
+    )]
+    pub disable_pre_execution: bool,
 }
 
 impl ZKDCAPRemoteAttestation {
@@ -323,6 +328,7 @@ fn run_zkdcap_remote_attestation<E: EnclaveCommandAPI<S>, S: CommitStore>(
         Address::from_hex_string(&cmd.enclave_key)?,
         mode,
         cmd.get_zkvm_program()?.as_ref(),
+        cmd.disable_pre_execution,
     )?;
     Ok(())
 }

app/src/commands/attestation.rs

@@ -24,6 +24,7 @@ reqwest = { version = "0.12.9", default-features = false, features = [
     "hickory-dns",
 ] }
 urlencoding = { version = "2" }
+anyhow = { version = "1" }
 
 dcap-quote-verifier = { git = "https://github.com/datachainlab/zkdcap", rev = "ff53632f49d2d8b00c9bcc456aeeeb139bacca9f" }
 

modules/remote-attestation/Cargo.toml

@@ -196,6 +196,10 @@ define_error! {
         DcapQuoteVerifier
         [TraceError<dcap_quote_verifier::Error>]
         |_| { "DCAP quote verifier error" },
+
+        Anyhow
+        [TraceError<anyhow::Error>]
+        |_| { "Anyhow error" },
     }
 }
 
@@ -216,3 +220,9 @@ impl From<reqwest::Error> for Error {
         Error::reqwest(e)
     }
 }
+
+impl From<zkvm::Error> for Error {
+    fn from(e: zkvm::Error) -> Self {
+        Error::zkvm(e)
+    }
+}

modules/remote-attestation/src/errors.rs

@@ -1,13 +1,14 @@
 use crate::{dcap::dcap_ra, errors::Error};
+use anyhow::anyhow;
 use attestation_report::{Risc0ZKVMProof, ZKDCAPQuote, ZKVMProof};
 use crypto::Address;
 use keymanager::EnclaveKeyManager;
 use lcp_types::Time;
 use log::*;
 use zkvm::{
     encode_seal,
-    prover::{prove, Risc0ProverMode},
-    risc0_zkvm::{compute_image_id, ExecutorEnv},
+    prover::{get_executor, prove, Risc0ProverMode},
+    risc0_zkvm::{compute_image_id, Executor, ExecutorEnv},
     verifier::verify_groth16_proof,
 };
 
@@ -16,12 +17,13 @@ pub fn run_zkdcap_ra(
     target_enclave_key: Address,
     prover_mode: Risc0ProverMode,
     elf: &[u8],
+    disable_pre_execution: bool,
 ) -> Result<(), Error> {
-    let image_id = compute_image_id(elf).unwrap();
+    let image_id = compute_image_id(elf)
+        .map_err(|e| Error::anyhow(anyhow!("cannot compute image id: {}", e)))?;
     info!(
-        "Run zkDCAP verification with prover_mode={} image_id={}",
-        prover_mode,
-        hex::encode(image_id.as_bytes())
+        "run zkDCAP verification with prover_mode={} image_id={} enclave_key={}",
+        prover_mode, image_id, target_enclave_key
     );
 
     let current_time = Time::now();
@@ -34,31 +36,50 @@ pub fn run_zkdcap_ra(
         current_time
     );
 
-    let env = ExecutorEnv::builder()
-        .write(&(
-            res.raw_quote.clone(),
-            res.collateral.to_bytes(),
-            current_time.as_unix_timestamp_secs(),
-        ))
-        .unwrap()
-        .build()
-        .unwrap();
+    if !disable_pre_execution {
+        info!("running pre-execution");
+        let res = get_executor()
+            .execute(
+                build_env(
+                    &res.raw_quote,
+                    &res.collateral.to_bytes(),
+                    current_time.as_unix_timestamp_secs(),
+                )?,
+                elf,
+            )
+            .map_err(|e| Error::anyhow(anyhow!("pre-execution failed: {}", e)))?;
+        info!(
+            "pre-execution done: exit_code={:?} cycles={} ",
+            res.exit_code,
+            res.cycles()
+        );
+    }
 
     info!("proving with prover mode: {:?}", prover_mode);
-    let prover_info = prove(&prover_mode, env, elf).unwrap();
+    let prover_info = prove(
+        &prover_mode,
+        build_env(
+            &res.raw_quote,
+            &res.collateral.to_bytes(),
+            current_time.as_unix_timestamp_secs(),
+        )?,
+        elf,
+    )?;
     info!("proving done: stats: {:?}", prover_info.stats);
 
-    prover_info.receipt.verify(image_id).unwrap();
+    prover_info
+        .receipt
+        .verify(image_id)
+        .map_err(|e| Error::anyhow(anyhow!("receipt verification failed: {}", e.to_string())))?;
     info!("receipt verified");
 
-    let seal = encode_seal(&prover_info.receipt).unwrap();
+    let seal = encode_seal(&prover_info.receipt)?;
     if let zkvm::risc0_zkvm::InnerReceipt::Groth16(_) = prover_info.receipt.inner {
         verify_groth16_proof(
             seal.clone(),
             image_id,
             prover_info.receipt.journal.bytes.clone(),
-        )
-        .unwrap();
+        )?;
     } else {
         assert!(
             prover_mode.is_dev_mode(),
@@ -90,3 +111,18 @@ pub fn run_zkdcap_ra(
 
     Ok(())
 }
+
+fn build_env<'a>(
+    quote: &[u8],
+    collateral: &[u8],
+    current_time: u64,
+) -> Result<ExecutorEnv<'a>, Error> {
+    ExecutorEnv::builder()
+        .write(&(quote, collateral, current_time))
+        .map_err(|e| Error::anyhow(anyhow!("cannot build env: {}", e)))
+        .and_then(|builder| {
+            builder
+                .build()
+                .map_err(|e| Error::anyhow(anyhow!("cannot build env: {}", e)))
+        })
+}

modules/remote-attestation/src/zkdcap.rs

@@ -1,6 +1,7 @@
 use crate::Error;
 use risc0_zkvm::{
-    BonsaiProver, ExecutorEnv, LocalProver, ProveInfo, Prover, ProverOpts, VerifierContext,
+    BonsaiProver, Executor, ExecutorEnv, LocalProver, ProveInfo, Prover, ProverOpts,
+    VerifierContext,
 };
 use std::fmt::Display;
 
@@ -112,3 +113,7 @@ impl BonsaiProverOptions {
             .ok_or_else(Error::missing_bonsai_api_key)
     }
 }
+
+pub fn get_executor() -> impl Executor {
+    LocalProver::new("local")
+}