From 43bd2a831485eb67c42a0842f9f8f32185cb3743 Mon Sep 17 00:00:00 2001 From: Rostyslav Tyshko Date: Fri, 7 Feb 2025 15:16:31 -0500 Subject: [PATCH] fix error handing in zkvm module --- zkvm/src/lib.rs | 99 ++++++++++++++++++++++++------------------------- 1 file changed, 49 insertions(+), 50 deletions(-) diff --git a/zkvm/src/lib.rs b/zkvm/src/lib.rs index aff0ced..19fa8be 100644 --- a/zkvm/src/lib.rs +++ b/zkvm/src/lib.rs @@ -56,7 +56,7 @@ pub fn prove_mint_utxo( let digest: UTXOPayload = receipt.journal.decode()?; - Ok((UTXO::create_utxo_from_payload(digest), receipt)) + Ok((UTXO::create_utxo_from_payload(digest).map_err(ExecutionFailureKind::write_error)?, receipt)) } pub fn prove_send_utxo( @@ -95,8 +95,8 @@ pub fn prove_send_utxo( Ok(( digest .into_iter() - .map(|(payload, addr)| (UTXO::create_utxo_from_payload(payload), addr)) - .collect(), + .map(|(payload, addr)| (UTXO::create_utxo_from_payload(payload).map(|sel| (sel, addr)))) + .collect::>>().map_err(ExecutionFailureKind::write_error)?, receipt, )) } @@ -144,12 +144,12 @@ pub fn prove_send_utxo_multiple_assets_one_receiver( .0 .into_iter() .map(|payload| UTXO::create_utxo_from_payload(payload)) - .collect(), + .collect::>>().map_err(ExecutionFailureKind::write_error)?, digest .1 .into_iter() .map(|payload| UTXO::create_utxo_from_payload(payload)) - .collect(), + .collect::>>().map_err(ExecutionFailureKind::write_error)?, receipt, )) } @@ -170,7 +170,7 @@ pub fn prove_send_utxo_shielded( asset: vec![], amount, privacy_flag: true, - }); + }).map_err(ExecutionFailureKind::write_error)?; let utxo_payload = temp_utxo_to_spend.into_payload(); let mut builder = ExecutorEnv::builder(); @@ -198,8 +198,8 @@ pub fn prove_send_utxo_shielded( Ok(( digest .into_iter() - .map(|(payload, addr)| (UTXO::create_utxo_from_payload(payload), addr)) - .collect(), + .map(|(payload, addr)| (UTXO::create_utxo_from_payload(payload).map(|sel| (sel, addr)))) + .collect::>>().map_err(ExecutionFailureKind::write_error)?, receipt, )) } @@ -280,24 +280,24 @@ pub fn prove_mint_utxo_multiple_assets( digest .into_iter() .map(UTXO::create_utxo_from_payload) - .collect(), + .collect::>>().map_err(ExecutionFailureKind::write_error)?, receipt, )) } -pub fn execute_mint_utxo(amount_to_mint: u128, owner: AccountAddress) -> UTXO { +pub fn execute_mint_utxo(amount_to_mint: u128, owner: AccountAddress) -> anyhow::Result { let mut builder = ExecutorEnv::builder(); - builder.write(&amount_to_mint).unwrap(); - builder.write(&owner).unwrap(); + builder.write(&amount_to_mint)?; + builder.write(&owner)?; - let env = builder.build().unwrap(); + let env = builder.build()?; let executor = default_executor(); - let receipt = executor.execute(env, test_methods::MINT_UTXO_ELF).unwrap(); + let receipt = executor.execute(env, test_methods::MINT_UTXO_ELF)?; - let digest: UTXOPayload = receipt.journal.decode().unwrap(); + let digest: UTXOPayload = receipt.journal.decode()?; UTXO::create_utxo_from_payload(digest) } @@ -305,76 +305,75 @@ pub fn execute_mint_utxo(amount_to_mint: u128, owner: AccountAddress) -> UTXO { pub fn execute_send_utxo( spent_utxo: UTXO, owners_parts: Vec<(u128, AccountAddress)>, -) -> (UTXO, Vec<(UTXO, AccountAddress)>) { +) -> anyhow::Result<(UTXO, Vec<(UTXO, AccountAddress)>)> { let mut builder = ExecutorEnv::builder(); let utxo_payload = spent_utxo.into_payload(); - builder.write(&utxo_payload).unwrap(); - builder.write(&owners_parts).unwrap(); + builder.write(&utxo_payload)?; + builder.write(&owners_parts)?; - let env = builder.build().unwrap(); + let env = builder.build()?; let executor = default_executor(); - let receipt = executor.execute(env, test_methods::SEND_UTXO_ELF).unwrap(); + let receipt = executor.execute(env, test_methods::SEND_UTXO_ELF)?; let digest: (UTXOPayload, Vec<(UTXOPayload, AccountAddress)>) = - receipt.journal.decode().unwrap(); + receipt.journal.decode()?; - ( - UTXO::create_utxo_from_payload(digest.0), + Ok(( + UTXO::create_utxo_from_payload(digest.0)?, digest .1 .into_iter() - .map(|(payload, addr)| (UTXO::create_utxo_from_payload(payload), addr)) - .collect(), - ) + .map(|(payload, addr)| (UTXO::create_utxo_from_payload(payload).map(|sel| (sel, addr)))) + .collect::>>().map_err(ExecutionFailureKind::write_error)?, + )) } -pub fn prove(input_vec: Vec, elf: &[u8]) -> (u64, Receipt) { +pub fn prove(input_vec: Vec, elf: &[u8]) -> anyhow::Result<(u64, Receipt)> { let mut builder = ExecutorEnv::builder(); for input in input_vec { - builder.write(&input).unwrap(); + builder.write(&input)?; } - let env = builder.build().unwrap(); + let env = builder.build()?; let prover = default_prover(); - let receipt = prover.prove(env, elf).unwrap().receipt; + let receipt = prover.prove(env, elf)?.receipt; - let digest = receipt.journal.decode().unwrap(); - (digest, receipt) + let digest = receipt.journal.decode()?; + Ok((digest, receipt)) } // This only executes the program and does not generate a receipt. pub fn execute serde::Deserialize<'de>>( input_vec: Vec, elf: &[u8], -) -> T { +) -> anyhow::Result { let mut builder = ExecutorEnv::builder(); for input in input_vec { - builder.write(&input).unwrap(); + builder.write(&input)?; } - let env = builder.build().unwrap(); + let env = builder.build()?; let exec = default_executor(); - let session = exec.execute(env, elf).unwrap(); + let session = exec.execute(env, elf)?; // We read the result committed to the journal by the guest code. - let result: T = session.journal.decode().unwrap(); + let result: T = session.journal.decode()?; - result + Ok(result) } -pub fn verify(receipt: Receipt, image_id: impl Into) { - receipt - .verify(image_id) - .expect("receipt verification failed"); +pub fn verify(receipt: Receipt, image_id: impl Into) -> anyhow::Result<()> { + Ok(receipt + .verify(image_id)?) } #[cfg(test)] @@ -389,7 +388,7 @@ mod tests { let message = 1; let message_2 = 2; - let (digest, receipt) = prove(vec![message, message_2], SUMMATION_ELF); + let (digest, receipt) = prove(vec![message, message_2], SUMMATION_ELF).unwrap(); verify(receipt, SUMMATION_ID); assert_eq!(digest, message + message_2); @@ -400,7 +399,7 @@ mod tests { let message = 123476; let message_2 = 2342384; - let (digest, receipt) = prove(vec![message, message_2], SUMMATION_ELF); + let (digest, receipt) = prove(vec![message, message_2], SUMMATION_ELF).unwrap(); verify(receipt, SUMMATION_ID); assert_eq!(digest, message + message_2); @@ -411,7 +410,7 @@ mod tests { let message = 1; let message_2 = 2; - let (digest, receipt) = prove(vec![message, message_2], MULTIPLICATION_ELF); + let (digest, receipt) = prove(vec![message, message_2], MULTIPLICATION_ELF).unwrap(); verify(receipt, MULTIPLICATION_ID); assert_eq!(digest, message * message_2); @@ -422,7 +421,7 @@ mod tests { let message = 3498; let message_2 = 438563; - let (digest, receipt) = prove(vec![message, message_2], MULTIPLICATION_ELF); + let (digest, receipt) = prove(vec![message, message_2], MULTIPLICATION_ELF).unwrap(); verify(receipt, MULTIPLICATION_ID); assert_eq!(digest, message * message_2); @@ -433,7 +432,7 @@ mod tests { let message: u64 = 1; let message_2: u64 = 2; - let result = execute(vec![message, message_2], SUMMATION_ELF); + let result = execute(vec![message, message_2], SUMMATION_ELF).unwrap(); assert_eq!(result, message + message_2); } @@ -442,7 +441,7 @@ mod tests { let message: u64 = 123476; let message_2: u64 = 2342384; - let result = execute(vec![message, message_2], SUMMATION_ELF); + let result = execute(vec![message, message_2], SUMMATION_ELF).unwrap(); assert_eq!(result, message + message_2); } @@ -451,7 +450,7 @@ mod tests { let message: u128 = 1; let message_2: u128 = 2; - let result = execute(vec![message, message_2], BIG_CALCULATION_ELF); + let result = execute(vec![message, message_2], BIG_CALCULATION_ELF).unwrap(); assert_eq!(result, big_calculation(message, message_2)); } @@ -460,7 +459,7 @@ mod tests { let message: u128 = 20; let message_2: u128 = 10; - let result = execute(vec![message, message_2], BIG_CALCULATION_ELF); + let result = execute(vec![message, message_2], BIG_CALCULATION_ELF).unwrap(); assert_eq!(result, big_calculation(message, message_2)); }