fix: write proofs to sp1 stdin

Author: MarcosNicolau

aggregation-mode/src/zk/aggregator.rs

@@ -1,13 +1,13 @@
-use crate::zk::backends::sp1::{self, SP1AggregatedProof};
-use serde::{Deserialize, Serialize};
+use crate::zk::backends::sp1::{self};
+
+use super::backends::sp1::{SP1AggregationInput, SP1Proof};
 
-#[derive(Serialize, Deserialize)]
 pub enum ProgramInput {
-    SP1(sp1_aggregator::Input),
+    SP1(SP1AggregationInput),
 }
 
 pub enum AggregatedProof {
-    SP1(SP1AggregatedProof),
+    SP1(SP1Proof),
 }
 
 pub struct ProgramOutput {
@@ -29,6 +29,7 @@ impl ProgramOutput {
 pub enum ProofAggregationError {
     SP1Verification(sp1_sdk::SP1VerificationError),
     SP1Proving,
+    UnsupportedProof,
 }
 
 pub fn aggregate_proofs(input: ProgramInput) -> Result<ProgramOutput, ProofAggregationError> {

aggregation-mode/src/zk/backends/sp1.rs

@@ -1,4 +1,8 @@
-use sp1_sdk::{Prover, ProverClient, SP1ProofWithPublicValues, SP1Stdin, SP1VerifyingKey};
+use alloy::primitives::Keccak256;
+use sp1_aggregator::{ProofInput, SP1ProofInput};
+use sp1_sdk::{
+    HashableKey, Prover, ProverClient, SP1ProofWithPublicValues, SP1Stdin, SP1VerifyingKey,
+};
 
 use crate::zk::aggregator::{AggregatedProof, ProgramOutput, ProofAggregationError};
 
@@ -7,28 +11,55 @@ const PROGRAM_ELF: &[u8] = include_bytes!("../../../zkvm/sp1/elf/sp1_aggregator_
 // TODO lock prover
 
 pub struct SP1Proof {
-    pub elf: Vec<u8>,
+    pub vk: SP1VerifyingKey,
     pub proof: SP1ProofWithPublicValues,
 }
 
 impl SP1Proof {
-    pub fn verifying_key(&self) -> SP1VerifyingKey {
-        let client = ProverClient::from_env();
-        let (_pk, vk) = client.setup(&self.elf);
-        vk
+    pub fn hash(&self) -> [u8; 32] {
+        let mut hasher = Keccak256::new();
+        for &word in &self.vk.hash_u32() {
+            hasher.update(word.to_le_bytes());
+        }
+        hasher.update(self.proof.public_values.as_slice());
+        hasher.finalize().into()
     }
 }
 
-pub struct SP1AggregatedProof {
-    pub proof: SP1ProofWithPublicValues,
-    pub vk: SP1VerifyingKey,
+pub struct SP1AggregationInput {
+    proofs: Vec<SP1Proof>,
+    merkle_root: [u8; 32],
 }
 
 pub(crate) fn aggregate_proofs(
-    input: sp1_aggregator::Input,
+    input: SP1AggregationInput,
 ) -> Result<ProgramOutput, ProofAggregationError> {
     let mut stdin = SP1Stdin::new();
-    stdin.write(&input);
+
+    let mut program_input = sp1_aggregator::Input {
+        proofs: vec![],
+        merkle_root: input.merkle_root,
+    };
+
+    // write vk + public inputs
+    for proof in input.proofs.iter() {
+        program_input
+            .proofs
+            .push(ProofInput::SP1Compressed(SP1ProofInput {
+                public_inputs: proof.proof.public_values.to_vec(),
+                vk: proof.vk.hash_u32(),
+            }));
+    }
+    stdin.write(&program_input);
+
+    // write proofs
+    for SP1Proof { proof, vk } in input.proofs {
+        // we only support sp1 Compressed proofs for now
+        let sp1_sdk::SP1Proof::Compressed(proof) = proof.proof else {
+            return Err(ProofAggregationError::UnsupportedProof);
+        };
+        stdin.write_proof(*proof, vk.vk);
+    }
 
     #[cfg(feature = "prove")]
     let client = ProverClient::from_env();
@@ -48,7 +79,7 @@ pub(crate) fn aggregate_proofs(
         .verify(&proof, &vk)
         .map_err(ProofAggregationError::SP1Verification)?;
 
-    let proof = SP1AggregatedProof { proof, vk };
+    let proof = SP1Proof { proof, vk };
 
     let output = ProgramOutput::new(AggregatedProof::SP1(proof));
 
@@ -59,10 +90,10 @@ pub enum SP1VerificationError {
     Verification(sp1_sdk::SP1VerificationError),
 }
 
-pub(crate) fn verify(proof: &SP1Proof) -> Result<(), SP1VerificationError> {
+pub(crate) fn verify(proof: &SP1Proof, elf: &[u8]) -> Result<(), SP1VerificationError> {
     let client = ProverClient::from_env();
 
-    let (_pk, vk) = client.setup(&proof.elf);
+    let (_pk, vk) = client.setup(elf);
     client
         .verify(&proof.proof, &vk)
         .map_err(SP1VerificationError::Verification)

aggregation-mode/src/zk/mod.rs

@@ -11,14 +11,22 @@ pub enum Proof {
     SP1(SP1Proof),
 }
 
+impl Proof {
+    pub fn hash(&self) -> [u8; 32] {
+        match self {
+            Proof::SP1(proof) => proof.hash(),
+        }
+    }
+}
+
 pub enum VerificationError {
     SP1(SP1VerificationError),
 }
 
 impl Proof {
-    pub fn verify(&self) -> Result<(), VerificationError> {
+    pub fn verify(&self, elf: &[u8]) -> Result<(), VerificationError> {
         match self {
-            Proof::SP1(proof) => sp1::verify(proof).map_err(VerificationError::SP1),
+            Proof::SP1(proof) => sp1::verify(proof, elf).map_err(VerificationError::SP1),
         }
     }
 }

aggregation-mode/zkvm/sp1/src/lib.rs

@@ -2,47 +2,37 @@ use serde::{Deserialize, Serialize};
 use sha3::{Digest, Keccak256};
 
 #[derive(Serialize, Deserialize)]
-pub struct SP1CompressedProof {
-    pub vk: Vec<u8>,
+pub struct SP1ProofInput {
+    pub vk: [u32; 8],
     pub public_inputs: Vec<u8>,
 }
 
-impl SP1CompressedProof {
-    pub fn vk(&self) -> [u32; 8] {
-        assert!(self.vk.len() >= 32, "vk must be at least 32 bytes long");
-
-        let mut bytes = [0_32; 8];
-
-        for (i, chunk) in self.vk.chunks_exact(4).enumerate() {
-            bytes[i] = u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]);
-        }
-
-        bytes
-    }
-
+impl SP1ProofInput {
     pub fn hash(&self) -> [u8; 32] {
         let mut hasher = Keccak256::new();
-        hasher.update(&self.vk);
+        for &word in &self.vk {
+            hasher.update(word.to_le_bytes());
+        }
         hasher.update(&self.public_inputs);
         hasher.finalize().into()
     }
 }
 
 #[derive(Serialize, Deserialize)]
-pub enum Proof {
-    SP1Compressed(SP1CompressedProof),
+pub enum ProofInput {
+    SP1Compressed(SP1ProofInput),
 }
 
-impl Proof {
+impl ProofInput {
     pub fn hash(&self) -> [u8; 32] {
         match self {
-            Proof::SP1Compressed(proof) => proof.hash(),
+            ProofInput::SP1Compressed(proof) => proof.hash(),
         }
     }
 }
 
 #[derive(Serialize, Deserialize)]
 pub struct Input {
-    pub proofs: Vec<Proof>,
+    pub proofs: Vec<ProofInput>,
     pub merkle_root: [u8; 32],
 }

aggregation-mode/zkvm/sp1/src/main.rs

@@ -3,7 +3,7 @@ sp1_zkvm::entrypoint!(main);
 
 use sha2::{Digest, Sha256};
 use sha3::Keccak256;
-use sp1_aggregator::{Input, Proof};
+use sp1_aggregator::{Input, ProofInput};
 
 fn combine_hashes(hash_a: &[u8; 32], hash_b: &[u8; 32]) -> [u8; 32] {
     let mut hasher = Keccak256::new();
@@ -13,7 +13,7 @@ fn combine_hashes(hash_a: &[u8; 32], hash_b: &[u8; 32]) -> [u8; 32] {
 }
 
 /// Computes the merkle root for the given proofs using the vk
-fn compute_merkle_root(proofs: &[Proof]) -> [u8; 32] {
+fn compute_merkle_root(proofs: &[ProofInput]) -> [u8; 32] {
     let mut leaves: Vec<[u8; 32]> = proofs
         .chunks(2)
         .map(|chunk| match chunk {
@@ -44,8 +44,8 @@ pub fn main() {
     // Verify the proofs.
     for proof in input.proofs.iter() {
         match proof {
-            Proof::SP1Compressed(proof) => {
-                let vkey = proof.vk();
+            ProofInput::SP1Compressed(proof) => {
+                let vkey = proof.vk;
                 let public_values = &proof.public_inputs;
                 let public_values_digest = Sha256::digest(public_values);
                 sp1_zkvm::lib::verify::verify_sp1_proof(&vkey, &public_values_digest.into());