fix: spawn blocking for verification (#697)

Co-authored-by: NicolasRampoldi <[email protected]>
yetanotherco · Aug 1, 2024 · ce153f7 · ce153f7
1 parent 83c4254
commit ce153f7
Show file tree

Hide file tree

Showing 5 changed files with 29 additions and 16 deletions.
diff --git a/batcher/Cargo.lock b/batcher/Cargo.lock
diff --git a/batcher/aligned-batcher/Cargo.toml b/batcher/aligned-batcher/Cargo.toml
@@ -19,14 +19,16 @@ bytes = "1.6.0"
 hex = "0.4.3"
 dotenv = "0.15.0"
 anyhow = "1.0.83"
-ethers = { tag = "v2.0.15-fix-reconnections", features = ["ws", "rustls"], git = "https://github.com/yetanotherco/ethers-rs.git" }
+ethers = { tag = "v2.0.15-fix-reconnections", features = [
+    "ws",
+    "rustls",
+], git = "https://github.com/yetanotherco/ethers-rs.git" }
 lambdaworks-crypto = { version = "0.7.0", features = ["serde"] }
 serde_yaml = "0.9.34"
 sp1-sdk = { git = "https://github.com/succinctlabs/sp1.git", rev = "v1.0.1" }
-risc0-zkvm = { git = "https://github.com/risc0/risc0", tag="v1.0.1" }
+risc0-zkvm = { git = "https://github.com/risc0/risc0", tag = "v1.0.1" }
 halo2curves = { version = "0.6.0", default-features = false }
-halo2_backend = { git = "https://github.com/yetanotherco/yet-another-halo2-fork.git", rev = "a3a56819d9183ac0b11c8d0543c7673c4a4c71a6"}
-halo2_proofs = { git = "https://github.com/yetanotherco/yet-another-halo2-fork.git", rev = "a3a56819d9183ac0b11c8d0543c7673c4a4c71a6"}
-lazy_static = "1.4.0"
+halo2_backend = { git = "https://github.com/yetanotherco/yet-another-halo2-fork.git", rev = "a3a56819d9183ac0b11c8d0543c7673c4a4c71a6" }
+halo2_proofs = { git = "https://github.com/yetanotherco/yet-another-halo2-fork.git", rev = "a3a56819d9183ac0b11c8d0543c7673c4a4c71a6" }
 bincode = "1.3.3"
-aligned-sdk = { path = "../aligned-sdk"}
+aligned-sdk = { path = "../aligned-sdk" }
diff --git a/batcher/aligned-batcher/src/lib.rs b/batcher/aligned-batcher/src/lib.rs
@@ -239,6 +239,11 @@ impl Batcher {
             serde_json::from_str(message.to_text().expect("Message is not text"))
                 .expect("Failed to deserialize task");
 
+        info!(
+            "Received message with nonce: {}",
+            U256::from_big_endian(client_msg.verification_data.nonce.as_slice())
+        );
+
         info!("Verifying message signature...");
         if let Ok(addr) = client_msg.verify_signature() {
             info!("Message signature verified");
@@ -270,7 +275,7 @@ impl Batcher {
 
                 // When pre-verification is enabled, batcher will verify proofs for faster feedback with clients
                 if self.pre_verification_is_enabled
-                    && !zk_utils::verify(&nonced_verification_data.verification_data)
+                    && !zk_utils::verify(&nonced_verification_data.verification_data).await
                 {
                     error!("Invalid proof detected. Verification failed.");
                     send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidProof).await;
@@ -703,7 +708,7 @@ impl Batcher {
         if client_msg.verification_data.verification_data.proof.len() <= self.max_proof_size {
             // When pre-verification is enabled, batcher will verify proofs for faster feedback with clients
             if self.pre_verification_is_enabled
-                && !zk_utils::verify(&client_msg.verification_data.verification_data)
+                && !zk_utils::verify(&client_msg.verification_data.verification_data).await
             {
                 error!("Invalid proof detected. Verification failed.");
                 send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidProof).await;

diff --git a/batcher/aligned-batcher/src/sp1/mod.rs b/batcher/aligned-batcher/src/sp1/mod.rs
@@ -1,16 +1,16 @@
-use lazy_static::lazy_static;
 use log::{debug, warn};
 use sp1_sdk::ProverClient;
+use std::sync::OnceLock;
 
-lazy_static! {
-    static ref SP1_PROVER_CLIENT: ProverClient = ProverClient::new();
-}
+static SP1_PROVER_CLIENT: OnceLock<ProverClient> = OnceLock::new();
 
 pub fn verify_sp1_proof(proof: &[u8], elf: &[u8]) -> bool {
     debug!("Verifying SP1 proof");
-    let (_pk, vk) = SP1_PROVER_CLIENT.setup(elf);
+    let prover_client = SP1_PROVER_CLIENT.get_or_init(ProverClient::new);
+
+    let (_pk, vk) = prover_client.setup(elf);
     if let Ok(proof) = bincode::deserialize(proof) {
-        let res = SP1_PROVER_CLIENT.verify(&proof, &vk).is_ok();
+        let res = prover_client.verify(&proof, &vk).is_ok();
         debug!("SP1 proof is valid: {}", res);
         if res {
             return true;

diff --git a/batcher/aligned-batcher/src/zk_utils/mod.rs b/batcher/aligned-batcher/src/zk_utils/mod.rs
@@ -6,7 +6,14 @@ use crate::sp1::verify_sp1_proof;
 use aligned_sdk::core::types::{ProvingSystemId, VerificationData};
 use log::{debug, warn};
 
-pub(crate) fn verify(verification_data: &VerificationData) -> bool {
+pub(crate) async fn verify(verification_data: &VerificationData) -> bool {
+    let verification_data = verification_data.clone();
+    tokio::task::spawn_blocking(move || verify_internal(&verification_data))
+        .await
+        .unwrap_or(false)
+}
+
+fn verify_internal(verification_data: &VerificationData) -> bool {
     match verification_data.proving_system {
         ProvingSystemId::SP1 => {
             if let Some(elf) = &verification_data.vm_program_code {