newton-core 0.4.16

/// SP1 attestation proof generation
pub mod attestation;

/// RISC0 module
pub mod risc0;

/// SP1 module
pub mod sp1;

use alloy::primitives::FixedBytes;
use risc0_zkvm::{Digest, Receipt};
use serde::{Deserialize, Serialize};
use sp1_sdk::{HashableKey, Prover, ProverClient, SP1ProofWithPublicValues, SP1VerifyingKey};
use std::path::PathBuf;

use crate::{
    generated::rego_verifier::IRegoVerifier::RegoContext,
    newton_prover_task_manager::INewtonProverTaskManager::{Task, TaskResponse},
};

/// SP1 Rego ELF binary
pub const SP1_REGO_ELF: &[u8] = include_bytes!("./elf/sp1-rego");

/// SP1 Attestation verification ELF binary
pub const SP1_ATTESTATION_ELF: &[u8] = include_bytes!("./elf/sp1-attestation");

/// Enum representing the available zkvm types
#[derive(strum_macros::Display, Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[strum(serialize_all = "lowercase")]
pub enum Zkvm {
    /// Sp1 zkvm type
    Sp1,
    /// Risc0 zkvm type
    Risc0,
}

/// Enum representing the available proof systems
#[derive(strum_macros::Display, Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[strum(serialize_all = "lowercase")]
pub enum ProofSystem {
    /// Groth16 proof system
    Groth16,
    /// Plonk proof system
    Plonk,
}

/// Enum representing the available proof structs
#[derive(Clone, Serialize, Deserialize)]
pub enum Proof {
    /// Sp1 proof
    Sp1 {
        /// Journal   
        journal: Vec<u8>,
        /// Verifying key
        vk: SP1VerifyingKey,
        /// Proof
        proof: SP1ProofWithPublicValues,
        /// Request ID
        request_id: FixedBytes<32>,
    },
    /// Risc0 proof
    Risc0 {
        /// Receipt
        receipt: Receipt,
        /// Image ID
        image_id: Digest,
        /// Seal
        seal: Vec<u8>,
    },
}

impl std::fmt::Debug for Proof {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Proof::Sp1 {
                journal,
                vk,
                proof,
                request_id,
            } => write!(
                f,
                "Proof::Sp1 {{ journal: {:?}, vk: {:?}, proof bytes: {:?}, request_id: {:?} }}",
                crate::hex!(journal.as_slice()),
                vk.bytes32().to_string().as_str(),
                crate::hex!(proof.bytes()),
                crate::hex!(request_id)
            ),
            Proof::Risc0 {
                receipt,
                image_id,
                seal,
            } => write!(
                f,
                "Proof::Risc0 {{ journal: {:?}, image_id: {:?}, seal: {:?} }}",
                crate::hex!(receipt.journal.bytes.as_slice()),
                crate::hex!(image_id),
                crate::hex!(seal.as_slice())
            ),
        }
    }
}

/// Proves the policy evaluation result using the specified zkvm
/// # Arguments
/// * `circuit` - The circuit to prove
/// * `zkvm` - The zkvm to use
/// * `task` - The task to prove
/// * `task_response` - The task response to prove
/// * `entrypoint` - The entrypoint to prove
/// # Returns
/// A `Result` containing the proof, or an error if the proof fails
pub async fn prove(
    circuit: Vec<u8>,
    zkvm: Zkvm,
    task: Task,
    task_response: TaskResponse,
    entrypoint: String,
) -> eyre::Result<Proof> {
    match zkvm {
        Zkvm::Sp1 => sp1::prove(circuit, task, task_response, entrypoint).await,
        Zkvm::Risc0 => risc0::prove(circuit, task, task_response, entrypoint).await,
    }
}

/// Verifies the zk proof
/// # Arguments
/// * `proof_bytes` - The proof to verify
/// * `proof_type` - The proof type
pub async fn verify_raw_proof(proof_bytes: Vec<u8>, vkey_bytes: Vec<u8>, proof_type: Zkvm) -> eyre::Result<bool> {
    let result = match proof_type {
        Zkvm::Sp1 => {
            let vk = serde_json::from_slice::<SP1VerifyingKey>(&vkey_bytes)?;
            let proof = serde_json::from_slice::<SP1ProofWithPublicValues>(&proof_bytes)?;
            ProverClient::from_env()
                .await
                .verify(&proof, &vk, None)
                .map_err(|e| eyre::eyre!("SP1 verification failed: {}", e))
        }
        Zkvm::Risc0 => {
            let receipt = serde_json::from_slice::<Receipt>(&proof_bytes)?;
            let image_id = serde_json::from_slice::<Digest>(&vkey_bytes)?;
            receipt
                .verify(image_id)
                .map_err(|e| eyre::eyre!("Risc0 verification failed: {}", e))
        }
    };
    Ok(result.is_ok())
}