diff --git a/Cargo.lock b/Cargo.lock index 0d555bf..6e3d8fc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5273,6 +5273,7 @@ name = "void-proof" version = "0.1.0" dependencies = [ "risc0-zkvm", + "serde", "serde_json", "thiserror", "tracing", diff --git a/crates/proof/Cargo.toml b/crates/proof/Cargo.toml index 05b6195..f9b226a 100644 --- a/crates/proof/Cargo.toml +++ b/crates/proof/Cargo.toml @@ -9,6 +9,7 @@ repository.workspace = true [dependencies] risc0-zkvm.workspace = true +serde.workspace = true serde_json.workspace = true thiserror = { workspace = true, optional = true } tracing = { workspace = true, optional = true } diff --git a/crates/proof/src/guest.rs b/crates/proof/src/guest.rs index bc2593e..d9b6830 100644 --- a/crates/proof/src/guest.rs +++ b/crates/proof/src/guest.rs @@ -1,20 +1,17 @@ +use crate::types::StateCommit; use risc0_zkvm::{guest, sha::Sha256}; use void_types::Block; -pub struct Output { - pub block_hash: [u8; 32], - pub block_height: u64, - pub pre_digest: D, - pub post_digest: D, -} +pub mod recursive; pub fn proof(state_transition_function: F) where W: From>, F: Fn(&Block, &mut W), D: for<'a> From<&'a W>, - Vec: From, + Vec: for<'a> From<&'a D>, { + // Read block from guest env and deserialize let block_bytes_len: u64 = guest::env::read(); let mut block_bytes = vec![0u8; block_bytes_len as usize]; guest::env::read_slice(&mut block_bytes); @@ -23,6 +20,7 @@ where let block = Block::from_bytes(&block_bytes).expect("Block failed to deserialize"); + // Read witness from guest env and deserialize let witness_bytes_len: u64 = guest::env::read(); let mut witness_bytes = vec![0u8; witness_bytes_len as usize]; guest::env::read_slice(&mut witness_bytes); @@ -35,52 +33,53 @@ where let post_digest = D::from(&witness); - let output = Output { + let output = StateCommit { block_hash: (*block_hash).into(), - // block_height: block.height, - block_height: 0, + parent_hash: block.parent_hash, + block_height: block.height, pre_digest, post_digest, }; - let output = to_output_bytes(output); + let mut bytes = Vec::new(); + to_output_bytes(&mut bytes, &output); - guest::env::commit_slice(&output); + guest::env::commit_slice(&bytes); } -pub fn to_output_bytes(output: Output) -> Vec +pub fn to_output_bytes(bytes: &mut Vec, output: &StateCommit) where - Vec: From, + Vec: for<'a> From<&'a D>, { - let mut bytes = Vec::new(); bytes.extend(&output.block_hash); + bytes.extend(&output.parent_hash); bytes.extend(&output.block_height.to_be_bytes()); - let pre_bytes: Vec = output.pre_digest.into(); + let pre_bytes: Vec = (&output.pre_digest).into(); bytes.extend(&(pre_bytes.len() as u64).to_be_bytes()); bytes.extend(&pre_bytes); - let post_bytes: Vec = output.post_digest.into(); + let post_bytes: Vec = (&output.post_digest).into(); bytes.extend(&(post_bytes.len() as u64).to_be_bytes()); bytes.extend(&post_bytes); - bytes } -pub fn from_output_bytes(bytes: &[u8]) -> Option> +pub fn from_output_bytes(bytes: &[u8]) -> Option> where D: for<'a> TryFrom<&'a [u8]>, { - if bytes.len() < 32 + 8 + 8 { + if bytes.len() < 32 + 32 + 8 + 8 { return None; } let block_hash: [u8; 32] = bytes[0..32].try_into().expect("Converting to sized array"); + let parent_hash: [u8; 32] = bytes[32..64].try_into().expect("Converting to sized array"); let block_height = - u64::from_be_bytes(bytes[32..40].try_into().expect("Converting to sized array")); + u64::from_be_bytes(bytes[64..72].try_into().expect("Converting to sized array")); let pre_len = - u64::from_be_bytes(bytes[40..48].try_into().expect("Converting to sized array")) as usize; - if bytes.len() < 48 + pre_len + 8 { + u64::from_be_bytes(bytes[72..80].try_into().expect("Converting to sized array")) as usize; + if bytes.len() < 80 + pre_len + 8 { return None; } - let pre_bytes = &bytes[48..48 + pre_len]; - let post_len_start = 48 + pre_len; + let pre_bytes = &bytes[80..80 + pre_len]; + let post_len_start = 80 + pre_len; let post_len = u64::from_be_bytes( bytes[post_len_start..post_len_start + 8] .try_into() @@ -90,8 +89,9 @@ where return None; } let post_bytes = &bytes[post_len_start + 8..post_len_start + 8 + post_len]; - Some(Output { + Some(StateCommit { block_hash, + parent_hash, block_height, // TODO: check that it is safe to lose the error pre_digest: D::try_from(pre_bytes).ok()?, diff --git a/crates/proof/src/guest/recursive.rs b/crates/proof/src/guest/recursive.rs new file mode 100644 index 0000000..85f5810 --- /dev/null +++ b/crates/proof/src/guest/recursive.rs @@ -0,0 +1,162 @@ +use crate::types::{RecursiveCommit, RecursiveType, StateCommit}; +use risc0_zkvm::guest::env; +use std::fmt::Debug; + +pub fn proof(state_image_id: [u32; 8]) +where + D: for<'a> TryFrom<&'a [u8]> + Debug + Default + PartialEq, + Vec: for<'a> From<&'a D>, +{ + let state_commits_count = env::read(); + let mut state_commits = Vec::with_capacity(state_commits_count); + + for _ in 0..state_commits_count { + let state_commit_len = env::read(); + + let mut state_commit_bytes = vec![0u8; state_commit_len]; + env::read_slice(&mut state_commit_bytes); + + let state_commit: StateCommit = super::from_output_bytes(&state_commit_bytes) + .expect("Failed to deserialize state commit"); + + state_commits.push(state_commit); + } + + let recursive_type: RecursiveType = env::read(); + + match recursive_type { + RecursiveType::Init => { + let image_id: [u32; 8] = env::read(); + + let first_state_commit = state_commits + .first() + .expect("Must have at least one state commit"); + + assert_eq!(first_state_commit.parent_hash, [0u8; 32]); + assert_eq!(first_state_commit.pre_digest, D::default()); + assert_eq!(first_state_commit.block_height, 0); + + let mut bytes = Vec::new(); + super::to_output_bytes(&mut bytes, first_state_commit); + + env::verify(state_image_id, &bytes).expect("init state proof failed to verify"); + + verify_state_steps(state_image_id, &state_commits); + + let recursive_commit = RecursiveCommit { + this_image_id: image_id, + state_commit: state_commits + .pop() + .expect("Must have at least one state commit"), + }; + + let mut bytes = Vec::new(); + to_output_bytes(&mut bytes, &recursive_commit); + env::commit_slice(&bytes[..]); + } + RecursiveType::Step => { + let recursive_commit_len = env::read(); + + let mut recursive_commit_bytes = vec![0u8; recursive_commit_len]; + env::read_slice(&mut recursive_commit_bytes); + + let recursive_commit: RecursiveCommit = from_output_bytes(&recursive_commit_bytes) + .expect("Failed to deserialize state commit"); + + let first_state_commit = state_commits + .first() + .expect("Must have at least one state commit"); + + verify_step( + state_image_id, + &recursive_commit.state_commit, + first_state_commit, + ); + verify_state_steps(state_image_id, &state_commits); + + let recursive_commit = RecursiveCommit { + this_image_id: recursive_commit.this_image_id, + state_commit: state_commits + .pop() + .expect("Must have at least one state commit"), + }; + + let mut bytes = Vec::new(); + to_output_bytes(&mut bytes, &recursive_commit); + env::commit_slice(&bytes[..]); + } + } +} + +// Verify a single step in the state commit chain. +fn verify_step(image_id: [u32; 8], parent: &StateCommit, commit: &StateCommit) +where + D: Debug + PartialEq, + Vec: for<'a> From<&'a D>, +{ + assert_eq!(commit.parent_hash, parent.block_hash); + assert_eq!(commit.pre_digest, parent.post_digest); + assert_eq!( + commit.block_height, + parent.block_height.checked_add(1).unwrap() + ); + + let mut bytes = Vec::new(); + super::to_output_bytes(&mut bytes, commit); + + env::verify(image_id, &bytes).expect("step state proof failed to verify"); +} + +// Verify a state commit chain. +fn verify_state_steps(image_id: [u32; 8], state_commits: &[StateCommit]) +where + D: Debug + PartialEq, + Vec: for<'a> From<&'a D>, +{ + for (parent, commit) in state_commits.iter().zip(state_commits.iter().skip(1)) { + verify_step(image_id, parent, commit); + } +} + +pub fn to_output_bytes(bytes: &mut Vec, commit: &RecursiveCommit) +where + Vec: for<'a> From<&'a D>, +{ + let mut image_id: [u8; 32] = [0u8; 32]; + for (input, output) in commit + .this_image_id + .iter() + .map(|e| e.to_be_bytes()) + .flatten() + .zip(image_id.iter_mut()) + { + *output = input; + } + bytes.extend(image_id); + + super::to_output_bytes(bytes, &commit.state_commit); +} + +pub fn from_output_bytes(bytes: &[u8]) -> Option> +where + D: for<'a> TryFrom<&'a [u8]>, +{ + if bytes.len() < 32 { + return None; + } + + let mut image_id: [u32; 8] = [0u32; 8]; + for (input, output) in bytes[0..32] + .chunks_exact(std::mem::size_of::()) + .zip(image_id.iter_mut()) + { + *output = u32::from_be_bytes(input.try_into().expect("The sizes match")); + } + + let state_commit = super::from_output_bytes(&bytes[32..])?; + + Some(RecursiveCommit { + this_image_id: image_id, + state_commit, + }) +} diff --git a/crates/proof/src/host.rs b/crates/proof/src/host.rs index dbff356..6a3b1e2 100644 --- a/crates/proof/src/host.rs +++ b/crates/proof/src/host.rs @@ -3,6 +3,8 @@ use thiserror::Error; use tracing::{debug, info, instrument, warn}; use void_types::Block; +pub mod recursive; + #[derive(Debug, Error)] pub enum ProverError { #[error("Failed to write to environment: {0}")] diff --git a/crates/proof/src/host/recursive.rs b/crates/proof/src/host/recursive.rs new file mode 100644 index 0000000..80141a7 --- /dev/null +++ b/crates/proof/src/host/recursive.rs @@ -0,0 +1,139 @@ +use crate::types::{RecursiveCommit, RecursiveType}; +use risc0_zkvm::{ExecutorEnv, Receipt, default_prover}; +use thiserror::Error; +use tracing::{debug, info}; + +pub type StateReceipt = Receipt; +pub type RecursiveReceipt = Receipt; + +pub struct Input { + pub recursive: RecursiveArgs, + pub state_receipts: Vec, +} + +pub enum RecursiveArgs { + Init(Init), + Step(Box), +} + +pub enum RecursiveInput { + Init(Init), + Step(Step), +} + +pub struct Init { + pub this_image_id: [u32; 8], +} + +pub struct Step { + pub previous_commit: RecursiveCommit, +} + +#[derive(Error, Debug)] +pub enum RecursiveProofError { + #[error("Error writing to environment: {0}")] + Write(String), + #[error("Failed to build environment: {0}")] + Build(String), + #[error("Failed to prove: {0}")] + Proof(String), + #[error("Failed to serialize receipt: {0}")] + Serialization(#[from] serde_json::Error), +} + +pub fn proof(input: Input, elf: &[u8]) -> Result, RecursiveProofError> +where + D: for<'a> TryFrom<&'a [u8]>, + Vec: From, +{ + info!("starting zkVM proof generation"); + let mut env = ExecutorEnv::builder(); + + // Write the number of state receipts to env. + env.write(&(input.state_receipts.len() as u64)) + .inspect_err(|e| { + tracing::error!( + "Failed to write number of state receipts to environment: {}", + e + ) + }) + .map_err(|e| RecursiveProofError::Write(e.to_string()))?; + + for state_receipt in input.state_receipts { + // Write state receipt bytes to env. + env.write(&(state_receipt.journal.bytes.len() as u64)) + .inspect_err(|e| { + tracing::error!( + "Failed to write state commit bytes length to environment: {}", + e + ) + }) + .map_err(|e| RecursiveProofError::Write(e.to_string()))?; + env.write_slice(&state_receipt.journal.bytes); + + // Add state receipt as assumption. + env.add_assumption(state_receipt); + } + + match input.recursive { + RecursiveArgs::Init(init) => { + // Write recursion type to env. + env.write(&RecursiveType::Init) + .inspect_err(|e| { + tracing::error!( + "Failed to write recursion type (Init) to environment: {}", + e + ) + }) + .map_err(|e| RecursiveProofError::Write(e.to_string()))?; + + // Write image id to env. + env.write(&init.this_image_id) + .inspect_err(|e| tracing::error!("Failed to write image id to environment: {}", e)) + .map_err(|e| RecursiveProofError::Write(e.to_string()))?; + } + RecursiveArgs::Step(recursive_receipt) => { + // Write recursion type to env. + env.write(&RecursiveType::Step) + .inspect_err(|e| { + tracing::error!( + "Failed to write recursion type (Step) to environment: {}", + e + ) + }) + .map_err(|e| RecursiveProofError::Write(e.to_string()))?; + + // Write recursive receipt bytes to env. + env.write(&(recursive_receipt.journal.bytes.len() as u64)) + .inspect_err(|e| { + tracing::error!( + "Failed to write recursive receipt length to environment: {}", + e + ) + }) + .map_err(|e| RecursiveProofError::Write(e.to_string()))?; + env.write_slice(&recursive_receipt.journal.bytes); + + // Add recursive receipt as assumption. + env.add_assumption(*recursive_receipt); + } + }; + + let env = env + .build() + .inspect_err(|e| tracing::error!("Failed to build executor environment: {}", e)) + .map_err(|e| RecursiveProofError::Build(e.to_string()))?; + + debug!("executor environment built"); + let prover = default_prover(); + + debug!("running prover"); + let receipt = prover + .prove(env, elf) + .inspect_err(|e| tracing::error!("Failed to prove: {}", e)) + .map_err(|e| RecursiveProofError::Proof(e.to_string()))? + .receipt; + + info!("proof generation completed"); + Ok(serde_json::to_vec(&receipt)?) +} diff --git a/crates/proof/src/lib.rs b/crates/proof/src/lib.rs index 170aff8..7295df8 100644 --- a/crates/proof/src/lib.rs +++ b/crates/proof/src/lib.rs @@ -2,3 +2,5 @@ pub mod guest; #[cfg(feature = "host")] pub mod host; + +pub mod types; diff --git a/crates/proof/src/types.rs b/crates/proof/src/types.rs new file mode 100644 index 0000000..1da1600 --- /dev/null +++ b/crates/proof/src/types.rs @@ -0,0 +1,18 @@ +pub struct RecursiveCommit { + pub this_image_id: [u32; 8], + pub state_commit: StateCommit, +} + +#[derive(serde::Serialize, serde::Deserialize, Debug, Clone, Copy, PartialEq)] +pub enum RecursiveType { + Init, + Step, +} + +pub struct StateCommit { + pub block_hash: [u8; 32], + pub parent_hash: [u8; 32], + pub block_height: u64, + pub pre_digest: D, + pub post_digest: D, +}