fix(raiko): fix r0 aggregation proof format (#386)

* fix r0 aggregation proof format * Update provers/risc0/driver/src/lib.rs Co-authored-by: Petar Vujović <[email protected]> * fix review comments * refine info print * refine info print * add timeout to sp1 wait_proof * refine info print --------- Co-authored-by: Petar Vujović <[email protected]>
taikoxyz · Oct 17, 2024 · 3cb6651 · 3cb6651
1 parent ec483b7
commit 3cb6651
Show file tree

Hide file tree

Showing 8 changed files with 85 additions and 45 deletions.
diff --git a/host/src/proof.rs b/host/src/proof.rs
@@ -211,8 +211,8 @@ impl ProofActor {
                 }
                 result = Self::handle_aggregate(request_clone, &opts) => {
                     match result {
-                        Ok(()) => {
-                            info!("Host handling message");
+                        Ok(status) => {
+                            info!("Host handling message: {status:?}");
                         }
                         Err(error) => {
                             error!("Worker failed due to: {error:?}");
@@ -231,20 +231,20 @@ impl ProofActor {
         while let Some(message) = self.receiver.recv().await {
             match message {
                 Message::Cancel(key) => {
-                    debug!("Message::Cancel task: {key:?}");
+                    debug!("Message::Cancel({key:?})");
                     if let Err(error) = self.cancel_task(key).await {
                         error!("Failed to cancel task: {error}")
                     }
                 }
                 Message::Task(proof_request) => {
-                    debug!("Message::Task proof_request: {proof_request:?}");
+                    debug!("Message::Task({proof_request:?})");
                     let running_task_count = self.running_tasks.lock().await.len();
                     if running_task_count < self.opts.concurrency_limit {
                         info!("Running task {proof_request:?}");
                         self.run_task(proof_request).await;
                     } else {
                         info!(
-                            "Task concurrency limit reached, current running {running_task_count:?}, pending: {:?}",
+                            "Task concurrency status: running:{running_task_count:?}, add {proof_request:?} to pending list[{:?}]",
                             self.pending_tasks.lock().await.len()
                         );
                         let mut pending_tasks = self.pending_tasks.lock().await;
@@ -253,9 +253,9 @@ impl ProofActor {
                 }
                 Message::TaskComplete(req) => {
                     // pop up pending task if any task complete
-                    debug!("Message::TaskComplete: {req:?}");
+                    debug!("Message::TaskComplete({req:?})");
                     info!(
-                        "task completed, current running {:?}, pending: {:?}",
+                        "task {req:?} completed, current running {:?}, pending: {:?}",
                         self.running_tasks.lock().await.len(),
                         self.pending_tasks.lock().await.len()
                     );
@@ -269,11 +269,13 @@ impl ProofActor {
                     }
                 }
                 Message::CancelAggregate(request) => {
+                    debug!("Message::CancelAggregate({request:?})");
                     if let Err(error) = self.cancel_aggregation_task(request).await {
                         error!("Failed to cancel task: {error}")
                     }
                 }
                 Message::Aggregate(request) => {
+                    debug!("Message::Aggregate({request:?})");
                     let permit = Arc::clone(&semaphore)
                         .acquire_owned()
                         .await

diff --git a/host/src/server/api/v3/proof/mod.rs b/host/src/server/api/v3/proof/mod.rs
@@ -1,18 +1,18 @@
-use axum::{debug_handler, extract::State, routing::post, Json, Router};
-use raiko_core::{
-    interfaces::{AggregationOnlyRequest, AggregationRequest, ProofRequest, ProofRequestOpt},
-    provider::get_task_data,
-};
-use raiko_tasks::{TaskDescriptor, TaskManager, TaskStatus};
-use utoipa::OpenApi;
-
 use crate::{
     interfaces::HostResult,
     metrics::{inc_current_req, inc_guest_req_count, inc_host_req_count},
     server::api::{v2, v3::Status},
     Message, ProverState,
 };
+use axum::{debug_handler, extract::State, routing::post, Json, Router};
+use raiko_core::{
+    interfaces::{AggregationOnlyRequest, AggregationRequest, ProofRequest, ProofRequestOpt},
+    provider::get_task_data,
+};
+use raiko_lib::prover::Proof;
+use raiko_tasks::{TaskDescriptor, TaskManager, TaskStatus};
 use tracing::{debug, info};
+use utoipa::OpenApi;
 
 mod aggregate;
 mod cancel;
@@ -125,8 +125,11 @@ async fn proof_handler(
         let mut proofs = Vec::with_capacity(tasks.len());
         for (task, req) in tasks {
             let raw_proof = manager.get_task_proof(&task).await?;
-            let proof = serde_json::from_slice(&raw_proof)?;
-            debug!("req: {req:?} gets proof: {proof:?}");
+            let proof: Proof = serde_json::from_slice(&raw_proof)?;
+            debug!(
+                "Aggregation sub-req: {req:?} gets proof {:?} with input {:?}.",
+                proof.proof, proof.input
+            );
             proofs.push(proof);
         }
 

diff --git a/provers/risc0/driver/src/lib.rs b/provers/risc0/driver/src/lib.rs
@@ -180,13 +180,18 @@ impl Prover for Risc0Prover {
             "Generate aggregatino receipt journal: {:?}",
             receipt.journal
         );
+        let block_proof_image_id = compute_image_id(RISC0_GUEST_ELF).unwrap();
         let aggregation_image_id = compute_image_id(RISC0_AGGREGATION_ELF).unwrap();
-        let enc_proof =
-            snarks::verify_groth16_snark_from_receipt(aggregation_image_id, receipt.clone())
-                .await
-                .map_err(|err| format!("Failed to verify SNARK: {err:?}"))?;
-        let snark_proof = format!("0x{}", hex::encode(enc_proof));
+        let proof_data = snarks::verify_aggregation_groth16_proof(
+            block_proof_image_id,
+            aggregation_image_id,
+            receipt.clone(),
+        )
+        .await
+        .map_err(|err| format!("Failed to verify SNARK: {err:?}"))?;
+        let snark_proof = alloy_primitives::hex::encode_prefixed(proof_data);
 
+        info!("Aggregation proof: {snark_proof:?}");
         let proof_gen_result = Ok(Risc0Response {
             proof: snark_proof,
             receipt: serde_json::to_string(&receipt).unwrap(),

diff --git a/provers/risc0/driver/src/methods/risc0_aggregation.rs b/provers/risc0/driver/src/methods/risc0_aggregation.rs
@@ -1,5 +1,5 @@
 pub const RISC0_AGGREGATION_ELF: &[u8] =
     include_bytes!("../../../guest/target/riscv32im-risc0-zkvm-elf/release/risc0-aggregation");
 pub const RISC0_AGGREGATION_ID: [u32; 8] = [
-    440526723, 3767976668, 67051936, 881100330, 2605787818, 1152192925, 943988177, 1141581874,
+    3593026424, 359928015, 3488866833, 2676323972, 1129344711, 55769507, 233041442, 3293280986,
 ];
diff --git a/provers/risc0/driver/src/methods/risc0_guest.rs b/provers/risc0/driver/src/methods/risc0_guest.rs
@@ -1,5 +1,5 @@
 pub const RISC0_GUEST_ELF: &[u8] =
     include_bytes!("../../../guest/target/riscv32im-risc0-zkvm-elf/release/risc0-guest");
 pub const RISC0_GUEST_ID: [u32; 8] = [
-    2426111784, 2252773481, 4093155148, 2853313326, 836865213, 1159934005, 790932950, 229907112,
+    2522428380, 1790994278, 397707036, 244564411, 3780865207, 1282154214, 1673205005, 3172292887,
 ];
diff --git a/provers/risc0/driver/src/snarks.rs b/provers/risc0/driver/src/snarks.rs
@@ -157,17 +157,53 @@ pub async fn verify_groth16_from_snark_receipt(
     let seal = encode(snark_receipt.snark.to_vec())?;
     let journal_digest = snark_receipt.journal.digest();
     let post_state_digest = snark_receipt.post_state_digest.digest();
-    verify_groth16_snark_impl(image_id, seal, journal_digest, post_state_digest).await
+    let encoded_proof =
+        verify_groth16_snark_impl(image_id, seal, journal_digest, post_state_digest).await?;
+    let proof = (encoded_proof, B256::from_slice(image_id.as_bytes()))
+        .abi_encode()
+        .iter()
+        .skip(32)
+        .copied()
+        .collect();
+    Ok(proof)
 }
 
-pub async fn verify_groth16_snark_from_receipt(
-    image_id: Digest,
+pub async fn verify_aggregation_groth16_proof(
+    block_proof_image_id: Digest,
+    aggregation_image_id: Digest,
     receipt: Receipt,
 ) -> Result<Vec<u8>> {
-    let seal = receipt.inner.groth16().unwrap().seal.clone();
+    let seal = receipt
+        .inner
+        .groth16()
+        .map_err(|e| anyhow::Error::msg(format!("receipt.inner.groth16() failed: {e:?}")))?
+        .seal
+        .clone();
     let journal_digest = receipt.journal.digest();
-    let post_state_digest = receipt.claim()?.as_value().unwrap().post.digest();
-    verify_groth16_snark_impl(image_id, seal, journal_digest, post_state_digest).await
+    let post_state_digest = receipt
+        .claim()?
+        .as_value()
+        .map_err(|e| anyhow::Error::msg(format!("receipt.claim()?.as_value() failed: {e:?}")))?
+        .post
+        .digest();
+    let encoded_proof = verify_groth16_snark_impl(
+        aggregation_image_id,
+        seal,
+        journal_digest,
+        post_state_digest,
+    )
+    .await?;
+    let proof = (
+        encoded_proof,
+        B256::from_slice(block_proof_image_id.as_bytes()),
+        B256::from_slice(aggregation_image_id.as_bytes()),
+    )
+        .abi_encode()
+        .iter()
+        .skip(32)
+        .copied()
+        .collect();
+    Ok(proof)
 }
 
 pub async fn verify_groth16_snark_impl(
@@ -209,14 +245,5 @@ pub async fn verify_groth16_snark_impl(
         tracing_err!("SNARK verification failed: {verify_call_res:?}!");
     }
 
-    Ok(make_risc0_groth16_proof(enc_seal, image_id))
-}
-
-pub fn make_risc0_groth16_proof(seal: Vec<u8>, image_id: Digest) -> Vec<u8> {
-    (seal, B256::from_slice(image_id.as_bytes()))
-        .abi_encode()
-        .iter()
-        .skip(32)
-        .copied()
-        .collect()
+    Ok(enc_seal)
 }
diff --git a/provers/sp1/driver/src/lib.rs b/provers/sp1/driver/src/lib.rs
@@ -24,6 +24,7 @@ use std::{
     borrow::BorrowMut,
     env, fs,
     path::{Path, PathBuf},
+    time::Duration,
 };
 use tracing::{debug, error, info};
 
@@ -170,7 +171,10 @@ impl Prover for Sp1Prover {
                 output.header.number
             );
             network_prover
-                .wait_proof::<sp1_sdk::SP1ProofWithPublicValues>(&proof_id, None)
+                .wait_proof::<sp1_sdk::SP1ProofWithPublicValues>(
+                    &proof_id,
+                    Some(Duration::from_secs(3600)),
+                )
                 .await
                 .map_err(|e| ProverError::GuestError(format!("Sp1: network proof failed {e:?}")))?
         };
@@ -253,9 +257,6 @@ impl Prover for Sp1Prover {
         _store: Option<&mut dyn IdWrite>,
     ) -> ProverResult<Proof> {
         let param = Sp1Param::deserialize(config.get("sp1").unwrap()).unwrap();
-
-        info!("aggregate proof with param: {param:?}");
-
         let block_inputs: Vec<B256> = input
             .proofs
             .iter()
@@ -312,6 +313,7 @@ impl Prover for Sp1Prover {
         let prove_result = client
             .prove(&pk, stdin)
             .plonk()
+            .timeout(Duration::from_secs(3600))
             .run()
             .expect("proving failed");
 

diff --git a/script/prove-block.sh b/script/prove-block.sh
@@ -144,13 +144,14 @@ for block in $(eval echo {$rangeStart..$rangeEnd}); do
 	fi
 
 	echo "- proving block $block"
-	curl --location --request POST 'http://localhost:8080/v3/proof' \
+	curl --location --request POST 'http://localhost:8080/v2/proof' \
 		--header 'Content-Type: application/json' \
 		--header 'Authorization: Bearer 4cbd753fbcbc2639de804f8ce425016a50e0ecd53db00cb5397912e83f5e570e' \
 		--data-raw "{
          \"network\": \"$chain\",
          \"l1_network\": \"$l1_network\",
          \"block_numbers\": [[$block, null], [$(($block+1)), null]],
+         \"block_number\": $block,
          \"prover\": \"$prover\",
          \"graffiti\": \"$graffiti\",
          $proofParam