dig-capsule 0.5.0

use crate::imp::core::{Bytes32, ChiaBlockRef, ExecutionProof};
use crate::imp::crypto::{bls, sha256};
use crate::imp::prover::chain::ChainSource;
use crate::imp::prover::commitment::{parse_public_input, signing_message};
use crate::imp::prover::error::{ProverError, Result};
use crate::imp::prover::mock::DEFAULT_FRESHNESS_WINDOW_SECS;
use crate::imp::prover::prover::{Prover, Verifier};
use crate::imp::prover::serving_inputs::ServingInputs;
use risc0_zkvm::{default_prover, ExecutorEnv, Receipt};
use serde::{Deserialize, Serialize};

// Generated by risc0-build (build.rs `embed_methods`): for the guest bin named
// `dig-capsule-serving-guest`, these are `DIG_CAPSULE_SERVING_GUEST_ELF` (&[u8]) and
// `DIG_CAPSULE_SERVING_GUEST_ID` ([u32; 8]).
include!(concat!(env!("OUT_DIR"), "/methods.rs"));

/// Wire input to the guest. MUST be byte-identical to the guest-side
/// `GuestInput` in `guest/src/main.rs`.
#[derive(Serialize, Deserialize)]
struct GuestInput {
    program_hash: [u8; 32],
    public_input: Vec<u8>,
    roothash: [u8; 32],
    chunks: Vec<Vec<u8>>,
}

/// Real ZK prover: runs the serving computation in a risc0 guest (§13.1-13.3).
pub struct Risc0Prover {
    node_secret: bls::SecretKey,
    node_pubkey: bls::PublicKey,
    chia_block: ChiaBlockRef,
}

impl Risc0Prover {
    pub fn new(
        node_secret: bls::SecretKey,
        node_pubkey: bls::PublicKey,
        chia_block: ChiaBlockRef,
    ) -> Self {
        Self {
            node_secret,
            node_pubkey,
            chia_block,
        }
    }
}

impl Prover for Risc0Prover {
    fn prove(
        &self,
        program_hash: Bytes32,
        public_input: &[u8],
        serving_inputs: &ServingInputs,
    ) -> Result<ExecutionProof> {
        let (_nonce, block) = parse_public_input(public_input)?;
        if block != self.chia_block {
            return Err(ProverError::Backend("public_input block mismatch".into()));
        }
        let guest_input = GuestInput {
            program_hash: program_hash.0,
            public_input: public_input.to_vec(),
            roothash: serving_inputs.roothash.0,
            chunks: serving_inputs.chunk_ciphertext.clone(),
        };
        let env = ExecutorEnv::builder()
            .write(&guest_input)
            .map_err(|e| ProverError::Backend(format!("env write: {e}")))?
            .build()
            .map_err(|e| ProverError::Backend(format!("env build: {e}")))?;
        let receipt = default_prover()
            .prove(env, DIG_CAPSULE_SERVING_GUEST_ELF)
            .map_err(|e| ProverError::Backend(format!("prove: {e}")))?
            .receipt;
        let proof = bincode::serialize(&receipt)
            .map_err(|e| ProverError::Backend(format!("receipt ser: {e}")))?;
        let public_output = serving_inputs.compute_public_output();
        let msg = signing_message(&proof, public_input);
        let node_sig = bls::bls_sign(&self.node_secret, &msg);
        Ok(ExecutionProof {
            program_hash,
            public_input: public_input.to_vec(),
            public_output,
            proof,
            chia_block: self.chia_block.clone(),
            node_pubkey: self.node_pubkey.to_bytes(),
            node_signature: node_sig,
        })
    }
}

/// Verifier for [`Risc0Prover`] proofs.
#[derive(Default)]
pub struct Risc0Verifier;

impl Verifier for Risc0Verifier {
    fn verify(
        &self,
        proof: &ExecutionProof,
        expected_program_hash: Bytes32,
        trusted_roots: &[Bytes32],
        chain: &dyn ChainSource,
    ) -> Result<()> {
        if proof.program_hash != expected_program_hash {
            return Err(ProverError::ProgramHashMismatch {
                expected: expected_program_hash.to_hex(),
                actual: proof.program_hash.to_hex(),
            });
        }
        let (_nonce, pi_block) = parse_public_input(&proof.public_input)?;
        if pi_block != proof.chia_block {
            return Err(ProverError::Codec(
                "public_input block != proof.chia_block".into(),
            ));
        }
        let receipt: Receipt = bincode::deserialize(&proof.proof)
            .map_err(|e| ProverError::ZkProofInvalid(format!("receipt de: {e}")))?;
        receipt
            .verify(DIG_CAPSULE_SERVING_GUEST_ID)
            .map_err(|e| ProverError::ZkProofInvalid(format!("receipt verify: {e}")))?;
        // Journal: (program_hash, public_input_hash, roothash, public_output)
        let (j_program_hash, j_pi_hash, _j_root, j_output): (
            [u8; 32],
            [u8; 32],
            [u8; 32],
            [u8; 32],
        ) = receipt
            .journal
            .decode()
            .map_err(|e| ProverError::ZkProofInvalid(format!("journal decode: {e}")))?;
        if j_program_hash != proof.program_hash.0 {
            return Err(ProverError::ProgramHashMismatch {
                expected: proof.program_hash.to_hex(),
                actual: Bytes32(j_program_hash).to_hex(),
            });
        }
        if j_pi_hash != sha256(&proof.public_input).0 {
            return Err(ProverError::PublicInputMismatch);
        }
        if j_output != proof.public_output.0 {
            return Err(ProverError::PublicOutputMismatch);
        }
        let msg = signing_message(&proof.proof, &proof.public_input);
        if !bls::bls_verify(&proof.node_pubkey, &msg, &proof.node_signature) {
            return Err(ProverError::NodeSignatureInvalid);
        }
        if trusted_roots.is_empty() {
            return Err(ProverError::UntrustedRoot(
                "no trusted roots provided".into(),
            ));
        }
        chain.verify_block(&proof.chia_block, DEFAULT_FRESHNESS_WINDOW_SECS)?;
        Ok(())
    }
}