zkboost support mock zkvm (#35)

* feat: add mock zkVM support in `zkboost-server`

* chore: loops to get chain config

* feat: add `mocknet` deployment example

* feat: add `GET /health` endpoint for EWS

* fix: verify test

Author: han0110

Cargo.lock

@@ -48,6 +48,7 @@ futures = "0.3"
 indexmap = "2.12"
 lru = "0.12"
 nix = "0.29"
+rand = "0.9.2"
 reqwest = "0.12"
 serde = "1.0"
 serde_json = "1"

Cargo.toml

@@ -1,7 +1,6 @@
-//! Relayer for execution proof.
+//! Execution witness sentry (EWS).
 //!
-//! This relayer orchestrates the complete workflow for generating proofs of
-//! execution proof:
+//! EWS orchestrates the complete workflow for generating proofs of execution proof:
 //!
 //! 1. Listen to new block from CL
 //! 2. Fetch execution witness from EL
@@ -12,7 +11,7 @@
 //! ## Architecture
 //!
 //! ```text
-//!   CL          Relayer               EL            Proof Engine
+//!   CL            EWS                 EL            Proof Engine
 //!   |              |                  |                  |
 //!   |--new block-->|                  |                  |
 //!   |              |                  |                  |
@@ -31,7 +30,7 @@
 //!   |              |                  |                  |
 //! ```
 
-use std::{num::NonZeroUsize, path::PathBuf, sync::Arc};
+use std::{num::NonZeroUsize, path::PathBuf, sync::Arc, time::Duration};
 
 use anyhow::bail;
 use clap::Parser;
@@ -50,6 +49,7 @@ use lru::LruCache;
 use tokio::{
     signal::unix::{SignalKind, signal},
     sync::Mutex,
+    time::sleep,
 };
 use tokio_util::sync::CancellationToken;
 use tracing::{info, warn};
@@ -99,21 +99,19 @@ async fn main() -> anyhow::Result<()> {
 
     // Get chain config.
 
-    let mut chain_config = None;
-    for el_client in &el_clients {
-        if let Ok(Some(c)) = el_client.get_chain_config().await {
-            chain_config = Some(c);
-            break;
-        } else {
+    let chain_config = 'outer: loop {
+        for el_client in &el_clients {
+            if let Ok(Some(chain_config)) = el_client.get_chain_config().await {
+                break 'outer chain_config;
+            }
             warn!(
                 name = %el_client.name(),
                 url = %el_client.url(),
                 "Failed to get chain config",
             )
-        };
-    }
-    let Some(chain_config) = chain_config else {
-        bail!("Failed to get chain config from any EL endpoint");
+        }
+
+        sleep(Duration::from_secs(2)).await;
     };
 
     // Initialize CL clients.
@@ -137,7 +135,7 @@ async fn main() -> anyhow::Result<()> {
                 }
             }
             Err(e) => {
-                warn!(name = %endpoint.name, error = %e, "Failed to check zkvm status");
+                warn!(name = %endpoint.name, error = %e, "Failed to check if client has zkVM enabled or not");
             }
         }
     }

crates/execution-witness-sentry/src/main.rs

@@ -33,7 +33,12 @@ use std::{
 
 use alloy_genesis::ChainConfig;
 use alloy_primitives::B256;
-use axum::{Json, Router, extract::State, http::StatusCode, routing::post};
+use axum::{
+    Json, Router,
+    extract::State,
+    http::StatusCode,
+    routing::{get, post},
+};
 use lru::LruCache;
 use reth_stateless::StatelessInput;
 use tokio::{
@@ -691,6 +696,7 @@ async fn spawn_webhook_server(
     Ok(tokio::spawn(async move {
         let app = Router::new()
             .route("/proofs", post(proof_webhook))
+            .route("/health", get(StatusCode::OK))
             .with_state(state)
             .layer(TraceLayer::new_for_http())
             // 10MB limit to account for the proof size

crates/execution-witness-sentry/src/service/proof.rs

@@ -104,6 +104,15 @@ pub enum zkVMConfig {
         /// Unique identifier for this program.
         program_id: ProgramID,
     },
+    /// Mock zkVM configuration
+    Mock {
+        /// Mock time in millisecond for proof generation.
+        mock_proving_time_ms: u64,
+        /// Mock size for proof.
+        mock_proof_size: u64,
+        /// Unique identifier for this program.
+        program_id: ProgramID,
+    },
 }
 
 impl zkVMConfig {
@@ -112,6 +121,7 @@ impl zkVMConfig {
         match self {
             Self::Docker { program_id, .. } => program_id,
             Self::External { program_id, .. } => program_id,
+            Self::Mock { program_id, .. } => program_id,
         }
     }
 }
@@ -150,6 +160,11 @@ mod test {
             [[zkvm]]
             endpoint = "http://remote:3000"
             program_id = "external-test"
+
+            [[zkvm]]
+            mock_proving_time_ms = 1000
+            mock_proof_size = 1024
+            program_id = "mock"
         "#;
         assert_eq!(Config::from_toml_str(toml).unwrap(), sample_config());
     }
@@ -178,6 +193,9 @@ mod test {
                 url: http://artifact
             - endpoint: "http://remote:3000"
               program_id: "external-test"
+            - mock_proving_time_ms: 1000
+              mock_proof_size: 1024
+              program_id: "mock"
         "#;
         assert_eq!(Config::from_yaml_str(yaml).unwrap(), sample_config());
     }
@@ -216,6 +234,11 @@ mod test {
                     endpoint: "http://remote:3000".to_string(),
                     program_id: "external-test".into(),
                 },
+                zkVMConfig::Mock {
+                    mock_proving_time_ms: 1000,
+                    mock_proof_size: 1 << 10,
+                    program_id: "mock".into(),
+                },
             ],
         }
     }

crates/zkboost/server-config/src/config.rs

@@ -13,6 +13,7 @@ anyhow.workspace = true
 axum = { workspace = true, features = ["macros"] }
 clap = { workspace = true, features = ["derive"] }
 futures.workspace = true
+rand.workspace = true
 reqwest = { workspace = true, features = ["json"] }
 sysinfo.workspace = true
 tokio = { workspace = true, features = ["full"] }

crates/zkboost/server/Cargo.toml

@@ -27,6 +27,7 @@ use crate::{
         execute::execute_program, info::get_server_info, prove::prove_program, verify::verify_proof,
     },
     metrics::http_metrics_middleware,
+    mock::MockzkVM,
     proof_service::{ProofMessage, ProofService},
 };
 
@@ -97,8 +98,7 @@ pub(crate) enum zkVMInstance {
         client: Arc<zkVMClient>,
     },
     /// Mock zkVM
-    #[cfg(test)]
-    Mock(crate::mock::MockzkVM),
+    Mock(MockzkVM),
 }
 
 impl zkVMInstance {
@@ -118,6 +118,11 @@ impl zkVMInstance {
         })
     }
 
+    /// Creates an mock zkVM instance.
+    pub(crate) fn mock(mock_proving_time_ms: u64, mock_proof_size: u64) -> Self {
+        Self::Mock(MockzkVM::new(mock_proving_time_ms, mock_proof_size))
+    }
+
     /// Executes the program with the given input.
     pub(crate) async fn execute(
         &self,
@@ -126,8 +131,7 @@ impl zkVMInstance {
         match self {
             Self::Docker { vm } => vm.execute_async(input).await,
             Self::External { client } => Ok(client.execute(input).await?),
-            #[cfg(test)]
-            Self::Mock(vm) => vm.execute(&input),
+            Self::Mock(vm) => vm.execute(&input).await,
         }
     }
 
@@ -140,8 +144,7 @@ impl zkVMInstance {
         match self {
             Self::Docker { vm } => vm.prove_async(input, proof_kind).await,
             Self::External { client } => Ok(client.prove(input, proof_kind).await?),
-            #[cfg(test)]
-            Self::Mock(vm) => vm.prove(&input, proof_kind),
+            Self::Mock(vm) => vm.prove(&input, proof_kind).await,
         }
     }
 
@@ -151,8 +154,7 @@ impl zkVMInstance {
         match self {
             Self::Docker { vm } => vm.verify_async(proof).await,
             Self::External { client } => Ok(client.verify(proof).await?),
-            #[cfg(test)]
-            Self::Mock(vm) => vm.verify(&proof),
+            Self::Mock(vm) => vm.verify(&proof).await,
         }
     }
 }
@@ -194,5 +196,10 @@ async fn init_zkvm(config: &zkVMConfig) -> anyhow::Result<zkVMInstance> {
             Ok(zkVMInstance::docker(zkvm))
         }
         zkVMConfig::External { endpoint, .. } => zkVMInstance::external(endpoint.clone()).await,
+        zkVMConfig::Mock {
+            mock_proving_time_ms,
+            mock_proof_size,
+            ..
+        } => Ok(zkVMInstance::mock(*mock_proving_time_ms, *mock_proof_size)),
     }
 }

crates/zkboost/server/src/app.rs

@@ -54,7 +54,7 @@ mod tests {
     use axum::{Json, extract::State, http::StatusCode};
     use zkboost_types::{ExecuteRequest, ProgramID};
 
-    use crate::{app::execute::execute_program, mock::mock_app_state};
+    use crate::{app::execute::execute_program, mock::tests::mock_app_state};
 
     #[tokio::test]
     async fn test_execute_success() {

crates/zkboost/server/src/app/execute.rs

@@ -42,7 +42,7 @@ mod tests {
     use axum::{Json, extract::State, http::StatusCode};
     use zkboost_types::{ProgramID, ProveRequest};
 
-    use crate::{app::prove::prove_program, mock::mock_app_state};
+    use crate::{app::prove::prove_program, mock::tests::mock_app_state};
 
     #[tokio::test]
     async fn test_prove_success() {

crates/zkboost/server/src/app/prove.rs

@@ -56,17 +56,24 @@ mod tests {
     use axum::{Json, extract::State, http::StatusCode};
     use zkboost_types::{ProgramID, VerifyRequest};
 
-    use crate::{app::verify::verify_proof, mock::mock_app_state};
+    use crate::{
+        app::{verify::verify_proof, zkVMInstance},
+        mock::tests::mock_app_state,
+    };
 
     #[tokio::test]
     async fn test_verify_valid_proof() {
         let program_id = ProgramID::from("mock_program_id");
         let state = mock_app_state(Some(&program_id));
 
+        let zkVMInstance::Mock(zkvm) = &state.programs[&program_id] else {
+            unreachable!()
+        };
+
         // Create a request with the mock proof that MockzkVM accepts
         let request = VerifyRequest {
             program_id: program_id.clone(),
-            proof: b"mock_proof".to_vec(),
+            proof: zkvm.random_proof(),
         };
 
         // Call the handler