diff --git a/nssa/core/src/program.rs b/nssa/core/src/program.rs index 3be22d4..7f6038f 100644 --- a/nssa/core/src/program.rs +++ b/nssa/core/src/program.rs @@ -26,23 +26,32 @@ pub struct ChainedCall { #[derive(Serialize, Deserialize, Clone)] #[cfg_attr(any(feature = "host", test), derive(Debug, PartialEq, Eq))] pub struct ProgramOutput { + pub instruction_data: InstructionData, pub pre_states: Vec, pub post_states: Vec, pub chained_call: Option, } -pub fn read_nssa_inputs() -> ProgramInput { +pub fn read_nssa_inputs() -> (ProgramInput, InstructionData) { let pre_states: Vec = env::read(); let instruction_words: InstructionData = env::read(); let instruction = T::deserialize(&mut Deserializer::new(instruction_words.as_ref())).unwrap(); - ProgramInput { - pre_states, - instruction, - } + ( + ProgramInput { + pre_states, + instruction, + }, + instruction_words, + ) } -pub fn write_nssa_outputs(pre_states: Vec, post_states: Vec) { +pub fn write_nssa_outputs( + instruction_data: InstructionData, + pre_states: Vec, + post_states: Vec, +) { let output = ProgramOutput { + instruction_data, pre_states, post_states, chained_call: None, @@ -51,11 +60,13 @@ pub fn write_nssa_outputs(pre_states: Vec, post_states: Vec } pub fn write_nssa_outputs_with_chained_call( + instruction_data: InstructionData, pre_states: Vec, post_states: Vec, chained_call: Option, ) { let output = ProgramOutput { + instruction_data, pre_states, post_states, chained_call, diff --git a/nssa/program_methods/guest/src/bin/authenticated_transfer.rs b/nssa/program_methods/guest/src/bin/authenticated_transfer.rs index df8a38e..2e05492 100644 --- a/nssa/program_methods/guest/src/bin/authenticated_transfer.rs +++ b/nssa/program_methods/guest/src/bin/authenticated_transfer.rs @@ -2,37 +2,42 @@ use nssa_core::{ account::{Account, AccountWithMetadata}, program::{ProgramInput, read_nssa_inputs, write_nssa_outputs}, }; +use risc0_zkvm::serde::to_vec; /// Initializes a default account under the ownership of this program. /// This is achieved by a noop. -fn initialize_account(pre_state: AccountWithMetadata) { +fn initialize_account(pre_state: AccountWithMetadata) -> (AccountWithMetadata, Account) { let account_to_claim = pre_state.account.clone(); let is_authorized = pre_state.is_authorized; // Continue only if the account to claim has default values if account_to_claim != Account::default() { - return; + panic!("Invalid input"); } // Continue only if the owner authorized this operation if !is_authorized { - return; + panic!("Invalid input"); } // Noop will result in account being claimed for this program - write_nssa_outputs(vec![pre_state], vec![account_to_claim]); + (pre_state, account_to_claim) } /// Transfers `balance_to_move` native balance from `sender` to `recipient`. -fn transfer(sender: AccountWithMetadata, recipient: AccountWithMetadata, balance_to_move: u128) { +fn transfer( + sender: AccountWithMetadata, + recipient: AccountWithMetadata, + balance_to_move: u128, +) -> (Vec, Vec) { // Continue only if the sender has authorized this operation if !sender.is_authorized { - return; + panic!("Invalid input"); } // Continue only if the sender has enough balance if sender.account.balance < balance_to_move { - return; + panic!("Invalid input"); } // Create accounts post states, with updated balances @@ -40,23 +45,30 @@ fn transfer(sender: AccountWithMetadata, recipient: AccountWithMetadata, balance let mut recipient_post = recipient.account.clone(); sender_post.balance -= balance_to_move; recipient_post.balance += balance_to_move; - - write_nssa_outputs(vec![sender, recipient], vec![sender_post, recipient_post]); + (vec![sender, recipient], vec![sender_post, recipient_post]) } /// A transfer of balance program. /// To be used both in public and private contexts. fn main() { // Read input accounts. - let ProgramInput { - pre_states, - instruction: balance_to_move, - } = read_nssa_inputs(); + let ( + ProgramInput { + pre_states, + instruction: balance_to_move, + }, + instruction_words, + ) = read_nssa_inputs(); match (pre_states.as_slice(), balance_to_move) { - ([account_to_claim], 0) => initialize_account(account_to_claim.clone()), + ([account_to_claim], 0) => { + let (pre, post) = initialize_account(account_to_claim.clone()); + write_nssa_outputs(instruction_words, vec![pre], vec![post]); + } ([sender, recipient], balance_to_move) => { - transfer(sender.clone(), recipient.clone(), balance_to_move) + let (pre_states, post_states) = + transfer(sender.clone(), recipient.clone(), balance_to_move); + write_nssa_outputs(instruction_words, pre_states, post_states); } _ => panic!("invalid params"), } diff --git a/nssa/program_methods/guest/src/bin/pinata.rs b/nssa/program_methods/guest/src/bin/pinata.rs index fbea167..d2cd80d 100644 --- a/nssa/program_methods/guest/src/bin/pinata.rs +++ b/nssa/program_methods/guest/src/bin/pinata.rs @@ -1,5 +1,8 @@ use nssa_core::program::{ProgramInput, read_nssa_inputs, write_nssa_outputs}; -use risc0_zkvm::sha::{Impl, Sha256}; +use risc0_zkvm::{ + serde::to_vec, + sha::{Impl, Sha256}, +}; const PRIZE: u128 = 150; @@ -44,10 +47,13 @@ impl Challenge { fn main() { // Read input accounts. // It is expected to receive only two accounts: [pinata_account, winner_account] - let ProgramInput { - pre_states, - instruction: solution, - } = read_nssa_inputs::(); + let ( + ProgramInput { + pre_states, + instruction: solution, + }, + instruction_data, + ) = read_nssa_inputs::(); let [pinata, winner] = match pre_states.try_into() { Ok(array) => array, @@ -66,5 +72,9 @@ fn main() { pinata_post.data = data.next_data().to_vec(); winner_post.balance += PRIZE; - write_nssa_outputs(vec![pinata, winner], vec![pinata_post, winner_post]); + write_nssa_outputs( + to_vec(&solution).unwrap(), + vec![pinata, winner], + vec![pinata_post, winner_post], + ); } diff --git a/nssa/program_methods/guest/src/bin/privacy_preserving_circuit.rs b/nssa/program_methods/guest/src/bin/privacy_preserving_circuit.rs index 223229f..5d13d46 100644 --- a/nssa/program_methods/guest/src/bin/privacy_preserving_circuit.rs +++ b/nssa/program_methods/guest/src/bin/privacy_preserving_circuit.rs @@ -8,7 +8,7 @@ use nssa_core::{ account::{Account, AccountId, AccountWithMetadata}, compute_digest_for_path, encryption::Ciphertext, - program::{DEFAULT_PROGRAM_ID, MAX_NUMBER_CHAINED_CALLS, ProgramOutput, validate_execution}, + program::{DEFAULT_PROGRAM_ID, MAX_NUMBER_CHAINED_CALLS, validate_execution}, }; fn main() { @@ -24,12 +24,30 @@ fn main() { let mut pre_states: Vec = Vec::new(); let mut state_diff: HashMap = HashMap::new(); - let mut program_output = program_outputs[0].clone(); + let num_calls = program_outputs.len(); + if num_calls > MAX_NUMBER_CHAINED_CALLS { + panic!("Max deapth is exceeded"); + } + + if program_outputs[num_calls - 1].chained_call.is_some() { + panic!("Call stack is incomplete"); + } + + for i in 0..(program_outputs.len() - 1) { + let Some(chained_call) = program_outputs[i].chained_call.clone() else { + panic!("Expected chained call"); + }; + + // Check that instruction data in caller is the instruction data in callee + if chained_call.instruction_data != program_outputs[i + 1].instruction_data { + panic!("Invalid instruction data"); + } + } + + for program_output in program_outputs { + let mut program_output = program_output.clone(); - for _i in 0..MAX_NUMBER_CHAINED_CALLS { // Check that `program_output` is consistent with the execution of the corresponding program. - // TODO: Program output should contain the instruction data to verify the chain of call si - // performed correctly. env::verify(program_id, &to_vec(&program_output).unwrap()).unwrap(); // Check that the program is well behaved. @@ -54,7 +72,11 @@ fn main() { .iter() .zip(&program_output.post_states) { - if !state_diff.contains_key(&pre.account_id) { + if let Some(account_pre) = state_diff.get(&pre.account_id) { + if account_pre != &pre.account { + panic!("Invalid input"); + } + } else { pre_states.push(pre.clone()); } state_diff.insert(pre.account_id.clone(), post.clone()); @@ -62,29 +84,6 @@ fn main() { if let Some(next_chained_call) = &program_output.chained_call { program_id = next_chained_call.program_id; - - // // Build post states with metadata for next call - // let mut post_states_with_metadata = Vec::new(); - // for (pre, post) in program_output - // .pre_states - // .iter() - // .zip(program_output.post_states) - // { - // let mut post_with_metadata = pre.clone(); - // post_with_metadata.account = post.clone(); - // post_states_with_metadata.push(post_with_metadata); - // } - - // input_pre_states = next_chained_call - // .account_indices - // .iter() - // .map(|&i| { - // post_states_with_metadata - // .get(i) - // .ok_or_else(|| NssaError::InvalidInput("Invalid account indices".into())) - // .cloned() - // }) - // .collect::, NssaError>>()?; } else { break; }; diff --git a/nssa/program_methods/guest/src/bin/token.rs b/nssa/program_methods/guest/src/bin/token.rs index e5680be..5be4f91 100644 --- a/nssa/program_methods/guest/src/bin/token.rs +++ b/nssa/program_methods/guest/src/bin/token.rs @@ -170,10 +170,13 @@ fn new_definition( type Instruction = [u8; 23]; fn main() { - let ProgramInput { - pre_states, - instruction, - } = read_nssa_inputs::(); + let ( + ProgramInput { + pre_states, + instruction, + }, + instruction_words, + ) = read_nssa_inputs::(); match instruction[0] { 0 => { @@ -184,7 +187,7 @@ fn main() { // Execute let post_states = new_definition(&pre_states, name, total_supply); - write_nssa_outputs(pre_states, post_states); + write_nssa_outputs(instruction_words, pre_states, post_states); } 1 => { // Parse instruction @@ -194,7 +197,7 @@ fn main() { // Execute let post_states = transfer(&pre_states, balance_to_move); - write_nssa_outputs(pre_states, post_states); + write_nssa_outputs(instruction_words, pre_states, post_states); } _ => panic!("Invalid instruction"), }; @@ -204,7 +207,7 @@ fn main() { mod tests { use nssa_core::account::{Account, AccountId, AccountWithMetadata}; - use crate::{new_definition, transfer, TOKEN_HOLDING_DATA_SIZE, TOKEN_HOLDING_TYPE}; + use crate::{TOKEN_HOLDING_DATA_SIZE, TOKEN_HOLDING_TYPE, new_definition, transfer}; #[should_panic(expected = "Invalid number of input accounts")] #[test] diff --git a/nssa/src/privacy_preserving_transaction/circuit.rs b/nssa/src/privacy_preserving_transaction/circuit.rs index 7628e78..96cf583 100644 --- a/nssa/src/privacy_preserving_transaction/circuit.rs +++ b/nssa/src/privacy_preserving_transaction/circuit.rs @@ -27,6 +27,7 @@ pub fn execute_and_prove( ) -> Result<(PrivacyPreservingCircuitOutput, Proof), NssaError> { let mut env_builder = ExecutorEnv::builder(); let mut program_outputs = Vec::new(); + for _i in 0..MAX_NUMBER_CHAINED_CALLS { let inner_receipt = execute_and_prove_program(program, pre_states, instruction_data)?; diff --git a/nssa/test_program_methods/guest/src/bin/burner.rs b/nssa/test_program_methods/guest/src/bin/burner.rs index 1ef7373..b5352a4 100644 --- a/nssa/test_program_methods/guest/src/bin/burner.rs +++ b/nssa/test_program_methods/guest/src/bin/burner.rs @@ -1,12 +1,15 @@ -use nssa_core::program::{read_nssa_inputs, write_nssa_outputs, ProgramInput}; +use nssa_core::program::{ProgramInput, read_nssa_inputs, write_nssa_outputs}; type Instruction = u128; fn main() { - let ProgramInput { - pre_states, - instruction: balance_to_burn, - } = read_nssa_inputs::(); + let ( + ProgramInput { + pre_states, + instruction: balance_to_burn, + }, + instruction_words, + ) = read_nssa_inputs::(); let [pre] = match pre_states.try_into() { Ok(array) => array, @@ -17,5 +20,5 @@ fn main() { let mut account_post = account_pre.clone(); account_post.balance -= balance_to_burn; - write_nssa_outputs(vec![pre], vec![account_post]); + write_nssa_outputs(instruction_words, vec![pre], vec![account_post]); } diff --git a/nssa/test_program_methods/guest/src/bin/chain_caller.rs b/nssa/test_program_methods/guest/src/bin/chain_caller.rs index dfd77b1..e75f17a 100644 --- a/nssa/test_program_methods/guest/src/bin/chain_caller.rs +++ b/nssa/test_program_methods/guest/src/bin/chain_caller.rs @@ -8,10 +8,13 @@ type Instruction = (u128, ProgramId); /// A program that calls another program. /// It permutes the order of the input accounts on the subsequent call fn main() { - let ProgramInput { - pre_states, - instruction: (balance, program_id), - } = read_nssa_inputs::(); + let ( + ProgramInput { + pre_states, + instruction: (balance, program_id), + }, + instruction_words, + ) = read_nssa_inputs::(); let [sender_pre, receiver_pre] = match pre_states.try_into() { Ok(array) => array, @@ -27,6 +30,7 @@ fn main() { }); write_nssa_outputs_with_chained_call( + instruction_words, vec![sender_pre.clone(), receiver_pre.clone()], vec![sender_pre.account, receiver_pre.account], chained_call, diff --git a/nssa/test_program_methods/guest/src/bin/data_changer.rs b/nssa/test_program_methods/guest/src/bin/data_changer.rs index c7d34a2..9cd8992 100644 --- a/nssa/test_program_methods/guest/src/bin/data_changer.rs +++ b/nssa/test_program_methods/guest/src/bin/data_changer.rs @@ -1,9 +1,9 @@ -use nssa_core::program::{read_nssa_inputs, write_nssa_outputs, ProgramInput}; +use nssa_core::program::{ProgramInput, read_nssa_inputs, write_nssa_outputs}; type Instruction = (); fn main() { - let ProgramInput { pre_states, .. } = read_nssa_inputs::(); + let (ProgramInput { pre_states, .. }, instruction_words) = read_nssa_inputs::(); let [pre] = match pre_states.try_into() { Ok(array) => array, @@ -14,5 +14,5 @@ fn main() { let mut account_post = account_pre.clone(); account_post.data.push(0); - write_nssa_outputs(vec![pre], vec![account_post]); + write_nssa_outputs(instruction_words, vec![pre], vec![account_post]); } diff --git a/nssa/test_program_methods/guest/src/bin/extra_output.rs b/nssa/test_program_methods/guest/src/bin/extra_output.rs index 3543d51..64fdea3 100644 --- a/nssa/test_program_methods/guest/src/bin/extra_output.rs +++ b/nssa/test_program_methods/guest/src/bin/extra_output.rs @@ -1,12 +1,12 @@ use nssa_core::{ account::Account, - program::{read_nssa_inputs, write_nssa_outputs, ProgramInput}, + program::{ProgramInput, read_nssa_inputs, write_nssa_outputs}, }; type Instruction = (); fn main() { - let ProgramInput { pre_states, .. } = read_nssa_inputs::(); + let (ProgramInput { pre_states, .. }, instruction_words) = read_nssa_inputs::(); let [pre] = match pre_states.try_into() { Ok(array) => array, @@ -15,5 +15,9 @@ fn main() { let account_pre = pre.account.clone(); - write_nssa_outputs(vec![pre], vec![account_pre, Account::default()]); + write_nssa_outputs( + instruction_words, + vec![pre], + vec![account_pre, Account::default()], + ); } diff --git a/nssa/test_program_methods/guest/src/bin/minter.rs b/nssa/test_program_methods/guest/src/bin/minter.rs index 2ec97a9..08bb488 100644 --- a/nssa/test_program_methods/guest/src/bin/minter.rs +++ b/nssa/test_program_methods/guest/src/bin/minter.rs @@ -1,9 +1,9 @@ -use nssa_core::program::{read_nssa_inputs, write_nssa_outputs, ProgramInput}; +use nssa_core::program::{ProgramInput, read_nssa_inputs, write_nssa_outputs}; type Instruction = (); fn main() { - let ProgramInput { pre_states, .. } = read_nssa_inputs::(); + let (ProgramInput { pre_states, .. }, instruction_words) = read_nssa_inputs::(); let [pre] = match pre_states.try_into() { Ok(array) => array, @@ -14,5 +14,5 @@ fn main() { let mut account_post = account_pre.clone(); account_post.balance += 1; - write_nssa_outputs(vec![pre], vec![account_post]); + write_nssa_outputs(instruction_words, vec![pre], vec![account_post]); } diff --git a/nssa/test_program_methods/guest/src/bin/missing_output.rs b/nssa/test_program_methods/guest/src/bin/missing_output.rs index 7b6016c..8c0a3f6 100644 --- a/nssa/test_program_methods/guest/src/bin/missing_output.rs +++ b/nssa/test_program_methods/guest/src/bin/missing_output.rs @@ -1,9 +1,9 @@ -use nssa_core::program::{read_nssa_inputs, write_nssa_outputs, ProgramInput}; +use nssa_core::program::{ProgramInput, read_nssa_inputs, write_nssa_outputs}; type Instruction = (); fn main() { - let ProgramInput { pre_states, .. } = read_nssa_inputs::(); + let (ProgramInput { pre_states, .. }, instruction_words) = read_nssa_inputs::(); let [pre1, pre2] = match pre_states.try_into() { Ok(array) => array, @@ -12,5 +12,5 @@ fn main() { let account_pre1 = pre1.account.clone(); - write_nssa_outputs(vec![pre1, pre2], vec![account_pre1]); + write_nssa_outputs(instruction_words, vec![pre1, pre2], vec![account_pre1]); } diff --git a/nssa/test_program_methods/guest/src/bin/nonce_changer.rs b/nssa/test_program_methods/guest/src/bin/nonce_changer.rs index b3b2599..c5d706b 100644 --- a/nssa/test_program_methods/guest/src/bin/nonce_changer.rs +++ b/nssa/test_program_methods/guest/src/bin/nonce_changer.rs @@ -3,7 +3,7 @@ use nssa_core::program::{read_nssa_inputs, write_nssa_outputs, ProgramInput}; type Instruction = (); fn main() { - let ProgramInput { pre_states, .. } = read_nssa_inputs::(); + let (ProgramInput { pre_states, .. } , instruction_words) = read_nssa_inputs::(); let [pre] = match pre_states.try_into() { Ok(array) => array, @@ -14,5 +14,5 @@ fn main() { let mut account_post = account_pre.clone(); account_post.nonce += 1; - write_nssa_outputs(vec![pre], vec![account_post]); + write_nssa_outputs(instruction_words, vec![pre], vec![account_post]); } diff --git a/nssa/test_program_methods/guest/src/bin/program_owner_changer.rs b/nssa/test_program_methods/guest/src/bin/program_owner_changer.rs index 49947cd..e8afc7c 100644 --- a/nssa/test_program_methods/guest/src/bin/program_owner_changer.rs +++ b/nssa/test_program_methods/guest/src/bin/program_owner_changer.rs @@ -1,9 +1,9 @@ -use nssa_core::program::{read_nssa_inputs, write_nssa_outputs, ProgramInput}; +use nssa_core::program::{ProgramInput, read_nssa_inputs, write_nssa_outputs}; type Instruction = (); fn main() { - let ProgramInput { pre_states, .. } = read_nssa_inputs::(); + let (ProgramInput { pre_states, .. }, instruction_words) = read_nssa_inputs::(); let [pre] = match pre_states.try_into() { Ok(array) => array, @@ -14,5 +14,5 @@ fn main() { let mut account_post = account_pre.clone(); account_post.program_owner = [0, 1, 2, 3, 4, 5, 6, 7]; - write_nssa_outputs(vec![pre], vec![account_post]); + write_nssa_outputs(instruction_words, vec![pre], vec![account_post]); } diff --git a/nssa/test_program_methods/guest/src/bin/simple_balance_transfer.rs b/nssa/test_program_methods/guest/src/bin/simple_balance_transfer.rs index 13263c5..e3e64a3 100644 --- a/nssa/test_program_methods/guest/src/bin/simple_balance_transfer.rs +++ b/nssa/test_program_methods/guest/src/bin/simple_balance_transfer.rs @@ -1,12 +1,15 @@ -use nssa_core::program::{read_nssa_inputs, write_nssa_outputs, ProgramInput}; +use nssa_core::program::{ProgramInput, read_nssa_inputs, write_nssa_outputs}; type Instruction = u128; fn main() { - let ProgramInput { - pre_states, - instruction: balance, - } = read_nssa_inputs::(); + let ( + ProgramInput { + pre_states, + instruction: balance, + }, + instruction_words, + ) = read_nssa_inputs::(); let [sender_pre, receiver_pre] = match pre_states.try_into() { Ok(array) => array, @@ -19,6 +22,7 @@ fn main() { receiver_post.balance += balance; write_nssa_outputs( + instruction_words, vec![sender_pre, receiver_pre], vec![sender_post, receiver_post], );