diff --git a/common/src/sequencer_client.rs b/common/src/sequencer_client.rs index d3c5f237..622c0c11 100644 --- a/common/src/sequencer_client.rs +++ b/common/src/sequencer_client.rs @@ -30,16 +30,25 @@ use crate::{ pub struct SequencerClient { pub client: reqwest::Client, pub sequencer_addr: String, + pub basic_auth: Option<(String, Option)>, } impl SequencerClient { pub fn new(sequencer_addr: String) -> Result { + Self::new_with_auth(sequencer_addr, None) + } + + pub fn new_with_auth( + sequencer_addr: String, + basic_auth: Option<(String, Option)>, + ) -> Result { Ok(Self { client: Client::builder() //Add more fiedls if needed .timeout(std::time::Duration::from_secs(60)) .build()?, sequencer_addr, + basic_auth, }) } @@ -51,13 +60,16 @@ impl SequencerClient { let request = rpc_primitives::message::Request::from_payload_version_2_0(method.to_string(), payload); - let call_builder = self.client.post(&self.sequencer_addr); + let mut call_builder = self.client.post(&self.sequencer_addr); + + if let Some((username, password)) = &self.basic_auth { + call_builder = call_builder.basic_auth(username, password.as_deref()); + } let call_res = call_builder.json(&request).send().await?; let response_vall = call_res.json::().await?; - // TODO: Actually why we need separation of `result` and `error` in rpc response? #[derive(Debug, Clone, Deserialize)] #[allow(dead_code)] pub struct SequencerRpcResponse { diff --git a/integration_tests/data_changer.bin b/integration_tests/data_changer.bin new file mode 100644 index 00000000..eb28a627 Binary files /dev/null and b/integration_tests/data_changer.bin differ diff --git a/integration_tests/src/data_changer.bin b/integration_tests/src/data_changer.bin deleted file mode 100644 index 6a36d52c..00000000 Binary files a/integration_tests/src/data_changer.bin and /dev/null differ diff --git a/integration_tests/src/lib.rs b/integration_tests/src/lib.rs index 31cd1771..401238fa 100644 --- a/integration_tests/src/lib.rs +++ b/integration_tests/src/lib.rs @@ -42,7 +42,7 @@ pub const ACC_RECEIVER_PRIVATE: &str = "AKTcXgJ1xoynta1Ec7y6Jso1z1JQtHqd7aPQ1h9e pub const TIME_TO_WAIT_FOR_BLOCK_SECONDS: u64 = 12; -pub const NSSA_PROGRAM_FOR_TEST_DATA_CHANGER: &[u8] = include_bytes!("data_changer.bin"); +pub const NSSA_PROGRAM_FOR_TEST_DATA_CHANGER: &str = "data_changer.bin"; fn make_public_account_input_from_str(account_id: &str) -> String { format!("Public/{account_id}") diff --git a/integration_tests/src/test_suite_map.rs b/integration_tests/src/test_suite_map.rs index 1c5f91fd..8012e17e 100644 --- a/integration_tests/src/test_suite_map.rs +++ b/integration_tests/src/test_suite_map.rs @@ -2,6 +2,7 @@ use std::{ collections::HashMap, path::PathBuf, pin::Pin, + str::FromStr, time::{Duration, Instant}, }; @@ -10,7 +11,7 @@ use anyhow::Result; use common::{PINATA_BASE58, sequencer_client::SequencerClient}; use key_protocol::key_management::key_tree::chain_index::ChainIndex; use log::info; -use nssa::{AccountId, ProgramDeploymentTransaction, program::Program}; +use nssa::{AccountId, program::Program}; use nssa_core::{NullifierPublicKey, encryption::shared_key_derivation::Secp256k1Point}; use sequencer_runner::startup_sequencer; use tempfile::TempDir; @@ -86,9 +87,7 @@ pub fn prepare_function_map() -> HashMap { #[nssa_integration_test] pub async fn test_success_move_to_another_account() { info!("########## test_success_move_to_another_account ##########"); - let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: ChainIndex::root(), - })); + let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public { cci: None })); let wallet_config = fetch_config().await.unwrap(); @@ -292,9 +291,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: definition_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Public { - cci: ChainIndex::root(), - }, + NewSubcommand::Public { cci: None }, ))) .await .unwrap() @@ -305,9 +302,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: supply_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Public { - cci: ChainIndex::root(), - }, + NewSubcommand::Public { cci: None }, ))) .await .unwrap() @@ -318,9 +313,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: recipient_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Public { - cci: ChainIndex::root(), - }, + NewSubcommand::Public { cci: None }, ))) .await .unwrap() @@ -356,8 +349,8 @@ pub fn prepare_function_map() -> HashMap { // The data of a token definition account has the following layout: // [ 0x00 || name (6 bytes) || total supply (little endian 16 bytes) ] assert_eq!( - definition_acc.data, - vec![ + definition_acc.data.as_ref(), + &[ 0, 65, 32, 78, 65, 77, 69, 37, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 ] ); @@ -375,11 +368,14 @@ pub fn prepare_function_map() -> HashMap { // The data of a token definition account has the following layout: // [ 0x01 || corresponding_token_definition_id (32 bytes) || balance (little endian 16 // bytes) ] First byte of the data equal to 1 means it's a token holding account - assert_eq!(supply_acc.data[0], 1); + assert_eq!(supply_acc.data.as_ref()[0], 1); // Bytes from 1 to 33 represent the id of the token this account is associated with. // In this example, this is a token account of the newly created token, so it is expected // to be equal to the account_id of the token definition account. - assert_eq!(&supply_acc.data[1..33], definition_account_id.to_bytes()); + assert_eq!( + &supply_acc.data.as_ref()[1..33], + definition_account_id.to_bytes() + ); assert_eq!( u128::from_le_bytes(supply_acc.data[33..].try_into().unwrap()), 37 @@ -441,20 +437,18 @@ pub fn prepare_function_map() -> HashMap { } /// This test creates a new private token using the token program. After creating the token, the - /// test executes a private token transfer to a new account. All accounts are owned except - /// definition. + /// test executes a private token transfer to a new account. All accounts are private owned + /// except definition which is public. #[nssa_integration_test] - pub async fn test_success_token_program_private_owned() { - info!("########## test_success_token_program_private_owned ##########"); + pub async fn test_success_token_program_private_owned_supply() { + info!("########## test_success_token_program_private_owned_supply ##########"); let wallet_config = fetch_config().await.unwrap(); // Create new account for the token definition (public) let SubcommandReturnValue::RegisterAccount { account_id: definition_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Public { - cci: ChainIndex::root(), - }, + NewSubcommand::Public { cci: None }, ))) .await .unwrap() @@ -465,9 +459,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: supply_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Private { - cci: ChainIndex::root(), - }, + NewSubcommand::Private { cci: None }, ))) .await .unwrap() @@ -478,9 +470,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: recipient_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Private { - cci: ChainIndex::root(), - }, + NewSubcommand::Private { cci: None }, ))) .await .unwrap() @@ -518,8 +508,8 @@ pub fn prepare_function_map() -> HashMap { // The data of a token definition account has the following layout: // [ 0x00 || name (6 bytes) || total supply (little endian 16 bytes) ] assert_eq!( - definition_acc.data, - vec![ + definition_acc.data.as_ref(), + &[ 0, 65, 32, 78, 65, 77, 69, 37, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 ] ); @@ -602,19 +592,104 @@ pub fn prepare_function_map() -> HashMap { assert!(verify_commitment_is_in_state(new_commitment2, &seq_client).await); } - /// This test creates a new private token using the token program. After creating the token, the - /// test executes a private token transfer to a new account. + /// This test creates a new private token using the token program. All accounts are private + /// owned except supply which is public. #[nssa_integration_test] - pub async fn test_success_token_program_private_claiming_path() { - info!("########## test_success_token_program_private_claiming_path ##########"); + pub async fn test_success_token_program_private_owned_definition() { + info!("########## test_success_token_program_private_owned_definition ##########"); let wallet_config = fetch_config().await.unwrap(); - // Create new account for the token definition (public) + // Create new account for the token definition (private) let SubcommandReturnValue::RegisterAccount { account_id: definition_account_id, + } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( + NewSubcommand::Private { + cci: Some(ChainIndex::root()), + }, + ))) + .await + .unwrap() + else { + panic!("invalid subcommand return value"); + }; + // Create new account for the token supply holder (public) + let SubcommandReturnValue::RegisterAccount { + account_id: supply_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( NewSubcommand::Public { - cci: ChainIndex::root(), + cci: Some(ChainIndex::root()), + }, + ))) + .await + .unwrap() + else { + panic!("invalid subcommand return value"); + }; + + // Create new token + let subcommand = TokenProgramAgnosticSubcommand::New { + definition_account_id: make_private_account_input_from_str( + &definition_account_id.to_string(), + ), + supply_account_id: make_public_account_input_from_str(&supply_account_id.to_string()), + name: "A NAME".to_string(), + total_supply: 37, + }; + + wallet::cli::execute_subcommand(Command::Token(subcommand)) + .await + .unwrap(); + + info!("Waiting for next block creation"); + tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; + + let seq_client = SequencerClient::new(wallet_config.sequencer_addr.clone()).unwrap(); + + let wallet_config = fetch_config().await.unwrap(); + let wallet_storage = WalletCore::start_from_config_update_chain(wallet_config) + .await + .unwrap(); + + let new_commitment1 = wallet_storage + .get_private_account_commitment(&definition_account_id) + .unwrap(); + assert!(verify_commitment_is_in_state(new_commitment1, &seq_client).await); + + // Check the status of the token definition account is the expected after the execution + let supply_acc = seq_client + .get_account(supply_account_id.to_string()) + .await + .unwrap() + .account; + + assert_eq!(supply_acc.program_owner, Program::token().id()); + // The data of a token definition account has the following layout: + // [ 0x00 || name (6 bytes) || total supply (little endian 16 bytes) ] + assert_eq!( + supply_acc.data.as_ref(), + &[ + 1, 128, 101, 5, 31, 43, 36, 97, 108, 164, 92, 25, 157, 173, 5, 14, 194, 121, 239, + 84, 19, 160, 243, 47, 193, 2, 250, 247, 232, 253, 191, 232, 173, 37, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 + ] + ); + } + + /// This test creates a new private token using the token program. All accounts are private + /// owned. + #[nssa_integration_test] + pub async fn test_success_token_program_private_owned_definition_and_supply() { + info!( + "########## test_success_token_program_private_owned_definition_and_supply ##########" + ); + let wallet_config = fetch_config().await.unwrap(); + + // Create new account for the token definition (private) + let SubcommandReturnValue::RegisterAccount { + account_id: definition_account_id, + } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( + NewSubcommand::Private { + cci: Some(ChainIndex::root()), }, ))) .await @@ -627,7 +702,7 @@ pub fn prepare_function_map() -> HashMap { account_id: supply_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( NewSubcommand::Private { - cci: ChainIndex::root(), + cci: Some(ChainIndex::root()), }, ))) .await @@ -635,13 +710,105 @@ pub fn prepare_function_map() -> HashMap { else { panic!("invalid subcommand return value"); }; + + // Create new token + let subcommand = TokenProgramAgnosticSubcommand::New { + definition_account_id: make_private_account_input_from_str( + &definition_account_id.to_string(), + ), + supply_account_id: make_private_account_input_from_str(&supply_account_id.to_string()), + name: "A NAME".to_string(), + total_supply: 37, + }; + + wallet::cli::execute_subcommand(Command::Token(subcommand)) + .await + .unwrap(); + + info!("Waiting for next block creation"); + tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; + + let seq_client = SequencerClient::new(wallet_config.sequencer_addr.clone()).unwrap(); + + let wallet_config = fetch_config().await.unwrap(); + let wallet_storage = WalletCore::start_from_config_update_chain(wallet_config) + .await + .unwrap(); + + let new_commitment1 = wallet_storage + .get_private_account_commitment(&definition_account_id) + .unwrap(); + assert!(verify_commitment_is_in_state(new_commitment1, &seq_client).await); + + let new_commitment2 = wallet_storage + .get_private_account_commitment(&supply_account_id) + .unwrap(); + assert!(verify_commitment_is_in_state(new_commitment2, &seq_client).await); + + let definition_acc = wallet_storage + .get_account_private(&definition_account_id) + .unwrap(); + let supply_acc = wallet_storage + .get_account_private(&supply_account_id) + .unwrap(); + + assert_eq!(definition_acc.program_owner, Program::token().id()); + // The data of a token definition account has the following layout: + // [ 0x00 || name (6 bytes) || total supply (little endian 16 bytes) ] + assert_eq!( + definition_acc.data.as_ref(), + &[ + 0, 65, 32, 78, 65, 77, 69, 37, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 + ] + ); + + assert_eq!(supply_acc.program_owner, Program::token().id()); + // The data of a token definition account has the following layout: + // [ 0x00 || name (6 bytes) || total supply (little endian 16 bytes) ] + assert_eq!( + supply_acc.data.as_ref(), + &[ + 1, 128, 101, 5, 31, 43, 36, 97, 108, 164, 92, 25, 157, 173, 5, 14, 194, 121, 239, + 84, 19, 160, 243, 47, 193, 2, 250, 247, 232, 253, 191, 232, 173, 37, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 + ] + ); + } + + /// This test creates a new private token using the token program. After creating the token, the + /// test executes a private token transfer to a new account. + #[nssa_integration_test] + pub async fn test_success_token_program_private_claiming_path() { + info!("########## test_success_token_program_private_claiming_path ##########"); + let wallet_config = fetch_config().await.unwrap(); + + // Create new account for the token definition (public) + let SubcommandReturnValue::RegisterAccount { + account_id: definition_account_id, + } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( + NewSubcommand::Public { cci: None }, + ))) + .await + .unwrap() + else { + panic!("invalid subcommand return value"); + }; + // Create new account for the token supply holder (private) + let SubcommandReturnValue::RegisterAccount { + account_id: supply_account_id, + } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( + NewSubcommand::Private { cci: None }, + ))) + .await + .unwrap() + else { + panic!("invalid subcommand return value"); + }; // Create new account for receiving a token transaction let SubcommandReturnValue::RegisterAccount { account_id: recipient_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Private { - cci: ChainIndex::root(), - }, + NewSubcommand::Private { cci: None }, ))) .await .unwrap() @@ -679,8 +846,8 @@ pub fn prepare_function_map() -> HashMap { // The data of a token definition account has the following layout: // [ 0x00 || name (6 bytes) || total supply (little endian 16 bytes) ] assert_eq!( - definition_acc.data, - vec![ + definition_acc.data.as_ref(), + &[ 0, 65, 32, 78, 65, 77, 69, 37, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 ] ); @@ -755,9 +922,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: definition_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Public { - cci: ChainIndex::root(), - }, + NewSubcommand::Public { cci: None }, ))) .await .unwrap() @@ -768,9 +933,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: supply_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Public { - cci: ChainIndex::root(), - }, + NewSubcommand::Public { cci: None }, ))) .await .unwrap() @@ -781,9 +944,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: recipient_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Private { - cci: ChainIndex::root(), - }, + NewSubcommand::Private { cci: None }, ))) .await .unwrap() @@ -821,8 +982,8 @@ pub fn prepare_function_map() -> HashMap { // The data of a token definition account has the following layout: // [ 0x00 || name (6 bytes) || total supply (little endian 16 bytes) ] assert_eq!( - definition_acc.data, - vec![ + definition_acc.data.as_ref(), + &[ 0, 65, 32, 78, 65, 77, 69, 37, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 ] ); @@ -897,9 +1058,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: definition_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Public { - cci: ChainIndex::root(), - }, + NewSubcommand::Public { cci: None }, ))) .await .unwrap() @@ -910,9 +1069,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: supply_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Private { - cci: ChainIndex::root(), - }, + NewSubcommand::Private { cci: None }, ))) .await .unwrap() @@ -923,9 +1080,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: recipient_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Public { - cci: ChainIndex::root(), - }, + NewSubcommand::Public { cci: None }, ))) .await .unwrap() @@ -963,8 +1118,8 @@ pub fn prepare_function_map() -> HashMap { // The data of a token definition account has the following layout: // [ 0x00 || name (6 bytes) || total supply (little endian 16 bytes) ] assert_eq!( - definition_acc.data, - vec![ + definition_acc.data.as_ref(), + &[ 0, 65, 32, 78, 65, 77, 69, 37, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 ] ); @@ -1126,9 +1281,8 @@ pub fn prepare_function_map() -> HashMap { ); let from: AccountId = ACC_SENDER_PRIVATE.parse().unwrap(); - let command = Command::Account(AccountSubcommand::New(NewSubcommand::Private { - cci: ChainIndex::root(), - })); + let command = + Command::Account(AccountSubcommand::New(NewSubcommand::Private { cci: None })); let sub_ret = wallet::cli::execute_subcommand(command).await.unwrap(); let SubcommandReturnValue::RegisterAccount { @@ -1441,15 +1595,18 @@ pub fn prepare_function_map() -> HashMap { #[nssa_integration_test] pub async fn test_program_deployment() { info!("########## test program deployment ##########"); - let bytecode = NSSA_PROGRAM_FOR_TEST_DATA_CHANGER.to_vec(); - let message = nssa::program_deployment_transaction::Message::new(bytecode.clone()); - let transaction = ProgramDeploymentTransaction::new(message); + + let binary_filepath: PathBuf = NSSA_PROGRAM_FOR_TEST_DATA_CHANGER.parse().unwrap(); + + let command = Command::DeployProgram { + binary_filepath: binary_filepath.clone(), + }; + + wallet::cli::execute_subcommand(command).await.unwrap(); let wallet_config = fetch_config().await.unwrap(); let seq_client = SequencerClient::new(wallet_config.sequencer_addr.clone()).unwrap(); - let _response = seq_client.send_tx_program(transaction).await.unwrap(); - info!("Waiting for next block creation"); tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; @@ -1457,13 +1614,15 @@ pub fn prepare_function_map() -> HashMap { // We pass an uninitialized account and we expect after execution to be owned by the data // changer program (NSSA account claiming mechanism) with data equal to [0] (due to program // logic) + // + let bytecode = std::fs::read(binary_filepath).unwrap(); let data_changer = Program::new(bytecode).unwrap(); let account_id: AccountId = "11".repeat(16).parse().unwrap(); let message = nssa::public_transaction::Message::try_new( data_changer.id(), vec![account_id], vec![], - (), + vec![0], ) .unwrap(); let witness_set = nssa::public_transaction::WitnessSet::for_message(&message, &[]); @@ -1480,7 +1639,7 @@ pub fn prepare_function_map() -> HashMap { .account; assert_eq!(post_state_account.program_owner, data_changer.id()); assert_eq!(post_state_account.balance, 0); - assert_eq!(post_state_account.data, vec![0]); + assert_eq!(post_state_account.data.as_ref(), &[0]); assert_eq!(post_state_account.nonce, 0); info!("Success!"); @@ -1489,9 +1648,7 @@ pub fn prepare_function_map() -> HashMap { #[nssa_integration_test] pub async fn test_authenticated_transfer_initialize_function() { info!("########## test initialize account for authenticated transfer ##########"); - let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: ChainIndex::root(), - })); + let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public { cci: None })); let SubcommandReturnValue::RegisterAccount { account_id } = wallet::cli::execute_subcommand(command).await.unwrap() else { @@ -1589,9 +1746,7 @@ pub fn prepare_function_map() -> HashMap { let SubcommandReturnValue::RegisterAccount { account_id: winner_account_id, } = wallet::cli::execute_subcommand(Command::Account(AccountSubcommand::New( - NewSubcommand::Private { - cci: ChainIndex::root(), - }, + NewSubcommand::Private { cci: None }, ))) .await .unwrap() @@ -1669,6 +1824,217 @@ pub fn prepare_function_map() -> HashMap { info!("Success!"); } + #[nssa_integration_test] + pub async fn test_keys_restoration() { + info!("########## test_keys_restoration ##########"); + let from: AccountId = ACC_SENDER_PRIVATE.parse().unwrap(); + + let command = Command::Account(AccountSubcommand::New(NewSubcommand::Private { + cci: Some(ChainIndex::root()), + })); + + let sub_ret = wallet::cli::execute_subcommand(command).await.unwrap(); + let SubcommandReturnValue::RegisterAccount { + account_id: to_account_id1, + } = sub_ret + else { + panic!("FAILED TO REGISTER ACCOUNT"); + }; + + let command = Command::Account(AccountSubcommand::New(NewSubcommand::Private { + cci: Some(ChainIndex::from_str("/0").unwrap()), + })); + + let sub_ret = wallet::cli::execute_subcommand(command).await.unwrap(); + let SubcommandReturnValue::RegisterAccount { + account_id: to_account_id2, + } = sub_ret + else { + panic!("FAILED TO REGISTER ACCOUNT"); + }; + + let command = Command::AuthTransfer(AuthTransferSubcommand::Send { + from: make_private_account_input_from_str(&from.to_string()), + to: Some(make_private_account_input_from_str( + &to_account_id1.to_string(), + )), + to_npk: None, + to_ipk: None, + amount: 100, + }); + + wallet::cli::execute_subcommand(command).await.unwrap(); + + let command = Command::AuthTransfer(AuthTransferSubcommand::Send { + from: make_private_account_input_from_str(&from.to_string()), + to: Some(make_private_account_input_from_str( + &to_account_id2.to_string(), + )), + to_npk: None, + to_ipk: None, + amount: 101, + }); + + wallet::cli::execute_subcommand(command).await.unwrap(); + + let from: AccountId = ACC_SENDER.parse().unwrap(); + + let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public { + cci: Some(ChainIndex::root()), + })); + + let sub_ret = wallet::cli::execute_subcommand(command).await.unwrap(); + let SubcommandReturnValue::RegisterAccount { + account_id: to_account_id3, + } = sub_ret + else { + panic!("FAILED TO REGISTER ACCOUNT"); + }; + + let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public { + cci: Some(ChainIndex::from_str("/0").unwrap()), + })); + + let sub_ret = wallet::cli::execute_subcommand(command).await.unwrap(); + let SubcommandReturnValue::RegisterAccount { + account_id: to_account_id4, + } = sub_ret + else { + panic!("FAILED TO REGISTER ACCOUNT"); + }; + + let command = Command::AuthTransfer(AuthTransferSubcommand::Send { + from: make_public_account_input_from_str(&from.to_string()), + to: Some(make_public_account_input_from_str( + &to_account_id3.to_string(), + )), + to_npk: None, + to_ipk: None, + amount: 102, + }); + + wallet::cli::execute_subcommand(command).await.unwrap(); + + let command = Command::AuthTransfer(AuthTransferSubcommand::Send { + from: make_public_account_input_from_str(&from.to_string()), + to: Some(make_public_account_input_from_str( + &to_account_id4.to_string(), + )), + to_npk: None, + to_ipk: None, + amount: 103, + }); + + wallet::cli::execute_subcommand(command).await.unwrap(); + + info!("########## PREPARATION END ##########"); + + wallet::cli::execute_keys_restoration("test_pass".to_string(), 10) + .await + .unwrap(); + + let wallet_config = fetch_config().await.unwrap(); + let wallet_storage = WalletCore::start_from_config_update_chain(wallet_config.clone()) + .await + .unwrap(); + + let acc1 = wallet_storage + .storage + .user_data + .private_key_tree + .get_node(to_account_id1) + .expect("Acc 1 should be restored"); + + let acc2 = wallet_storage + .storage + .user_data + .private_key_tree + .get_node(to_account_id2) + .expect("Acc 2 should be restored"); + + let _ = wallet_storage + .storage + .user_data + .public_key_tree + .get_node(to_account_id3) + .expect("Acc 3 should be restored"); + + let _ = wallet_storage + .storage + .user_data + .public_key_tree + .get_node(to_account_id4) + .expect("Acc 4 should be restored"); + + assert_eq!( + acc1.value.1.program_owner, + Program::authenticated_transfer_program().id() + ); + assert_eq!( + acc2.value.1.program_owner, + Program::authenticated_transfer_program().id() + ); + + assert_eq!(acc1.value.1.balance, 100); + assert_eq!(acc2.value.1.balance, 101); + + info!("########## TREE CHECKS END ##########"); + + let command = Command::AuthTransfer(AuthTransferSubcommand::Send { + from: make_private_account_input_from_str(&to_account_id1.to_string()), + to: Some(make_private_account_input_from_str( + &to_account_id2.to_string(), + )), + to_npk: None, + to_ipk: None, + amount: 10, + }); + + wallet::cli::execute_subcommand(command).await.unwrap(); + + let command = Command::AuthTransfer(AuthTransferSubcommand::Send { + from: make_public_account_input_from_str(&to_account_id3.to_string()), + to: Some(make_public_account_input_from_str( + &to_account_id4.to_string(), + )), + to_npk: None, + to_ipk: None, + amount: 11, + }); + + wallet::cli::execute_subcommand(command).await.unwrap(); + + let wallet_config = fetch_config().await.unwrap(); + let seq_client = SequencerClient::new(wallet_config.sequencer_addr.clone()).unwrap(); + let wallet_storage = WalletCore::start_from_config_update_chain(wallet_config.clone()) + .await + .unwrap(); + + let comm1 = wallet_storage + .get_private_account_commitment(&to_account_id1) + .expect("Acc 1 commitment should exist"); + let comm2 = wallet_storage + .get_private_account_commitment(&to_account_id2) + .expect("Acc 2 commitment should exist"); + + assert!(verify_commitment_is_in_state(comm1, &seq_client).await); + assert!(verify_commitment_is_in_state(comm2, &seq_client).await); + + let acc3 = seq_client + .get_account_balance(to_account_id3.to_string()) + .await + .expect("Acc 3 must be present in public state"); + let acc4 = seq_client + .get_account_balance(to_account_id4.to_string()) + .await + .expect("Acc 4 must be present in public state"); + + assert_eq!(acc3.balance, 91); + assert_eq!(acc4.balance, 114); + + info!("Success!"); + } + println!("{function_map:#?}"); function_map diff --git a/integration_tests/src/tps_test_utils.rs b/integration_tests/src/tps_test_utils.rs index d06a08f6..6f597e21 100644 --- a/integration_tests/src/tps_test_utils.rs +++ b/integration_tests/src/tps_test_utils.rs @@ -8,7 +8,8 @@ use nssa::{ public_transaction as putx, }; use nssa_core::{ - MembershipProof, NullifierPublicKey, account::AccountWithMetadata, + MembershipProof, NullifierPublicKey, + account::{AccountWithMetadata, data::Data}, encryption::IncomingViewingPublicKey, }; use sequencer_core::config::{AccountInitialData, CommitmentsInitialData, SequencerConfig}; @@ -90,7 +91,7 @@ impl TpsTestManager { balance: 100, nonce: 0xdeadbeef, program_owner: Program::authenticated_transfer_program().id(), - data: vec![], + data: Data::default(), }; let initial_commitment = CommitmentsInitialData { npk: sender_npk, @@ -129,7 +130,7 @@ fn build_privacy_transaction() -> PrivacyPreservingTransaction { balance: 100, nonce: 0xdeadbeef, program_owner: program.id(), - data: vec![], + data: Data::default(), }, true, AccountId::from(&sender_npk), @@ -167,7 +168,7 @@ fn build_privacy_transaction() -> PrivacyPreservingTransaction { (recipient_npk.clone(), recipient_ss), ], &[(sender_nsk, proof)], - &program, + &program.into(), ) .unwrap(); let message = pptx::message::Message::try_from_circuit_output( diff --git a/key_protocol/Cargo.toml b/key_protocol/Cargo.toml index a5625159..103a1dee 100644 --- a/key_protocol/Cargo.toml +++ b/key_protocol/Cargo.toml @@ -16,6 +16,7 @@ bip39.workspace = true hmac-sha512.workspace = true thiserror.workspace = true nssa-core = { path = "../nssa/core", features = ["host"] } +itertools.workspace = true [dependencies.common] path = "../common" diff --git a/key_protocol/src/key_management/key_tree/chain_index.rs b/key_protocol/src/key_management/key_tree/chain_index.rs index e46fc0f0..6dbaf9a9 100644 --- a/key_protocol/src/key_management/key_tree/chain_index.rs +++ b/key_protocol/src/key_management/key_tree/chain_index.rs @@ -1,8 +1,9 @@ use std::{fmt::Display, str::FromStr}; +use itertools::Itertools; use serde::{Deserialize, Serialize}; -#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize)] +#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, Hash)] pub struct ChainIndex(Vec); #[derive(thiserror::Error, Debug)] @@ -77,12 +78,82 @@ impl ChainIndex { ChainIndex(chain) } + pub fn previous_in_line(&self) -> Option { + let mut chain = self.0.clone(); + if let Some(last_p) = chain.last_mut() { + *last_p = last_p.checked_sub(1)?; + } + + Some(ChainIndex(chain)) + } + + pub fn parent(&self) -> Option { + if self.0.is_empty() { + None + } else { + Some(ChainIndex(self.0[..(self.0.len() - 1)].to_vec())) + } + } + pub fn nth_child(&self, child_id: u32) -> ChainIndex { let mut chain = self.0.clone(); chain.push(child_id); ChainIndex(chain) } + + pub fn depth(&self) -> u32 { + self.0.iter().map(|cci| cci + 1).sum() + } + + fn collapse_back(&self) -> Option { + let mut res = self.parent()?; + + let last_mut = res.0.last_mut()?; + *last_mut += *(self.0.last()?) + 1; + + Some(res) + } + + fn shuffle_iter(&self) -> impl Iterator { + self.0 + .iter() + .permutations(self.0.len()) + .unique() + .map(|item| ChainIndex(item.into_iter().cloned().collect())) + } + + pub fn chain_ids_at_depth(depth: usize) -> impl Iterator { + let mut stack = vec![ChainIndex(vec![0; depth])]; + let mut cumulative_stack = vec![ChainIndex(vec![0; depth])]; + + while let Some(id) = stack.pop() { + if let Some(collapsed_id) = id.collapse_back() { + for id in collapsed_id.shuffle_iter() { + stack.push(id.clone()); + cumulative_stack.push(id); + } + } + } + + cumulative_stack.into_iter().unique() + } + + pub fn chain_ids_at_depth_rev(depth: usize) -> impl Iterator { + let mut stack = vec![ChainIndex(vec![0; depth])]; + let mut cumulative_stack = vec![ChainIndex(vec![0; depth])]; + + while let Some(id) = stack.pop() { + if let Some(collapsed_id) = id.collapse_back() { + for id in collapsed_id.shuffle_iter() { + stack.push(id.clone()); + cumulative_stack.push(id); + } + } + } + + cumulative_stack.into_iter().rev().unique() + } } #[cfg(test)] @@ -145,4 +216,83 @@ mod tests { assert_eq!(string_index, "/5/7/8".to_string()); } + + #[test] + fn test_prev_in_line() { + let chain_id = ChainIndex(vec![1, 7, 3]); + + let prev_chain_id = chain_id.previous_in_line().unwrap(); + + assert_eq!(prev_chain_id, ChainIndex(vec![1, 7, 2])) + } + + #[test] + fn test_prev_in_line_no_prev() { + let chain_id = ChainIndex(vec![1, 7, 0]); + + let prev_chain_id = chain_id.previous_in_line(); + + assert_eq!(prev_chain_id, None) + } + + #[test] + fn test_parent() { + let chain_id = ChainIndex(vec![1, 7, 3]); + + let parent_chain_id = chain_id.parent().unwrap(); + + assert_eq!(parent_chain_id, ChainIndex(vec![1, 7])) + } + + #[test] + fn test_parent_no_parent() { + let chain_id = ChainIndex(vec![]); + + let parent_chain_id = chain_id.parent(); + + assert_eq!(parent_chain_id, None) + } + + #[test] + fn test_parent_root() { + let chain_id = ChainIndex(vec![1]); + + let parent_chain_id = chain_id.parent().unwrap(); + + assert_eq!(parent_chain_id, ChainIndex::root()) + } + + #[test] + fn test_collapse_back() { + let chain_id = ChainIndex(vec![1, 1]); + + let collapsed = chain_id.collapse_back().unwrap(); + + assert_eq!(collapsed, ChainIndex(vec![3])) + } + + #[test] + fn test_collapse_back_one() { + let chain_id = ChainIndex(vec![1]); + + let collapsed = chain_id.collapse_back(); + + assert_eq!(collapsed, None) + } + + #[test] + fn test_collapse_back_root() { + let chain_id = ChainIndex(vec![]); + + let collapsed = chain_id.collapse_back(); + + assert_eq!(collapsed, None) + } + + #[test] + fn test_shuffle() { + for id in ChainIndex::chain_ids_at_depth(5) { + println!("{id}"); + } + } } diff --git a/key_protocol/src/key_management/key_tree/mod.rs b/key_protocol/src/key_management/key_tree/mod.rs index 1319a5a8..389580b6 100644 --- a/key_protocol/src/key_management/key_tree/mod.rs +++ b/key_protocol/src/key_management/key_tree/mod.rs @@ -1,5 +1,10 @@ -use std::collections::{BTreeMap, HashMap}; +use std::{ + collections::{BTreeMap, HashMap}, + sync::Arc, +}; +use anyhow::Result; +use common::sequencer_client::SequencerClient; use serde::{Deserialize, Serialize}; use crate::key_management::{ @@ -15,6 +20,8 @@ pub mod keys_private; pub mod keys_public; pub mod traits; +pub const DEPTH_SOFT_CAP: u32 = 20; + #[derive(Debug, Serialize, Deserialize, Clone)] pub struct KeyTree { pub key_map: BTreeMap, @@ -96,21 +103,53 @@ impl KeyTree { } } - pub fn generate_new_node(&mut self, parent_cci: ChainIndex) -> Option { - let father_keys = self.key_map.get(&parent_cci)?; + pub fn generate_new_node( + &mut self, + parent_cci: &ChainIndex, + ) -> Option<(nssa::AccountId, ChainIndex)> { + let parent_keys = self.key_map.get(parent_cci)?; let next_child_id = self - .find_next_last_child_of_id(&parent_cci) + .find_next_last_child_of_id(parent_cci) .expect("Can be None only if parent is not present"); let next_cci = parent_cci.nth_child(next_child_id); - let child_keys = father_keys.nth_child(next_child_id); - + let child_keys = parent_keys.nth_child(next_child_id); let account_id = child_keys.account_id(); self.key_map.insert(next_cci.clone(), child_keys); - self.account_id_map.insert(account_id, next_cci); + self.account_id_map.insert(account_id, next_cci.clone()); - Some(account_id) + Some((account_id, next_cci)) + } + + fn find_next_slot_layered(&self) -> ChainIndex { + let mut depth = 1; + + 'outer: loop { + for chain_id in ChainIndex::chain_ids_at_depth_rev(depth) { + if !self.key_map.contains_key(&chain_id) { + break 'outer chain_id; + } + } + depth += 1; + } + } + + pub fn fill_node(&mut self, chain_index: &ChainIndex) -> Option<(nssa::AccountId, ChainIndex)> { + let parent_keys = self.key_map.get(&chain_index.parent()?)?; + let child_id = *chain_index.chain().last()?; + + let child_keys = parent_keys.nth_child(child_id); + let account_id = child_keys.account_id(); + + self.key_map.insert(chain_index.clone(), child_keys); + self.account_id_map.insert(account_id, chain_index.clone()); + + Some((account_id, chain_index.clone())) + } + + pub fn generate_new_node_layered(&mut self) -> Option<(nssa::AccountId, ChainIndex)> { + self.fill_node(&self.find_next_slot_layered()) } pub fn get_node(&self, account_id: nssa::AccountId) -> Option<&N> { @@ -129,11 +168,164 @@ impl KeyTree { self.account_id_map.insert(account_id, chain_index.clone()); self.key_map.insert(chain_index, node); } + + pub fn remove(&mut self, addr: nssa::AccountId) -> Option { + let chain_index = self.account_id_map.remove(&addr).unwrap(); + self.key_map.remove(&chain_index) + } + + /// Populates tree with children. + /// + /// For given `depth` adds children to a tree such that their `ChainIndex::depth(&self) < + /// depth`. + /// + /// Tree must be empty before start + pub fn generate_tree_for_depth(&mut self, depth: u32) { + let mut id_stack = vec![ChainIndex::root()]; + + while let Some(curr_id) = id_stack.pop() { + let mut next_id = curr_id.nth_child(0); + + while (next_id.depth()) < depth { + self.generate_new_node(&curr_id); + id_stack.push(next_id.clone()); + next_id = next_id.next_in_line(); + } + } + } +} + +impl KeyTree { + /// Cleanup of all non-initialized accounts in a private tree + /// + /// For given `depth` checks children to a tree such that their `ChainIndex::depth(&self) < + /// depth`. + /// + /// If account is default, removes them. + /// + /// Chain must be parsed for accounts beforehand + /// + /// Fast, leaves gaps between accounts + pub fn cleanup_tree_remove_uninit_for_depth(&mut self, depth: u32) { + let mut id_stack = vec![ChainIndex::root()]; + + while let Some(curr_id) = id_stack.pop() { + if let Some(node) = self.key_map.get(&curr_id) + && node.value.1 == nssa::Account::default() + && curr_id != ChainIndex::root() + { + let addr = node.account_id(); + self.remove(addr); + } + + let mut next_id = curr_id.nth_child(0); + + while (next_id.depth()) < depth { + id_stack.push(next_id.clone()); + next_id = next_id.next_in_line(); + } + } + } + + /// Cleanup of non-initialized accounts in a private tree + /// + /// If account is default, removes them, stops at first non-default account. + /// + /// Walks through tree in lairs of same depth using `ChainIndex::chain_ids_at_depth()` + /// + /// Chain must be parsed for accounts beforehand + /// + /// Slow, maintains tree consistency. + pub fn cleanup_tree_remove_uninit_layered(&mut self, depth: u32) { + 'outer: for i in (1..(depth as usize)).rev() { + println!("Cleanup of tree at depth {i}"); + for id in ChainIndex::chain_ids_at_depth(i) { + if let Some(node) = self.key_map.get(&id) { + if node.value.1 == nssa::Account::default() { + let addr = node.account_id(); + self.remove(addr); + } else { + break 'outer; + } + } + } + } + } +} + +impl KeyTree { + /// Cleanup of all non-initialized accounts in a public tree + /// + /// For given `depth` checks children to a tree such that their `ChainIndex::depth(&self) < + /// depth`. + /// + /// If account is default, removes them. + /// + /// Fast, leaves gaps between accounts + pub async fn cleanup_tree_remove_ininit_for_depth( + &mut self, + depth: u32, + client: Arc, + ) -> Result<()> { + let mut id_stack = vec![ChainIndex::root()]; + + while let Some(curr_id) = id_stack.pop() { + if let Some(node) = self.key_map.get(&curr_id) { + let address = node.account_id(); + let node_acc = client.get_account(address.to_string()).await?.account; + + if node_acc == nssa::Account::default() && curr_id != ChainIndex::root() { + self.remove(address); + } + } + + let mut next_id = curr_id.nth_child(0); + + while (next_id.depth()) < depth { + id_stack.push(next_id.clone()); + next_id = next_id.next_in_line(); + } + } + + Ok(()) + } + + /// Cleanup of non-initialized accounts in a public tree + /// + /// If account is default, removes them, stops at first non-default account. + /// + /// Walks through tree in lairs of same depth using `ChainIndex::chain_ids_at_depth()` + /// + /// Slow, maintains tree consistency. + pub async fn cleanup_tree_remove_uninit_layered( + &mut self, + depth: u32, + client: Arc, + ) -> Result<()> { + 'outer: for i in (1..(depth as usize)).rev() { + println!("Cleanup of tree at depth {i}"); + for id in ChainIndex::chain_ids_at_depth(i) { + if let Some(node) = self.key_map.get(&id) { + let address = node.account_id(); + let node_acc = client.get_account(address.to_string()).await?.account; + + if node_acc == nssa::Account::default() { + let addr = node.account_id(); + self.remove(addr); + } else { + break 'outer; + } + } + } + } + + Ok(()) + } } #[cfg(test)] mod tests { - use std::str::FromStr; + use std::{collections::HashSet, str::FromStr}; use nssa::AccountId; @@ -162,7 +354,7 @@ mod tests { fn test_small_key_tree() { let seed_holder = seed_holder_for_tests(); - let mut tree = KeyTreePublic::new(&seed_holder); + let mut tree = KeyTreePrivate::new(&seed_holder); let next_last_child_for_parent_id = tree .find_next_last_child_of_id(&ChainIndex::root()) @@ -170,7 +362,7 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 0); - tree.generate_new_node(ChainIndex::root()).unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); assert!( tree.key_map @@ -183,12 +375,12 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 1); - tree.generate_new_node(ChainIndex::root()).unwrap(); - tree.generate_new_node(ChainIndex::root()).unwrap(); - tree.generate_new_node(ChainIndex::root()).unwrap(); - tree.generate_new_node(ChainIndex::root()).unwrap(); - tree.generate_new_node(ChainIndex::root()).unwrap(); - tree.generate_new_node(ChainIndex::root()).unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); let next_last_child_for_parent_id = tree .find_next_last_child_of_id(&ChainIndex::root()) @@ -201,7 +393,7 @@ mod tests { fn test_key_tree_can_not_make_child_keys() { let seed_holder = seed_holder_for_tests(); - let mut tree = KeyTreePublic::new(&seed_holder); + let mut tree = KeyTreePrivate::new(&seed_holder); let next_last_child_for_parent_id = tree .find_next_last_child_of_id(&ChainIndex::root()) @@ -209,7 +401,7 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 0); - tree.generate_new_node(ChainIndex::root()).unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); assert!( tree.key_map @@ -222,7 +414,7 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 1); - let key_opt = tree.generate_new_node(ChainIndex::from_str("/3").unwrap()); + let key_opt = tree.generate_new_node(&ChainIndex::from_str("/3").unwrap()); assert_eq!(key_opt, None); } @@ -239,7 +431,7 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 0); - tree.generate_new_node(ChainIndex::root()).unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); assert!( tree.key_map @@ -252,7 +444,7 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 1); - tree.generate_new_node(ChainIndex::root()).unwrap(); + tree.generate_new_node(&ChainIndex::root()).unwrap(); assert!( tree.key_map @@ -265,7 +457,7 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 2); - tree.generate_new_node(ChainIndex::from_str("/0").unwrap()) + tree.generate_new_node(&ChainIndex::from_str("/0").unwrap()) .unwrap(); let next_last_child_for_parent_id = tree @@ -279,7 +471,7 @@ mod tests { .contains_key(&ChainIndex::from_str("/0/0").unwrap()) ); - tree.generate_new_node(ChainIndex::from_str("/0").unwrap()) + tree.generate_new_node(&ChainIndex::from_str("/0").unwrap()) .unwrap(); let next_last_child_for_parent_id = tree @@ -293,7 +485,7 @@ mod tests { .contains_key(&ChainIndex::from_str("/0/1").unwrap()) ); - tree.generate_new_node(ChainIndex::from_str("/0").unwrap()) + tree.generate_new_node(&ChainIndex::from_str("/0").unwrap()) .unwrap(); let next_last_child_for_parent_id = tree @@ -307,7 +499,7 @@ mod tests { .contains_key(&ChainIndex::from_str("/0/2").unwrap()) ); - tree.generate_new_node(ChainIndex::from_str("/0/1").unwrap()) + tree.generate_new_node(&ChainIndex::from_str("/0/1").unwrap()) .unwrap(); assert!( @@ -321,4 +513,94 @@ mod tests { assert_eq!(next_last_child_for_parent_id, 1); } + + #[test] + fn test_tree_balancing_automatic() { + let seed_holder = seed_holder_for_tests(); + + let mut tree = KeyTreePublic::new(&seed_holder); + + for _ in 0..100 { + tree.generate_new_node_layered().unwrap(); + } + + let next_slot = tree.find_next_slot_layered(); + + assert_eq!(next_slot, ChainIndex::from_str("/0/0/2/1").unwrap()); + } + + #[test] + fn test_cleanup() { + let seed_holder = seed_holder_for_tests(); + + let mut tree = KeyTreePrivate::new(&seed_holder); + tree.generate_tree_for_depth(10); + + let acc = tree + .key_map + .get_mut(&ChainIndex::from_str("/1").unwrap()) + .unwrap(); + acc.value.1.balance = 2; + + let acc = tree + .key_map + .get_mut(&ChainIndex::from_str("/2").unwrap()) + .unwrap(); + acc.value.1.balance = 3; + + let acc = tree + .key_map + .get_mut(&ChainIndex::from_str("/0/1").unwrap()) + .unwrap(); + acc.value.1.balance = 5; + + let acc = tree + .key_map + .get_mut(&ChainIndex::from_str("/1/0").unwrap()) + .unwrap(); + acc.value.1.balance = 6; + + tree.cleanup_tree_remove_uninit_layered(10); + + let mut key_set_res = HashSet::new(); + key_set_res.insert("/0".to_string()); + key_set_res.insert("/1".to_string()); + key_set_res.insert("/2".to_string()); + key_set_res.insert("/".to_string()); + key_set_res.insert("/0/0".to_string()); + key_set_res.insert("/0/1".to_string()); + key_set_res.insert("/1/0".to_string()); + + let mut key_set = HashSet::new(); + + for key in tree.key_map.keys() { + key_set.insert(key.to_string()); + } + + assert_eq!(key_set, key_set_res); + + let acc = tree + .key_map + .get(&ChainIndex::from_str("/1").unwrap()) + .unwrap(); + assert_eq!(acc.value.1.balance, 2); + + let acc = tree + .key_map + .get(&ChainIndex::from_str("/2").unwrap()) + .unwrap(); + assert_eq!(acc.value.1.balance, 3); + + let acc = tree + .key_map + .get(&ChainIndex::from_str("/0/1").unwrap()) + .unwrap(); + assert_eq!(acc.value.1.balance, 5); + + let acc = tree + .key_map + .get(&ChainIndex::from_str("/1/0").unwrap()) + .unwrap(); + assert_eq!(acc.value.1.balance, 6); + } } diff --git a/key_protocol/src/key_protocol_core/mod.rs b/key_protocol/src/key_protocol_core/mod.rs index ac5ee481..b46c46c9 100644 --- a/key_protocol/src/key_protocol_core/mod.rs +++ b/key_protocol/src/key_protocol_core/mod.rs @@ -89,9 +89,18 @@ impl NSSAUserData { /// Returns the account_id of new account pub fn generate_new_public_transaction_private_key( &mut self, - parent_cci: ChainIndex, - ) -> nssa::AccountId { - self.public_key_tree.generate_new_node(parent_cci).unwrap() + parent_cci: Option, + ) -> (nssa::AccountId, ChainIndex) { + match parent_cci { + Some(parent_cci) => self + .public_key_tree + .generate_new_node(&parent_cci) + .expect("Parent must be present in a tree"), + None => self + .public_key_tree + .generate_new_node_layered() + .expect("Search for new node slot failed"), + } } /// Returns the signing key for public transaction signatures @@ -113,9 +122,18 @@ impl NSSAUserData { /// Returns the account_id of new account pub fn generate_new_privacy_preserving_transaction_key_chain( &mut self, - parent_cci: ChainIndex, - ) -> nssa::AccountId { - self.private_key_tree.generate_new_node(parent_cci).unwrap() + parent_cci: Option, + ) -> (nssa::AccountId, ChainIndex) { + match parent_cci { + Some(parent_cci) => self + .private_key_tree + .generate_new_node(&parent_cci) + .expect("Parent must be present in a tree"), + None => self + .private_key_tree + .generate_new_node_layered() + .expect("Search for new node slot failed"), + } } /// Returns the signing key for public transaction signatures @@ -169,16 +187,8 @@ mod tests { fn test_new_account() { let mut user_data = NSSAUserData::default(); - let account_id_pub = - user_data.generate_new_public_transaction_private_key(ChainIndex::root()); - let account_id_private = - user_data.generate_new_privacy_preserving_transaction_key_chain(ChainIndex::root()); - - let is_private_key_generated = user_data - .get_pub_account_signing_key(&account_id_pub) - .is_some(); - - assert!(is_private_key_generated); + let (account_id_private, _) = user_data + .generate_new_privacy_preserving_transaction_key_chain(Some(ChainIndex::root())); let is_key_chain_generated = user_data.get_private_account(&account_id_private).is_some(); diff --git a/nssa/core/Cargo.toml b/nssa/core/Cargo.toml index 0e16a3f5..80fe7df9 100644 --- a/nssa/core/Cargo.toml +++ b/nssa/core/Cargo.toml @@ -6,7 +6,7 @@ edition = "2024" [dependencies] risc0-zkvm = { version = "3.0.3", features = ['std'] } serde = { version = "1.0", default-features = false } -thiserror = { version = "2.0.12", optional = true } +thiserror = { version = "2.0.12" } bytemuck = { version = "1.13", optional = true } chacha20 = { version = "0.9", default-features = false } k256 = { version = "0.13.3", optional = true } @@ -14,6 +14,9 @@ base58 = { version = "0.2.0", optional = true } anyhow = { version = "1.0.98", optional = true } borsh = "1.5.7" +[dev-dependencies] +serde_json = "1.0.81" + [features] default = [] -host = ["thiserror", "bytemuck", "k256", "base58", "anyhow"] +host = ["dep:bytemuck", "dep:k256", "dep:base58", "dep:anyhow"] diff --git a/nssa/core/src/account.rs b/nssa/core/src/account.rs index f32d05d0..c152581d 100644 --- a/nssa/core/src/account.rs +++ b/nssa/core/src/account.rs @@ -4,12 +4,14 @@ use std::{fmt::Display, str::FromStr}; #[cfg(feature = "host")] use base58::{FromBase58, ToBase58}; use borsh::{BorshDeserialize, BorshSerialize}; +pub use data::Data; use serde::{Deserialize, Serialize}; use crate::program::ProgramId; +pub mod data; + pub type Nonce = u128; -pub type Data = Vec; /// Account to be used both in public and private contexts #[derive( @@ -23,8 +25,8 @@ pub struct Account { pub nonce: Nonce, } -#[derive(Serialize, Deserialize, Clone)] -#[cfg_attr(any(feature = "host", test), derive(Debug, PartialEq, Eq))] +#[derive(Serialize, Deserialize, Clone, PartialEq, Eq)] +#[cfg_attr(any(feature = "host", test), derive(Debug))] pub struct AccountWithMetadata { pub account: Account, pub is_authorized: bool, @@ -139,7 +141,10 @@ mod tests { let account = Account { program_owner: [1, 2, 3, 4, 5, 6, 7, 8], balance: 1337, - data: b"testing_account_with_metadata_constructor".to_vec(), + data: b"testing_account_with_metadata_constructor" + .to_vec() + .try_into() + .unwrap(), nonce: 0xdeadbeef, }; let fingerprint = AccountId::new([8; 32]); diff --git a/nssa/core/src/account/data.rs b/nssa/core/src/account/data.rs new file mode 100644 index 00000000..974cb06c --- /dev/null +++ b/nssa/core/src/account/data.rs @@ -0,0 +1,174 @@ +use std::ops::Deref; + +use borsh::{BorshDeserialize, BorshSerialize}; +use serde::{Deserialize, Serialize}; + +pub const DATA_MAX_LENGTH_IN_BYTES: usize = 100 * 1024; // 100 KiB + +#[derive(Default, Clone, PartialEq, Eq, Serialize, BorshSerialize)] +#[cfg_attr(any(feature = "host", test), derive(Debug))] +pub struct Data(Vec); + +impl Data { + pub fn into_inner(self) -> Vec { + self.0 + } + + #[cfg(feature = "host")] + pub fn from_cursor( + cursor: &mut std::io::Cursor<&[u8]>, + ) -> Result { + use std::io::Read as _; + + let mut u32_bytes = [0u8; 4]; + cursor.read_exact(&mut u32_bytes)?; + let data_length = u32::from_le_bytes(u32_bytes); + if data_length as usize > DATA_MAX_LENGTH_IN_BYTES { + return Err( + std::io::Error::new(std::io::ErrorKind::InvalidData, DataTooBigError).into(), + ); + } + + let mut data = vec![0; data_length as usize]; + cursor.read_exact(&mut data)?; + Ok(Self(data)) + } +} + +#[derive(Debug, thiserror::Error, Clone, Copy, PartialEq, Eq)] +#[error("data length exceeds maximum allowed length of {DATA_MAX_LENGTH_IN_BYTES} bytes")] +pub struct DataTooBigError; + +impl From for Vec { + fn from(data: Data) -> Self { + data.0 + } +} + +impl TryFrom> for Data { + type Error = DataTooBigError; + + fn try_from(value: Vec) -> Result { + if value.len() > DATA_MAX_LENGTH_IN_BYTES { + Err(DataTooBigError) + } else { + Ok(Self(value)) + } + } +} + +impl Deref for Data { + type Target = [u8]; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +impl AsRef<[u8]> for Data { + fn as_ref(&self) -> &[u8] { + &self.0 + } +} + +impl<'de> Deserialize<'de> for Data { + fn deserialize(deserializer: D) -> Result + where + D: serde::Deserializer<'de>, + { + /// Data deserialization visitor. + /// + /// Compared to a simple deserialization into a `Vec`, this visitor enforces + /// early length check defined by [`DATA_MAX_LENGTH_IN_BYTES`]. + struct DataVisitor; + + impl<'de> serde::de::Visitor<'de> for DataVisitor { + type Value = Data; + + fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + write!( + formatter, + "a byte array with length not exceeding {} bytes", + DATA_MAX_LENGTH_IN_BYTES + ) + } + + fn visit_seq(self, mut seq: A) -> Result + where + A: serde::de::SeqAccess<'de>, + { + let mut vec = + Vec::with_capacity(seq.size_hint().unwrap_or(0).min(DATA_MAX_LENGTH_IN_BYTES)); + + while let Some(value) = seq.next_element()? { + if vec.len() >= DATA_MAX_LENGTH_IN_BYTES { + return Err(serde::de::Error::custom(DataTooBigError)); + } + vec.push(value); + } + + Ok(Data(vec)) + } + } + + deserializer.deserialize_seq(DataVisitor) + } +} + +impl BorshDeserialize for Data { + fn deserialize_reader(reader: &mut R) -> std::io::Result { + // Implementation adapted from `impl BorshDeserialize for Vec` + + let len = u32::deserialize_reader(reader)?; + match len { + 0 => Ok(Self::default()), + len if len as usize > DATA_MAX_LENGTH_IN_BYTES => Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + DataTooBigError, + )), + len => { + let vec_bytes = u8::vec_from_reader(len, reader)? + .expect("can't be None in current borsh crate implementation"); + Ok(Self(vec_bytes)) + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_data_max_length_allowed() { + let max_vec = vec![0u8; DATA_MAX_LENGTH_IN_BYTES]; + let result = Data::try_from(max_vec); + assert!(result.is_ok()); + } + + #[test] + fn test_data_too_big_error() { + let big_vec = vec![0u8; DATA_MAX_LENGTH_IN_BYTES + 1]; + let result = Data::try_from(big_vec); + assert!(matches!(result, Err(DataTooBigError))); + } + + #[test] + fn test_borsh_deserialize_exceeding_limit_error() { + let too_big_data = vec![0u8; DATA_MAX_LENGTH_IN_BYTES + 1]; + let mut serialized = Vec::new(); + <_ as BorshSerialize>::serialize(&too_big_data, &mut serialized).unwrap(); + + let result = ::deserialize(&mut serialized.as_ref()); + assert!(result.is_err()); + } + + #[test] + fn test_json_deserialize_exceeding_limit_error() { + let data = vec![0u8; DATA_MAX_LENGTH_IN_BYTES + 1]; + let json = serde_json::to_string(&data).unwrap(); + + let result: Result = serde_json::from_str(&json); + assert!(result.is_err()); + } +} diff --git a/nssa/core/src/circuit_io.rs b/nssa/core/src/circuit_io.rs index e1afe106..848fe3e6 100644 --- a/nssa/core/src/circuit_io.rs +++ b/nssa/core/src/circuit_io.rs @@ -10,7 +10,7 @@ use crate::{ #[derive(Serialize, Deserialize)] pub struct PrivacyPreservingCircuitInput { - pub program_output: ProgramOutput, + pub program_outputs: Vec, pub visibility_mask: Vec, pub private_account_nonces: Vec, pub private_account_keys: Vec<(NullifierPublicKey, SharedSecretKey)>, @@ -54,7 +54,7 @@ mod tests { Account { program_owner: [1, 2, 3, 4, 5, 6, 7, 8], balance: 12345678901234567890, - data: b"test data".to_vec(), + data: b"test data".to_vec().try_into().unwrap(), nonce: 18446744073709551614, }, true, @@ -64,7 +64,7 @@ mod tests { Account { program_owner: [9, 9, 9, 8, 8, 8, 7, 7], balance: 123123123456456567112, - data: b"test data".to_vec(), + data: b"test data".to_vec().try_into().unwrap(), nonce: 9999999999999999999999, }, false, @@ -74,7 +74,7 @@ mod tests { public_post_states: vec![Account { program_owner: [1, 2, 3, 4, 5, 6, 7, 8], balance: 100, - data: b"post state data".to_vec(), + data: b"post state data".to_vec().try_into().unwrap(), nonce: 18446744073709551615, }], ciphertexts: vec![Ciphertext(vec![255, 255, 1, 1, 2, 2])], diff --git a/nssa/core/src/encoding.rs b/nssa/core/src/encoding.rs index 3a8a1285..24ac050c 100644 --- a/nssa/core/src/encoding.rs +++ b/nssa/core/src/encoding.rs @@ -26,12 +26,14 @@ impl Account { bytes.extend_from_slice(&self.nonce.to_le_bytes()); let data_length: u32 = self.data.len() as u32; bytes.extend_from_slice(&data_length.to_le_bytes()); - bytes.extend_from_slice(self.data.as_slice()); + bytes.extend_from_slice(self.data.as_ref()); bytes } #[cfg(feature = "host")] pub fn from_cursor(cursor: &mut Cursor<&[u8]>) -> Result { + use crate::account::data::Data; + let mut u32_bytes = [0u8; 4]; let mut u128_bytes = [0u8; 16]; @@ -51,10 +53,7 @@ impl Account { let nonce = u128::from_le_bytes(u128_bytes); // data - cursor.read_exact(&mut u32_bytes)?; - let data_length = u32::from_le_bytes(u32_bytes); - let mut data = vec![0; data_length as usize]; - cursor.read_exact(&mut data)?; + let data = Data::from_cursor(cursor)?; Ok(Self { program_owner, @@ -149,7 +148,7 @@ mod tests { program_owner: [1, 2, 3, 4, 5, 6, 7, 8], balance: 123456789012345678901234567890123456, nonce: 42, - data: b"hola mundo".to_vec(), + data: b"hola mundo".to_vec().try_into().unwrap(), }; // program owner || balance || nonce || data_len || data @@ -210,7 +209,7 @@ mod tests { program_owner: [1, 2, 3, 4, 5, 6, 7, 8], balance: 123456789012345678901234567890123456, nonce: 42, - data: b"hola mundo".to_vec(), + data: b"hola mundo".to_vec().try_into().unwrap(), }; let bytes = account.to_bytes(); let mut cursor = Cursor::new(bytes.as_ref()); diff --git a/nssa/core/src/program.rs b/nssa/core/src/program.rs index 2d38fa4d..26ee8deb 100644 --- a/nssa/core/src/program.rs +++ b/nssa/core/src/program.rs @@ -1,3 +1,5 @@ +use std::collections::HashSet; + use risc0_zkvm::{DeserializeOwned, guest::env, serde::Deserializer}; use serde::{Deserialize, Serialize}; @@ -8,6 +10,7 @@ use crate::account::{Account, AccountWithMetadata}; pub type ProgramId = [u32; 8]; pub type InstructionData = Vec; pub const DEFAULT_PROGRAM_ID: ProgramId = [0; 8]; +pub const MAX_NUMBER_CHAINED_CALLS: usize = 10; pub struct ProgramInput { pub pre_states: Vec, @@ -54,7 +57,9 @@ impl From<(&ProgramId, &PdaSeed)> for AccountId { #[derive(Serialize, Deserialize, Clone)] #[cfg_attr(any(feature = "host", test), derive(Debug, PartialEq, Eq))] pub struct ChainedCall { + /// The program ID of the program to execute pub program_id: ProgramId, + /// The instruction data to pass pub instruction_data: InstructionData, pub pre_states: Vec, pub pda_seeds: Vec, @@ -111,26 +116,34 @@ impl AccountPostState { #[derive(Serialize, Deserialize, Clone)] #[cfg_attr(any(feature = "host", test), derive(Debug, PartialEq, Eq))] pub struct ProgramOutput { + /// The instruction data the program received to produce this output + pub instruction_data: InstructionData, + /// The account pre states the program received to produce this output pub pre_states: Vec, pub post_states: Vec, pub chained_calls: Vec, } -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( + instruction_data: InstructionData, pre_states: Vec, post_states: Vec, ) { let output = ProgramOutput { + instruction_data, pre_states, post_states, chained_calls: Vec::new(), @@ -139,11 +152,13 @@ pub fn write_nssa_outputs( } pub fn write_nssa_outputs_with_chained_call( + instruction_data: InstructionData, pre_states: Vec, post_states: Vec, chained_calls: Vec, ) { let output = ProgramOutput { + instruction_data, pre_states, post_states, chained_calls, @@ -162,32 +177,37 @@ pub fn validate_execution( post_states: &[AccountPostState], executing_program_id: ProgramId, ) -> bool { - // 1. Lengths must match + // 1. Check account ids are all different + if !validate_uniqueness_of_account_ids(pre_states) { + return false; + } + + // 2. Lengths must match if pre_states.len() != post_states.len() { return false; } for (pre, post) in pre_states.iter().zip(post_states) { - // 2. Nonce must remain unchanged + // 3. Nonce must remain unchanged if pre.account.nonce != post.account.nonce { return false; } - // 3. Program ownership changes are not allowed + // 4. Program ownership changes are not allowed if pre.account.program_owner != post.account.program_owner { return false; } let account_program_owner = pre.account.program_owner; - // 4. Decreasing balance only allowed if owned by executing program + // 5. Decreasing balance only allowed if owned by executing program if post.account.balance < pre.account.balance && account_program_owner != executing_program_id { return false; } - // 5. Data changes only allowed if owned by executing program or if account pre state has + // 6. Data changes only allowed if owned by executing program or if account pre state has // default values if pre.account.data != post.account.data && pre.account != Account::default() @@ -196,16 +216,27 @@ pub fn validate_execution( return false; } - // 6. If a post state has default program owner, the pre state must have been a default + // 7. If a post state has default program owner, the pre state must have been a default // account if post.account.program_owner == DEFAULT_PROGRAM_ID && pre.account != Account::default() { return false; } } - // 7. Total balance is preserved - let total_balance_pre_states: u128 = pre_states.iter().map(|pre| pre.account.balance).sum(); - let total_balance_post_states: u128 = post_states.iter().map(|post| post.account.balance).sum(); + // 8. Total balance is preserved + + let Some(total_balance_pre_states) = + WrappedBalanceSum::from_balances(pre_states.iter().map(|pre| pre.account.balance)) + else { + return false; + }; + + let Some(total_balance_post_states) = + WrappedBalanceSum::from_balances(post_states.iter().map(|post| post.account.balance)) + else { + return false; + }; + if total_balance_pre_states != total_balance_post_states { return false; } @@ -213,6 +244,44 @@ pub fn validate_execution( true } +fn validate_uniqueness_of_account_ids(pre_states: &[AccountWithMetadata]) -> bool { + let number_of_accounts = pre_states.len(); + let number_of_account_ids = pre_states + .iter() + .map(|account| &account.account_id) + .collect::>() + .len(); + + number_of_accounts == number_of_account_ids +} + +/// Representation of a number as `lo + hi * 2^128`. +#[derive(PartialEq, Eq)] +struct WrappedBalanceSum { + lo: u128, + hi: u128, +} + +impl WrappedBalanceSum { + /// Constructs a [`WrappedBalanceSum`] from an iterator of balances. + /// + /// Returns [`None`] if balance sum overflows `lo + hi * 2^128` representation, which is not + /// expected in practical scenarios. + fn from_balances(balances: impl Iterator) -> Option { + let mut wrapped = WrappedBalanceSum { lo: 0, hi: 0 }; + + for balance in balances { + let (new_sum, did_overflow) = wrapped.lo.overflowing_add(balance); + if did_overflow { + wrapped.hi = wrapped.hi.checked_add(1)?; + } + wrapped.lo = new_sum; + } + + Some(wrapped) + } +} + #[cfg(test)] mod tests { use super::*; @@ -222,7 +291,7 @@ mod tests { let account = Account { program_owner: [1, 2, 3, 4, 5, 6, 7, 8], balance: 1337, - data: vec![0xde, 0xad, 0xbe, 0xef], + data: vec![0xde, 0xad, 0xbe, 0xef].try_into().unwrap(), nonce: 10, }; @@ -237,7 +306,7 @@ mod tests { let account = Account { program_owner: [1, 2, 3, 4, 5, 6, 7, 8], balance: 1337, - data: vec![0xde, 0xad, 0xbe, 0xef], + data: vec![0xde, 0xad, 0xbe, 0xef].try_into().unwrap(), nonce: 10, }; @@ -252,7 +321,7 @@ mod tests { let mut account = Account { program_owner: [1, 2, 3, 4, 5, 6, 7, 8], balance: 1337, - data: vec![0xde, 0xad, 0xbe, 0xef], + data: vec![0xde, 0xad, 0xbe, 0xef].try_into().unwrap(), nonce: 10, }; diff --git a/nssa/program_methods/guest/Cargo.lock b/nssa/program_methods/guest/Cargo.lock index 563e8b96..2c293ecb 100644 --- a/nssa/program_methods/guest/Cargo.lock +++ b/nssa/program_methods/guest/Cargo.lock @@ -1578,6 +1578,7 @@ dependencies = [ "chacha20", "risc0-zkvm", "serde", + "thiserror", ] [[package]] diff --git a/nssa/program_methods/guest/src/bin/authenticated_transfer.rs b/nssa/program_methods/guest/src/bin/authenticated_transfer.rs index 50afa507..fe02d065 100644 --- a/nssa/program_methods/guest/src/bin/authenticated_transfer.rs +++ b/nssa/program_methods/guest/src/bin/authenticated_transfer.rs @@ -6,34 +6,37 @@ use nssa_core::{ }; /// Initializes a default account under the ownership of this program. -fn initialize_account(pre_state: AccountWithMetadata) { +fn initialize_account(pre_state: AccountWithMetadata) -> AccountPostState { let account_to_claim = AccountPostState::new_claimed(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() != &Account::default() { - return; + panic!("Account must be uninitialized"); } // 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]); + 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 { // 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 @@ -57,23 +60,31 @@ fn transfer(sender: AccountWithMetadata, recipient: AccountWithMetadata, balance } }; - write_nssa_outputs(vec![sender, recipient], vec![sender_post, recipient_post]); + 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()), + let post_states = match (pre_states.as_slice(), balance_to_move) { + ([account_to_claim], 0) => { + let post = initialize_account(account_to_claim.clone()); + vec![post] + } ([sender, recipient], balance_to_move) => { transfer(sender.clone(), recipient.clone(), balance_to_move) } _ => panic!("invalid params"), - } + }; + + write_nssa_outputs(instruction_words, pre_states, post_states); } diff --git a/nssa/program_methods/guest/src/bin/pinata.rs b/nssa/program_methods/guest/src/bin/pinata.rs index 50aac7b0..a0c46a1a 100644 --- a/nssa/program_methods/guest/src/bin/pinata.rs +++ b/nssa/program_methods/guest/src/bin/pinata.rs @@ -44,10 +44,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_words, + ) = read_nssa_inputs::(); let [pinata, winner] = match pre_states.try_into() { Ok(array) => array, @@ -63,10 +66,15 @@ fn main() { let mut pinata_post = pinata.account.clone(); let mut winner_post = winner.account.clone(); pinata_post.balance -= PRIZE; - pinata_post.data = data.next_data().to_vec(); + pinata_post.data = data + .next_data() + .to_vec() + .try_into() + .expect("33 bytes should fit into Data"); winner_post.balance += PRIZE; write_nssa_outputs( + instruction_words, vec![pinata, winner], vec![ AccountPostState::new(pinata_post), diff --git a/nssa/program_methods/guest/src/bin/pinata_token.rs b/nssa/program_methods/guest/src/bin/pinata_token.rs index be661c2b..f988be9e 100644 --- a/nssa/program_methods/guest/src/bin/pinata_token.rs +++ b/nssa/program_methods/guest/src/bin/pinata_token.rs @@ -1,8 +1,14 @@ -use nssa_core::program::{ - read_nssa_inputs, write_nssa_outputs_with_chained_call, AccountPostState, ChainedCall, PdaSeed, ProgramInput +use nssa_core::{ + account::Data, + program::{ + AccountPostState, ChainedCall, PdaSeed, ProgramInput, read_nssa_inputs, + write_nssa_outputs_with_chained_call, + }, +}; +use risc0_zkvm::{ + serde::to_vec, + sha::{Impl, Sha256}, }; -use risc0_zkvm::serde::to_vec; -use risc0_zkvm::sha::{Impl, Sha256}; const PRIZE: u128 = 150; @@ -35,22 +41,26 @@ impl Challenge { digest[..difficulty].iter().all(|&b| b == 0) } - fn next_data(self) -> [u8; 33] { + fn next_data(self) -> Data { let mut result = [0; 33]; result[0] = self.difficulty; result[1..].copy_from_slice(Impl::hash_bytes(&self.seed).as_bytes()); - result + result.to_vec().try_into().expect("should fit") } } /// A pinata program fn main() { // Read input accounts. - // It is expected to receive three accounts: [pinata_definition, pinata_token_holding, winner_token_holding] - let ProgramInput { - pre_states, - instruction: solution, - } = read_nssa_inputs::(); + // It is expected to receive three accounts: [pinata_definition, pinata_token_holding, + // winner_token_holding] + let ( + ProgramInput { + pre_states, + instruction: solution, + }, + instruction_words, + ) = read_nssa_inputs::(); let [ pinata_definition, @@ -70,7 +80,7 @@ fn main() { let mut pinata_definition_post = pinata_definition.account.clone(); let pinata_token_holding_post = pinata_token_holding.account.clone(); let winner_token_holding_post = winner_token_holding.account.clone(); - pinata_definition_post.data = data.next_data().to_vec(); + pinata_definition_post.data = data.next_data(); let mut instruction_data: [u8; 23] = [0; 23]; instruction_data[0] = 1; @@ -83,11 +93,15 @@ fn main() { let chained_calls = vec![ChainedCall { program_id: pinata_token_holding_post.program_owner, instruction_data: to_vec(&instruction_data).unwrap(), - pre_states: vec![pinata_token_holding_for_chain_call, winner_token_holding.clone()], + pre_states: vec![ + pinata_token_holding_for_chain_call, + winner_token_holding.clone(), + ], pda_seeds: vec![PdaSeed::new([0; 32])], }]; write_nssa_outputs_with_chained_call( + instruction_words, vec![ pinata_definition, pinata_token_holding, 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 7813fa58..29162db9 100644 --- a/nssa/program_methods/guest/src/bin/privacy_preserving_circuit.rs +++ b/nssa/program_methods/guest/src/bin/privacy_preserving_circuit.rs @@ -1,49 +1,119 @@ -use std::collections::HashSet; - -use risc0_zkvm::{guest::env, serde::to_vec}; +use std::collections::HashMap; use nssa_core::{ - Commitment, CommitmentSetDigest, DUMMY_COMMITMENT_HASH, EncryptionScheme, - Nullifier, NullifierPublicKey, PrivacyPreservingCircuitInput, PrivacyPreservingCircuitOutput, + Commitment, CommitmentSetDigest, DUMMY_COMMITMENT_HASH, EncryptionScheme, Nullifier, + NullifierPublicKey, PrivacyPreservingCircuitInput, PrivacyPreservingCircuitOutput, account::{Account, AccountId, AccountWithMetadata}, compute_digest_for_path, encryption::Ciphertext, - program::{DEFAULT_PROGRAM_ID, ProgramOutput, validate_execution}, + program::{DEFAULT_PROGRAM_ID, MAX_NUMBER_CHAINED_CALLS, validate_execution}, }; +use risc0_zkvm::{guest::env, serde::to_vec}; fn main() { let PrivacyPreservingCircuitInput { - program_output, + program_outputs, visibility_mask, private_account_nonces, private_account_keys, private_account_auth, - program_id, + mut program_id, } = env::read(); - // Check that `program_output` is consistent with the execution of the corresponding program. - env::verify(program_id, &to_vec(&program_output).unwrap()).unwrap(); + let mut pre_states: Vec = Vec::new(); + let mut state_diff: HashMap = HashMap::new(); - let ProgramOutput { - pre_states, - post_states, - chained_calls, - } = program_output; - - // TODO: implement chained calls for privacy preserving transactions - if !chained_calls.is_empty() { - panic!("Privacy preserving transactions do not support yet chained calls.") + let num_calls = program_outputs.len(); + if num_calls > MAX_NUMBER_CHAINED_CALLS { + panic!("Max chained calls depth is exceeded"); } - // Check that there are no repeated account ids - if !validate_uniqueness_of_account_ids(&pre_states) { - panic!("Repeated account ids found") + let Some(last_program_call) = program_outputs.last() else { + panic!("Program outputs is empty") + }; + + if !last_program_call.chained_calls.is_empty() { + panic!("Call stack is incomplete"); } - // Check that the program is well behaved. - // See the # Programs section for the definition of the `validate_execution` method. - if !validate_execution(&pre_states, &post_states, program_id) { - panic!("Bad behaved program"); + for window in program_outputs.windows(2) { + let caller = &window[0]; + let callee = &window[1]; + + if caller.chained_calls.len() > 1 { + panic!("Privacy Multi-chained calls are not supported yet"); + } + + // TODO: Modify when multi-chain calls are supported in the circuit + let Some(caller_chained_call) = &caller.chained_calls.first() else { + panic!("Expected chained call"); + }; + + // Check that instruction data in caller is the instruction data in callee + if caller_chained_call.instruction_data != callee.instruction_data { + panic!("Invalid instruction data"); + } + + // Check that account pre_states in caller are the ones in calle + if caller_chained_call.pre_states != callee.pre_states { + panic!("Invalid pre states"); + } + } + + for (i, program_output) in program_outputs.iter().enumerate() { + let mut program_output = program_output.clone(); + + // Check that `program_output` is consistent with the execution of the corresponding program. + let program_output_words = + &to_vec(&program_output).expect("program_output must be serializable"); + env::verify(program_id, program_output_words) + .expect("program output must match the program's execution"); + + // Check that the program is well behaved. + // See the # Programs section for the definition of the `validate_execution` method. + if !validate_execution( + &program_output.pre_states, + &program_output.post_states, + program_id, + ) { + panic!("Bad behaved program"); + } + + // The invoked program claims the accounts with default program id. + for post in program_output + .post_states + .iter_mut() + .filter(|post| post.requires_claim()) + { + // The invoked program can only claim accounts with default program id. + if post.account().program_owner == DEFAULT_PROGRAM_ID { + post.account_mut().program_owner = program_id; + } else { + panic!("Cannot claim an initialized account") + } + } + + for (pre, post) in program_output + .pre_states + .iter() + .zip(&program_output.post_states) + { + 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.account().clone()); + } + + // TODO: Modify when multi-chain calls are supported in the circuit + if let Some(next_chained_call) = &program_output.chained_calls.first() { + program_id = next_chained_call.program_id; + } else if i != program_outputs.len() - 1 { + panic!("Inner call without a chained call found") + }; } let n_accounts = pre_states.len(); @@ -70,10 +140,8 @@ fn main() { // Public account public_pre_states.push(pre_states[i].clone()); - let mut post = post_states[i].account().clone(); - if pre_states[i].is_authorized { - post.nonce += 1; - } + let mut post = state_diff.get(&pre_states[i].account_id).unwrap().clone(); + if post.program_owner == DEFAULT_PROGRAM_ID { // Claim account post.program_owner = program_id; @@ -126,7 +194,8 @@ fn main() { } // Update post-state with new nonce - let mut post_with_updated_values = post_states[i].account().clone(); + let mut post_with_updated_values = + state_diff.get(&pre_states[i].account_id).unwrap().clone(); post_with_updated_values.nonce = *new_nonce; if post_with_updated_values.program_owner == DEFAULT_PROGRAM_ID { @@ -175,14 +244,3 @@ fn main() { env::commit(&output); } - -fn validate_uniqueness_of_account_ids(pre_states: &[AccountWithMetadata]) -> bool { - let number_of_accounts = pre_states.len(); - let number_of_account_ids = pre_states - .iter() - .map(|account| account.account_id.clone()) - .collect::>() - .len(); - - number_of_accounts == number_of_account_ids -} diff --git a/nssa/src/privacy_preserving_transaction/circuit.rs b/nssa/src/privacy_preserving_transaction/circuit.rs index eeba6928..95933a32 100644 --- a/nssa/src/privacy_preserving_transaction/circuit.rs +++ b/nssa/src/privacy_preserving_transaction/circuit.rs @@ -1,9 +1,11 @@ +use std::collections::HashMap; + use borsh::{BorshDeserialize, BorshSerialize}; use nssa_core::{ MembershipProof, NullifierPublicKey, NullifierSecretKey, PrivacyPreservingCircuitInput, PrivacyPreservingCircuitOutput, SharedSecretKey, account::AccountWithMetadata, - program::{InstructionData, ProgramOutput}, + program::{InstructionData, ProgramId, ProgramOutput}, }; use risc0_zkvm::{ExecutorEnv, InnerReceipt, Receipt, default_prover}; @@ -11,12 +13,35 @@ use crate::{ error::NssaError, program::Program, program_methods::{PRIVACY_PRESERVING_CIRCUIT_ELF, PRIVACY_PRESERVING_CIRCUIT_ID}, + state::MAX_NUMBER_CHAINED_CALLS, }; /// Proof of the privacy preserving execution circuit #[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)] pub struct Proof(pub(crate) Vec); +#[derive(Clone)] +pub struct ProgramWithDependencies { + pub program: Program, + // TODO: avoid having a copy of the bytecode of each dependency. + pub dependencies: HashMap, +} + +impl ProgramWithDependencies { + pub fn new(program: Program, dependencies: HashMap) -> Self { + Self { + program, + dependencies, + } + } +} + +impl From for ProgramWithDependencies { + fn from(program: Program) -> Self { + ProgramWithDependencies::new(program, HashMap::new()) + } +} + /// Generates a proof of the execution of a NSSA program inside the privacy preserving execution /// circuit pub fn execute_and_prove( @@ -26,27 +51,64 @@ pub fn execute_and_prove( private_account_nonces: &[u128], private_account_keys: &[(NullifierPublicKey, SharedSecretKey)], private_account_auth: &[(NullifierSecretKey, MembershipProof)], - program: &Program, + program_with_dependencies: &ProgramWithDependencies, ) -> Result<(PrivacyPreservingCircuitOutput, Proof), NssaError> { - let inner_receipt = execute_and_prove_program(program, pre_states, instruction_data)?; + let mut program = &program_with_dependencies.program; + let dependencies = &program_with_dependencies.dependencies; + let mut instruction_data = instruction_data.clone(); + let mut pre_states = pre_states.to_vec(); + let mut env_builder = ExecutorEnv::builder(); + let mut program_outputs = Vec::new(); - let program_output: ProgramOutput = inner_receipt - .journal - .decode() - .map_err(|e| NssaError::ProgramOutputDeserializationError(e.to_string()))?; + for _i in 0..MAX_NUMBER_CHAINED_CALLS { + let inner_receipt = execute_and_prove_program(program, &pre_states, &instruction_data)?; + + let program_output: ProgramOutput = inner_receipt + .journal + .decode() + .map_err(|e| NssaError::ProgramOutputDeserializationError(e.to_string()))?; + + // TODO: remove clone + program_outputs.push(program_output.clone()); + + // Prove circuit. + env_builder.add_assumption(inner_receipt); + + // TODO: Remove when multi-chain calls are supported in the circuit + assert!(program_output.chained_calls.len() <= 1); + // TODO: Modify when multi-chain calls are supported in the circuit + if let Some(next_call) = program_output.chained_calls.first() { + program = dependencies + .get(&next_call.program_id) + .ok_or(NssaError::InvalidProgramBehavior)?; + instruction_data = next_call.instruction_data.clone(); + // 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.account().clone(); + post_states_with_metadata.push(post_with_metadata); + } + + pre_states = next_call.pre_states.clone(); + } else { + break; + } + } let circuit_input = PrivacyPreservingCircuitInput { - program_output, + program_outputs, visibility_mask: visibility_mask.to_vec(), private_account_nonces: private_account_nonces.to_vec(), private_account_keys: private_account_keys.to_vec(), private_account_auth: private_account_auth.to_vec(), - program_id: program.id(), + program_id: program_with_dependencies.program.id(), }; - // Prove circuit. - let mut env_builder = ExecutorEnv::builder(); - env_builder.add_assumption(inner_receipt); env_builder.write(&circuit_input).unwrap(); let env = env_builder.build().unwrap(); let prover = default_prover(); @@ -95,7 +157,7 @@ impl Proof { mod tests { use nssa_core::{ Commitment, DUMMY_COMMITMENT_HASH, EncryptionScheme, Nullifier, - account::{Account, AccountId, AccountWithMetadata}, + account::{Account, AccountId, AccountWithMetadata, data::Data}, }; use super::*; @@ -133,15 +195,15 @@ mod tests { let expected_sender_post = Account { program_owner: program.id(), balance: 100 - balance_to_move, - nonce: 1, - data: vec![], + nonce: 0, + data: Data::default(), }; let expected_recipient_post = Account { program_owner: program.id(), balance: balance_to_move, nonce: 0xdeadbeef, - data: vec![], + data: Data::default(), }; let expected_sender_pre = sender.clone(); @@ -156,7 +218,7 @@ mod tests { &[0xdeadbeef], &[(recipient_keys.npk(), shared_secret.clone())], &[], - &Program::authenticated_transfer_program(), + &Program::authenticated_transfer_program().into(), ) .unwrap(); @@ -191,7 +253,7 @@ mod tests { balance: 100, nonce: 0xdeadbeef, program_owner: program.id(), - data: vec![], + data: Data::default(), }, true, AccountId::from(&sender_keys.npk()), @@ -257,7 +319,7 @@ mod tests { sender_keys.nsk, commitment_set.get_proof_for(&commitment_sender).unwrap(), )], - &program, + &program.into(), ) .unwrap(); diff --git a/nssa/src/program.rs b/nssa/src/program.rs index b2561304..1865248a 100644 --- a/nssa/src/program.rs +++ b/nssa/src/program.rs @@ -14,7 +14,7 @@ use crate::{ /// TODO: Make this variable when fees are implemented const MAX_NUM_CYCLES_PUBLIC_EXECUTION: u64 = 1024 * 1024 * 32; // 32M cycles -#[derive(Debug, PartialEq, Eq)] +#[derive(Clone, Debug, PartialEq, Eq)] pub struct Program { id: ProgramId, elf: Vec, @@ -221,6 +221,13 @@ mod tests { elf: CLAIMER_ELF.to_vec(), } } + + pub fn modified_transfer_program() -> Self { + use test_program_methods::MODIFIED_TRANSFER_ELF; + // This unwrap won't panic since the `MODIFIED_TRANSFER_ELF` comes from risc0 build of + // `program_methods` + Self::new(MODIFIED_TRANSFER_ELF.to_vec()).unwrap() + } } #[test] diff --git a/nssa/src/state.rs b/nssa/src/state.rs index 409ceca9..86df3a5e 100644 --- a/nssa/src/state.rs +++ b/nssa/src/state.rs @@ -154,6 +154,12 @@ impl V02State { *current_account = post; } + // 5. Increment nonces for public signers + for account_id in tx.signer_account_ids() { + let current_account = self.get_account_by_id_mut(account_id); + current_account.nonce += 1; + } + Ok(()) } @@ -234,7 +240,7 @@ impl V02State { program_owner: Program::pinata().id(), balance: 1500, // Difficulty: 3 - data: vec![3; 33], + data: vec![3; 33].try_into().expect("should fit"), nonce: 0, }, ); @@ -248,7 +254,7 @@ impl V02State { Account { program_owner: Program::pinata_token().id(), // Difficulty: 3 - data: vec![3; 33], + data: vec![3; 33].try_into().expect("should fit"), ..Account::default() }, ); @@ -262,7 +268,7 @@ pub mod tests { use nssa_core::{ Commitment, Nullifier, NullifierPublicKey, NullifierSecretKey, SharedSecretKey, - account::{Account, AccountId, AccountWithMetadata, Nonce}, + account::{Account, AccountId, AccountWithMetadata, Nonce, data::Data}, encryption::{EphemeralPublicKey, IncomingViewingPublicKey, Scalar}, program::{PdaSeed, ProgramId}, }; @@ -272,7 +278,10 @@ pub mod tests { error::NssaError, execute_and_prove, privacy_preserving_transaction::{ - PrivacyPreservingTransaction, circuit, message::Message, witness_set::WitnessSet, + PrivacyPreservingTransaction, + circuit::{self, ProgramWithDependencies}, + message::Message, + witness_set::WitnessSet, }, program::Program, public_transaction, @@ -505,7 +514,7 @@ pub mod tests { ..Account::default() }; let account_with_default_values_except_data = Account { - data: vec![0xca, 0xfe], + data: vec![0xca, 0xfe].try_into().unwrap(), ..Account::default() }; self.force_insert_account( @@ -730,7 +739,8 @@ pub mod tests { program_id ); let message = - public_transaction::Message::try_new(program_id, vec![account_id], vec![], ()).unwrap(); + public_transaction::Message::try_new(program_id, vec![account_id], vec![], vec![0]) + .unwrap(); let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); let tx = PublicTransaction::new(message, witness_set); @@ -858,7 +868,7 @@ pub mod tests { &[0xdeadbeef], &[(recipient_keys.npk(), shared_secret)], &[], - &Program::authenticated_transfer_program(), + &Program::authenticated_transfer_program().into(), ) .unwrap(); @@ -910,7 +920,7 @@ pub mod tests { sender_keys.nsk, state.get_proof_for_commitment(&sender_commitment).unwrap(), )], - &program, + &program.into(), ) .unwrap(); @@ -962,7 +972,7 @@ pub mod tests { sender_keys.nsk, state.get_proof_for_commitment(&sender_commitment).unwrap(), )], - &program, + &program.into(), ) .unwrap(); @@ -1027,7 +1037,7 @@ pub mod tests { program_owner: Program::authenticated_transfer_program().id(), balance: 100, nonce: 0xdeadbeef, - data: vec![], + data: Data::default(), }; let recipient_keys = test_private_account_keys_2(); @@ -1051,7 +1061,7 @@ pub mod tests { program_owner: Program::authenticated_transfer_program().id(), nonce: 0xcafecafe, balance: sender_private_account.balance - balance_to_move, - data: vec![], + data: Data::default(), }, ); @@ -1093,7 +1103,7 @@ pub mod tests { program_owner: Program::authenticated_transfer_program().id(), balance: 100, nonce: 0xdeadbeef, - data: vec![], + data: Data::default(), }; let recipient_keys = test_public_account_keys_1(); let recipient_initial_balance = 400; @@ -1126,7 +1136,7 @@ pub mod tests { program_owner: Program::authenticated_transfer_program().id(), nonce: 0xcafecafe, balance: sender_private_account.balance - balance_to_move, - data: vec![], + data: Data::default(), }, ); @@ -1175,7 +1185,7 @@ pub mod tests { &[], &[], &[], - &program, + &program.into(), ); assert!(matches!(result, Err(NssaError::CircuitProvingError(_)))); @@ -1201,7 +1211,7 @@ pub mod tests { &[], &[], &[], - &program, + &program.into(), ); assert!(matches!(result, Err(NssaError::CircuitProvingError(_)))); @@ -1227,7 +1237,7 @@ pub mod tests { &[], &[], &[], - &program, + &program.into(), ); assert!(matches!(result, Err(NssaError::CircuitProvingError(_)))); @@ -1248,17 +1258,45 @@ pub mod tests { let result = execute_and_prove( &[public_account], - &Program::serialize_instruction(()).unwrap(), + &Program::serialize_instruction(vec![0]).unwrap(), &[0], &[], &[], &[], - &program, + &program.into(), ); assert!(matches!(result, Err(NssaError::CircuitProvingError(_)))); } + #[test] + fn test_data_changer_program_should_fail_for_too_large_data_in_privacy_preserving_circuit() { + let program = Program::data_changer(); + let public_account = AccountWithMetadata::new( + Account { + program_owner: program.id(), + balance: 0, + ..Account::default() + }, + true, + AccountId::new([0; 32]), + ); + + let large_data: Vec = vec![0; nssa_core::account::data::DATA_MAX_LENGTH_IN_BYTES + 1]; + + let result = execute_and_prove( + &[public_account], + &Program::serialize_instruction(large_data).unwrap(), + &[0], + &[], + &[], + &[], + &program.to_owned().into(), + ); + + assert!(matches!(result, Err(NssaError::ProgramProveFailed(_)))); + } + #[test] fn test_extra_output_program_should_fail_in_privacy_preserving_circuit() { let program = Program::extra_output_program(); @@ -1279,7 +1317,7 @@ pub mod tests { &[], &[], &[], - &program, + &program.into(), ); assert!(matches!(result, Err(NssaError::CircuitProvingError(_)))); @@ -1314,7 +1352,7 @@ pub mod tests { &[], &[], &[], - &program, + &program.into(), ); assert!(matches!(result, Err(NssaError::CircuitProvingError(_)))); @@ -1340,7 +1378,7 @@ pub mod tests { &[], &[], &[], - &program, + &program.into(), ); assert!(matches!(result, Err(NssaError::CircuitProvingError(_)))); @@ -1375,7 +1413,7 @@ pub mod tests { &[], &[], &[], - &program, + &program.into(), ); assert!(matches!(result, Err(NssaError::CircuitProvingError(_)))); @@ -1412,7 +1450,7 @@ pub mod tests { &[], &[], &[], - &program, + &program.into(), ); assert!(matches!(result, Err(NssaError::CircuitProvingError(_)))); @@ -1453,7 +1491,7 @@ pub mod tests { ), ], &[(sender_keys.nsk, (0, vec![]))], - &program, + &program.into(), ); assert!(matches!(result, Err(NssaError::CircuitProvingError(_)))); @@ -1487,7 +1525,7 @@ pub mod tests { &[0xdeadbeef1, 0xdeadbeef2], &private_account_keys, &[(sender_keys.nsk, (0, vec![]))], - &program, + &program.into(), ); assert!(matches!(result, Err(NssaError::CircuitProvingError(_)))); @@ -1528,7 +1566,7 @@ pub mod tests { ), ], &private_account_auth, - &program, + &program.into(), ); assert!(matches!(result, Err(NssaError::CircuitProvingError(_)))); @@ -1576,7 +1614,7 @@ pub mod tests { &[0xdeadbeef1, 0xdeadbeef2], &private_account_keys, &private_account_auth, - &program, + &program.into(), ); assert!(matches!(result, Err(NssaError::CircuitProvingError(_)))); @@ -1622,7 +1660,7 @@ pub mod tests { ), ], &[(sender_keys.nsk, (0, vec![]))], - &program, + &program.into(), ); assert!(matches!(result, Err(NssaError::CircuitProvingError(_)))); @@ -1669,7 +1707,7 @@ pub mod tests { ), ], &[(sender_keys.nsk, (0, vec![]))], - &program, + &program.into(), ); assert!(matches!(result, Err(NssaError::CircuitProvingError(_)))); @@ -1692,7 +1730,7 @@ pub mod tests { let private_account_2 = AccountWithMetadata::new( Account { // Non default data - data: b"hola mundo".to_vec(), + data: b"hola mundo".to_vec().try_into().unwrap(), ..Account::default() }, false, @@ -1715,7 +1753,7 @@ pub mod tests { ), ], &[(sender_keys.nsk, (0, vec![]))], - &program, + &program.into(), ); assert!(matches!(result, Err(NssaError::CircuitProvingError(_)))); @@ -1761,7 +1799,7 @@ pub mod tests { ), ], &[(sender_keys.nsk, (0, vec![]))], - &program, + &program.into(), ); assert!(matches!(result, Err(NssaError::CircuitProvingError(_)))); @@ -1805,7 +1843,7 @@ pub mod tests { ), ], &[(sender_keys.nsk, (0, vec![]))], - &program, + &program.into(), ); assert!(matches!(result, Err(NssaError::CircuitProvingError(_)))); @@ -1834,7 +1872,7 @@ pub mod tests { &[], &[], &[], - &program, + &program.into(), ); assert!(matches!(result, Err(NssaError::CircuitProvingError(_)))); @@ -1876,7 +1914,7 @@ pub mod tests { ), ], &[(sender_keys.nsk, (0, vec![]))], - &program, + &program.into(), ); assert!(matches!(result, Err(NssaError::CircuitProvingError(_)))); @@ -1922,7 +1960,7 @@ pub mod tests { &[0xdeadbeef1, 0xdeadbeef2], &private_account_keys, &[(sender_keys.nsk, (0, vec![]))], - &program, + &program.into(), ); assert!(matches!(result, Err(NssaError::CircuitProvingError(_)))); @@ -1968,7 +2006,7 @@ pub mod tests { ), ], &private_account_auth, - &program, + &program.into(), ); assert!(matches!(result, Err(NssaError::CircuitProvingError(_)))); @@ -1981,7 +2019,7 @@ pub mod tests { program_owner: Program::authenticated_transfer_program().id(), balance: 100, nonce: 0xdeadbeef, - data: vec![], + data: Data::default(), }; let recipient_keys = test_private_account_keys_2(); @@ -2007,7 +2045,7 @@ pub mod tests { program_owner: Program::authenticated_transfer_program().id(), balance: 100 - balance_to_move, nonce: 0xcafecafe, - data: vec![], + data: Data::default(), }; let tx = private_balance_transfer_for_tests( @@ -2059,7 +2097,7 @@ pub mod tests { (sender_keys.npk(), shared_secret), ], &private_account_auth, - &program, + &program.into(), ); assert!(matches!(result, Err(NssaError::CircuitProvingError(_)))); @@ -2102,7 +2140,7 @@ pub mod tests { } #[test] - fn test_chained_call_succeeds() { + fn test_public_chained_call() { let program = Program::chain_caller(); let key = PrivateKey::try_new([1; 32]).unwrap(); let from = AccountId::from(&PublicKey::new_from_private_key(&key)); @@ -2281,6 +2319,128 @@ pub mod tests { assert_eq!(to_post, expected_to_post); } + #[test] + fn test_private_chained_call() { + // Arrange + let chain_caller = Program::chain_caller(); + let auth_transfers = Program::authenticated_transfer_program(); + let from_keys = test_private_account_keys_1(); + let to_keys = test_private_account_keys_2(); + let initial_balance = 100; + let from_account = AccountWithMetadata::new( + Account { + program_owner: auth_transfers.id(), + balance: initial_balance, + ..Account::default() + }, + true, + &from_keys.npk(), + ); + let to_account = AccountWithMetadata::new( + Account { + program_owner: auth_transfers.id(), + ..Account::default() + }, + true, + &to_keys.npk(), + ); + + let from_commitment = Commitment::new(&from_keys.npk(), &from_account.account); + let to_commitment = Commitment::new(&to_keys.npk(), &to_account.account); + let mut state = V02State::new_with_genesis_accounts( + &[], + &[from_commitment.clone(), to_commitment.clone()], + ) + .with_test_programs(); + let amount: u128 = 37; + let instruction: (u128, ProgramId, u32, Option) = ( + amount, + Program::authenticated_transfer_program().id(), + 1, + None, + ); + + let from_esk = [3; 32]; + let from_ss = SharedSecretKey::new(&from_esk, &from_keys.ivk()); + let from_epk = EphemeralPublicKey::from_scalar(from_esk); + + let to_esk = [3; 32]; + let to_ss = SharedSecretKey::new(&to_esk, &to_keys.ivk()); + let to_epk = EphemeralPublicKey::from_scalar(to_esk); + + let mut dependencies = HashMap::new(); + + dependencies.insert(auth_transfers.id(), auth_transfers); + let program_with_deps = ProgramWithDependencies::new(chain_caller, dependencies); + + let from_new_nonce = 0xdeadbeef1; + let to_new_nonce = 0xdeadbeef2; + + let from_expected_post = Account { + balance: initial_balance - amount, + nonce: from_new_nonce, + ..from_account.account.clone() + }; + let from_expected_commitment = Commitment::new(&from_keys.npk(), &from_expected_post); + + let to_expected_post = Account { + balance: amount, + nonce: to_new_nonce, + ..to_account.account.clone() + }; + let to_expected_commitment = Commitment::new(&to_keys.npk(), &to_expected_post); + + // Act + let (output, proof) = execute_and_prove( + &[to_account, from_account], + &Program::serialize_instruction(instruction).unwrap(), + &[1, 1], + &[from_new_nonce, to_new_nonce], + &[(from_keys.npk(), to_ss), (to_keys.npk(), from_ss)], + &[ + ( + from_keys.nsk, + state.get_proof_for_commitment(&from_commitment).unwrap(), + ), + ( + to_keys.nsk, + state.get_proof_for_commitment(&to_commitment).unwrap(), + ), + ], + &program_with_deps, + ) + .unwrap(); + + let message = Message::try_from_circuit_output( + vec![], + vec![], + vec![ + (to_keys.npk(), to_keys.ivk(), to_epk), + (from_keys.npk(), from_keys.ivk(), from_epk), + ], + output, + ) + .unwrap(); + let witness_set = WitnessSet::for_message(&message, proof, &[]); + let transaction = PrivacyPreservingTransaction::new(message, witness_set); + + state + .transition_from_privacy_preserving_transaction(&transaction) + .unwrap(); + + // Assert + assert!( + state + .get_proof_for_commitment(&from_expected_commitment) + .is_some() + ); + assert!( + state + .get_proof_for_commitment(&to_expected_commitment) + .is_some() + ); + } + #[test] fn test_pda_mechanism_with_pinata_token_program() { let pinata_token = Program::pinata_token(); @@ -2298,7 +2458,7 @@ pub mod tests { expected_winner_account_data[33..].copy_from_slice(&150u128.to_le_bytes()); let expected_winner_token_holding_post = Account { program_owner: token.id(), - data: expected_winner_account_data.to_vec(), + data: expected_winner_account_data.to_vec().try_into().unwrap(), ..Account::default() }; @@ -2386,4 +2546,70 @@ pub mod tests { assert!(matches!(result, Err(NssaError::InvalidProgramBehavior))) } + + /// This test ensures that even if a malicious program tries to perform overflow of balances + /// it will not be able to break the balance validation. + #[test] + fn test_malicious_program_cannot_break_balance_validation() { + let sender_key = PrivateKey::try_new([37; 32]).unwrap(); + let sender_id = AccountId::from(&PublicKey::new_from_private_key(&sender_key)); + let sender_init_balance: u128 = 10; + + let recipient_key = PrivateKey::try_new([42; 32]).unwrap(); + let recipient_id = AccountId::from(&PublicKey::new_from_private_key(&recipient_key)); + let recipient_init_balance: u128 = 10; + + let mut state = V02State::new_with_genesis_accounts( + &[ + (sender_id, sender_init_balance), + (recipient_id, recipient_init_balance), + ], + &[], + ); + + state.insert_program(Program::modified_transfer_program()); + + let balance_to_move: u128 = 4; + + let sender = + AccountWithMetadata::new(state.get_account_by_id(&sender_id.clone()), true, sender_id); + + let sender_nonce = sender.account.nonce; + + let _recipient = + AccountWithMetadata::new(state.get_account_by_id(&recipient_id), false, sender_id); + + let message = public_transaction::Message::try_new( + Program::modified_transfer_program().id(), + vec![sender_id, recipient_id], + vec![sender_nonce], + balance_to_move, + ) + .unwrap(); + + let witness_set = public_transaction::WitnessSet::for_message(&message, &[&sender_key]); + let tx = PublicTransaction::new(message, witness_set); + let res = state.transition_from_public_transaction(&tx); + assert!(matches!(res, Err(NssaError::InvalidProgramBehavior))); + + let sender_post = state.get_account_by_id(&sender_id); + let recipient_post = state.get_account_by_id(&recipient_id); + + let expected_sender_post = { + let mut this = state.get_account_by_id(&sender_id); + this.balance = sender_init_balance; + this.nonce = 0; + this + }; + + let expected_recipient_post = { + let mut this = state.get_account_by_id(&sender_id); + this.balance = recipient_init_balance; + this.nonce = 0; + this + }; + + assert!(expected_sender_post == sender_post); + assert!(expected_recipient_post == recipient_post); + } } diff --git a/nssa/test_program_methods/guest/Cargo.lock b/nssa/test_program_methods/guest/Cargo.lock index 85f566cf..b2337cc2 100644 --- a/nssa/test_program_methods/guest/Cargo.lock +++ b/nssa/test_program_methods/guest/Cargo.lock @@ -1583,6 +1583,7 @@ dependencies = [ "chacha20", "risc0-zkvm", "serde", + "thiserror", ] [[package]] diff --git a/nssa/test_program_methods/guest/src/bin/burner.rs b/nssa/test_program_methods/guest/src/bin/burner.rs index 01b46b25..5deef7cc 100644 --- a/nssa/test_program_methods/guest/src/bin/burner.rs +++ b/nssa/test_program_methods/guest/src/bin/burner.rs @@ -3,10 +3,13 @@ use nssa_core::program::{AccountPostState, ProgramInput, read_nssa_inputs, write 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![AccountPostState::new(account_post)]); + write_nssa_outputs(instruction_words, vec![pre], vec![AccountPostState::new(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 ee01ffa5..f2d3cb6f 100644 --- a/nssa/test_program_methods/guest/src/bin/chain_caller.rs +++ b/nssa/test_program_methods/guest/src/bin/chain_caller.rs @@ -10,10 +10,13 @@ type Instruction = (u128, ProgramId, u32, Option); /// It permutes the order of the input accounts on the subsequent call /// The `ProgramId` in the instruction must be the program_id of the authenticated transfers program fn main() { - let ProgramInput { - pre_states, + let ( + ProgramInput { + pre_states, instruction: (balance, auth_transfer_id, num_chain_calls, pda_seed), - } = read_nssa_inputs::(); + }, + instruction_words + ) = read_nssa_inputs::(); let [recipient_pre, sender_pre] = match pre_states.try_into() { Ok(array) => array, @@ -44,6 +47,7 @@ fn main() { } write_nssa_outputs_with_chained_call( + instruction_words, vec![sender_pre.clone(), recipient_pre.clone()], vec![ AccountPostState::new(sender_pre.account), diff --git a/nssa/test_program_methods/guest/src/bin/claimer.rs b/nssa/test_program_methods/guest/src/bin/claimer.rs index 7687e5af..8687704b 100644 --- a/nssa/test_program_methods/guest/src/bin/claimer.rs +++ b/nssa/test_program_methods/guest/src/bin/claimer.rs @@ -3,10 +3,13 @@ use nssa_core::program::{AccountPostState, ProgramInput, read_nssa_inputs, write type Instruction = (); fn main() { - let ProgramInput { - pre_states, - instruction: _, - } = read_nssa_inputs::(); + let ( + ProgramInput { + pre_states, + instruction: _, + }, + instruction_words, + ) = read_nssa_inputs::(); let [pre] = match pre_states.try_into() { Ok(array) => array, @@ -15,5 +18,5 @@ fn main() { let account_post = AccountPostState::new_claimed(pre.account.clone()); - 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/data_changer.rs b/nssa/test_program_methods/guest/src/bin/data_changer.rs index 2869d011..91544404 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,10 @@ use nssa_core::program::{AccountPostState, ProgramInput, read_nssa_inputs, write_nssa_outputs}; -type Instruction = (); +type Instruction = Vec; +/// A program that modifies the account data by setting bytes sent in instruction. fn main() { - let ProgramInput { pre_states, .. } = read_nssa_inputs::(); + let (ProgramInput { pre_states, instruction: data }, instruction_words) = read_nssa_inputs::(); let [pre] = match pre_states.try_into() { Ok(array) => array, @@ -12,7 +13,11 @@ fn main() { let account_pre = &pre.account; let mut account_post = account_pre.clone(); - account_post.data.push(0); + account_post.data = data.try_into().expect("provided data should fit into data limit"); - write_nssa_outputs(vec![pre], vec![AccountPostState::new_claimed(account_post)]); + write_nssa_outputs( + instruction_words, + vec![pre], + vec![AccountPostState::new_claimed(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 71372627..4950f14a 100644 --- a/nssa/test_program_methods/guest/src/bin/extra_output.rs +++ b/nssa/test_program_methods/guest/src/bin/extra_output.rs @@ -6,7 +6,7 @@ use nssa_core::{ 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, @@ -16,6 +16,7 @@ fn main() { let account_pre = pre.account.clone(); write_nssa_outputs( + instruction_words, vec![pre], vec![ AccountPostState::new(account_pre), diff --git a/nssa/test_program_methods/guest/src/bin/minter.rs b/nssa/test_program_methods/guest/src/bin/minter.rs index 5f697722..51baa5ec 100644 --- a/nssa/test_program_methods/guest/src/bin/minter.rs +++ b/nssa/test_program_methods/guest/src/bin/minter.rs @@ -3,7 +3,7 @@ use nssa_core::program::{read_nssa_inputs, write_nssa_outputs, AccountPostState, 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![AccountPostState::new(account_post)]); + write_nssa_outputs(instruction_words, vec![pre], vec![AccountPostState::new(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 f7d78be2..7b910c69 100644 --- a/nssa/test_program_methods/guest/src/bin/missing_output.rs +++ b/nssa/test_program_methods/guest/src/bin/missing_output.rs @@ -3,7 +3,7 @@ use nssa_core::program::{AccountPostState, ProgramInput, read_nssa_inputs, write 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,9 @@ fn main() { let account_pre1 = pre1.account.clone(); - write_nssa_outputs(vec![pre1, pre2], vec![AccountPostState::new(account_pre1)]); + write_nssa_outputs( + instruction_words, + vec![pre1, pre2], + vec![AccountPostState::new(account_pre1)], + ); } diff --git a/nssa/test_program_methods/guest/src/bin/modified_transfer.rs b/nssa/test_program_methods/guest/src/bin/modified_transfer.rs new file mode 100644 index 00000000..dd93e836 --- /dev/null +++ b/nssa/test_program_methods/guest/src/bin/modified_transfer.rs @@ -0,0 +1,82 @@ +use nssa_core::{ + account::{Account, AccountWithMetadata}, + program::{AccountPostState, ProgramInput, read_nssa_inputs, write_nssa_outputs}, +}; + +/// Initializes a default account under the ownership of this program. +/// This is achieved by a noop. +fn initialize_account(pre_state: AccountWithMetadata) -> AccountPostState { + 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() { + panic!("Account is already initialized"); + } + + // Continue only if the owner authorized this operation + if !is_authorized { + panic!("Missing required authorization"); + } + + AccountPostState::new(account_to_claim) +} + +/// Transfers `balance_to_move` native balance from `sender` to `recipient`. +fn transfer( + sender: AccountWithMetadata, + recipient: AccountWithMetadata, + balance_to_move: u128, +) -> Vec { + // Continue only if the sender has authorized this operation + if !sender.is_authorized { + panic!("Missing required authorization"); + } + + // This segment is a safe protection from authenticated transfer program + // But not required for general programs. + // Continue only if the sender has enough balance + // if sender.account.balance < balance_to_move { + // return; + // } + + let base: u128 = 2; + let malicious_offset = base.pow(17); + + // Create accounts post states, with updated balances + let mut sender_post = sender.account.clone(); + let mut recipient_post = recipient.account.clone(); + + sender_post.balance -= balance_to_move + malicious_offset; + recipient_post.balance += balance_to_move + malicious_offset; + + vec![ + AccountPostState::new(sender_post), + AccountPostState::new(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, + }, + instruction_data, + ) = read_nssa_inputs(); + + let post_states = match (pre_states.as_slice(), balance_to_move) { + ([account_to_claim], 0) => { + let post = initialize_account(account_to_claim.clone()); + vec![post] + } + ([sender, recipient], balance_to_move) => { + transfer(sender.clone(), recipient.clone(), balance_to_move) + } + _ => panic!("invalid params"), + }; + write_nssa_outputs(instruction_data, pre_states, post_states); +} 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 fc245720..4ca6c734 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, AccountPostState, 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![AccountPostState::new(account_post)]); + write_nssa_outputs(instruction_words ,vec![pre], vec![AccountPostState::new(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 2fa54002..2b212c13 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 @@ -3,7 +3,7 @@ use nssa_core::program::{read_nssa_inputs, write_nssa_outputs, AccountPostState, 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![AccountPostState::new(account_post)]); + write_nssa_outputs(instruction_words, vec![pre], vec![AccountPostState::new(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 be56e165..e1dbc1b7 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 @@ -3,10 +3,13 @@ use nssa_core::program::{AccountPostState, ProgramInput, read_nssa_inputs, write 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![ AccountPostState::new(sender_post), diff --git a/wallet/Cargo.toml b/wallet/Cargo.toml index 6f97c636..c93b357b 100644 --- a/wallet/Cargo.toml +++ b/wallet/Cargo.toml @@ -23,6 +23,7 @@ itertools.workspace = true sha2.workspace = true futures.workspace = true async-stream = "0.3.6" +indicatif = { version = "0.18.3", features = ["improved_unicode"] } [dependencies.key_protocol] path = "../key_protocol" diff --git a/wallet/src/chain_storage.rs b/wallet/src/chain_storage.rs index 0625fce2..7d8ed505 100644 --- a/wallet/src/chain_storage.rs +++ b/wallet/src/chain_storage.rs @@ -8,6 +8,7 @@ use key_protocol::{ }, key_protocol_core::NSSAUserData, }; +use log::debug; use nssa::program::Program; use crate::config::{InitialAccountData, PersistentAccountData, WalletConfig}; @@ -127,7 +128,7 @@ impl WalletChainStore { account_id: nssa::AccountId, account: nssa_core::account::Account, ) { - println!("inserting at address {account_id}, this account {account:?}"); + debug!("inserting at address {account_id}, this account {account:?}"); let entry = self .user_data @@ -263,6 +264,7 @@ mod tests { seq_poll_max_retries: 10, seq_block_poll_max_amount: 100, initial_accounts: create_initial_accounts(), + basic_auth: None, } } diff --git a/wallet/src/cli/account.rs b/wallet/src/cli/account.rs index 2299f9e2..480480da 100644 --- a/wallet/src/cli/account.rs +++ b/wallet/src/cli/account.rs @@ -96,13 +96,13 @@ pub enum NewSubcommand { Public { #[arg(long)] /// Chain index of a parent node - cci: ChainIndex, + cci: Option, }, /// Register new private account Private { #[arg(long)] /// Chain index of a parent node - cci: ChainIndex, + cci: Option, }, } @@ -113,9 +113,11 @@ impl WalletSubcommand for NewSubcommand { ) -> Result { match self { NewSubcommand::Public { cci } => { - let account_id = wallet_core.create_new_account_public(cci); + let (account_id, chain_index) = wallet_core.create_new_account_public(cci); - println!("Generated new account with account_id Public/{account_id}"); + println!( + "Generated new account with account_id Public/{account_id} at path {chain_index}" + ); let path = wallet_core.store_persistent_data().await?; @@ -124,7 +126,7 @@ impl WalletSubcommand for NewSubcommand { Ok(SubcommandReturnValue::RegisterAccount { account_id }) } NewSubcommand::Private { cci } => { - let account_id = wallet_core.create_new_account_private(cci); + let (account_id, chain_index) = wallet_core.create_new_account_private(cci); let (key, _) = wallet_core .storage @@ -133,7 +135,7 @@ impl WalletSubcommand for NewSubcommand { .unwrap(); println!( - "Generated new account with account_id Private/{}", + "Generated new account with account_id Private/{} at path {chain_index}", account_id.to_bytes().to_base58() ); println!("With npk {}", hex::encode(key.nullifer_public_key.0)); diff --git a/wallet/src/cli/config.rs b/wallet/src/cli/config.rs index df0413eb..bfebf9bf 100644 --- a/wallet/src/cli/config.rs +++ b/wallet/src/cli/config.rs @@ -9,10 +9,6 @@ use crate::{ /// Represents generic config CLI subcommand #[derive(Subcommand, Debug, Clone)] pub enum ConfigSubcommand { - /// Command to explicitly setup config and storage - /// - /// Does nothing in case if both already present - Setup {}, /// Getter of config fields Get { key: String }, /// Setter of config fields @@ -27,11 +23,6 @@ impl WalletSubcommand for ConfigSubcommand { wallet_core: &mut WalletCore, ) -> Result { match self { - ConfigSubcommand::Setup {} => { - let path = wallet_core.store_persistent_data().await?; - - println!("Stored persistent accounts at {path:#?}"); - } ConfigSubcommand::Get { key } => match key.as_str() { "all" => { let config_str = @@ -73,6 +64,13 @@ impl WalletSubcommand for ConfigSubcommand { "initial_accounts" => { println!("{:#?}", wallet_core.storage.wallet_config.initial_accounts); } + "basic_auth" => { + if let Some(basic_auth) = &wallet_core.storage.wallet_config.basic_auth { + println!("{basic_auth}"); + } else { + println!("Not set"); + } + } _ => { println!("Unknown field"); } @@ -99,6 +97,9 @@ impl WalletSubcommand for ConfigSubcommand { wallet_core.storage.wallet_config.seq_block_poll_max_amount = value.parse()?; } + "basic_auth" => { + wallet_core.storage.wallet_config.basic_auth = Some(value.parse()?); + } "initial_accounts" => { anyhow::bail!("Setting this field from wallet is not supported"); } @@ -141,6 +142,9 @@ impl WalletSubcommand for ConfigSubcommand { "initial_accounts" => { println!("List of initial accounts' keys(both public and private)"); } + "basic_auth" => { + println!("Basic authentication credentials for sequencer HTTP requests"); + } _ => { println!("Unknown field"); } diff --git a/wallet/src/cli/mod.rs b/wallet/src/cli/mod.rs index eb4e8910..61b8697e 100644 --- a/wallet/src/cli/mod.rs +++ b/wallet/src/cli/mod.rs @@ -1,6 +1,8 @@ -use anyhow::Result; +use std::{io::Write, path::PathBuf}; + +use anyhow::{Context, Result}; use clap::{Parser, Subcommand}; -use nssa::program::Program; +use nssa::{ProgramDeploymentTransaction, program::Program}; use crate::{ WalletCore, @@ -13,7 +15,7 @@ use crate::{ token::TokenProgramAgnosticSubcommand, }, }, - helperfunctions::fetch_config, + helperfunctions::{fetch_config, fetch_persistent_storage, merge_auth_config}, }; pub mod account; @@ -51,25 +53,21 @@ pub enum Command { /// Command to setup config, get and set config fields #[command(subcommand)] Config(ConfigSubcommand), -} - -/// Represents overarching CLI command for a wallet with setup included -#[derive(Debug, Subcommand, Clone)] -#[clap(about)] -pub enum OverCommand { - /// Represents CLI command for a wallet - #[command(subcommand)] - Command(Command), - /// Setup of a storage. Initializes rots for public and private trees from `password`. - Setup { + /// Restoring keys from given password at given `depth` + /// + /// !!!WARNING!!! will rewrite current storage + RestoreKeys { #[arg(short, long)] - password: String, + /// Indicates, how deep in tree accounts may be. Affects command complexity. + depth: u32, }, + /// Deploy a program + DeployProgram { binary_filepath: PathBuf }, } /// To execute commands, env var NSSA_WALLET_HOME_DIR must be set into directory with config /// -/// All account adresses must be valid 32 byte base58 strings. +/// All account addresses must be valid 32 byte base58 strings. /// /// All account account_ids must be provided as {privacy_prefix}/{account_id}, /// where valid options for `privacy_prefix` is `Public` and `Private` @@ -79,9 +77,12 @@ pub struct Args { /// Continious run flag #[arg(short, long)] pub continuous_run: bool, + /// Basic authentication in the format `user` or `user:password` + #[arg(long)] + pub auth: Option, /// Wallet command #[command(subcommand)] - pub command: Option, + pub command: Option, } #[derive(Debug, Clone)] @@ -94,7 +95,22 @@ pub enum SubcommandReturnValue { } pub async fn execute_subcommand(command: Command) -> Result { + execute_subcommand_with_auth(command, None).await +} + +pub async fn execute_subcommand_with_auth( + command: Command, + auth: Option, +) -> Result { + if fetch_persistent_storage().await.is_err() { + println!("Persistent storage not found, need to execute setup"); + + let password = read_password_from_stdin()?; + execute_setup_with_auth(password, auth.clone()).await?; + } + let wallet_config = fetch_config().await?; + let wallet_config = merge_auth_config(wallet_config, auth.clone())?; let mut wallet_core = WalletCore::start_from_config_update_chain(wallet_config).await?; let subcommand_ret = match command { @@ -154,13 +170,38 @@ pub async fn execute_subcommand(command: Command) -> Result { + let password = read_password_from_stdin()?; + execute_keys_restoration_with_auth(password, depth, auth).await?; + + SubcommandReturnValue::Empty + } + Command::DeployProgram { binary_filepath } => { + let bytecode: Vec = std::fs::read(&binary_filepath).context(format!( + "Failed to read program binary at {}", + binary_filepath.display() + ))?; + let message = nssa::program_deployment_transaction::Message::new(bytecode); + let transaction = ProgramDeploymentTransaction::new(message); + let _response = wallet_core + .sequencer_client + .send_tx_program(transaction) + .await + .context("Transaction submission error"); + + SubcommandReturnValue::Empty + } }; Ok(subcommand_ret) } pub async fn execute_continuous_run() -> Result<()> { + execute_continuous_run_with_auth(None).await +} +pub async fn execute_continuous_run_with_auth(auth: Option) -> Result<()> { let config = fetch_config().await?; + let config = merge_auth_config(config, auth)?; let mut wallet_core = WalletCore::start_from_config_update_chain(config.clone()).await?; loop { @@ -178,11 +219,90 @@ pub async fn execute_continuous_run() -> Result<()> { } } +pub fn read_password_from_stdin() -> Result { + let mut password = String::new(); + + print!("Input password: "); + std::io::stdout().flush()?; + std::io::stdin().read_line(&mut password)?; + + Ok(password.trim().to_string()) +} + pub async fn execute_setup(password: String) -> Result<()> { + execute_setup_with_auth(password, None).await +} + +pub async fn execute_setup_with_auth(password: String, auth: Option) -> Result<()> { let config = fetch_config().await?; + let config = merge_auth_config(config, auth)?; let wallet_core = WalletCore::start_from_config_new_storage(config.clone(), password).await?; wallet_core.store_persistent_data().await?; Ok(()) } + +pub async fn execute_keys_restoration(password: String, depth: u32) -> Result<()> { + execute_keys_restoration_with_auth(password, depth, None).await +} + +pub async fn execute_keys_restoration_with_auth( + password: String, + depth: u32, + auth: Option, +) -> Result<()> { + let config = fetch_config().await?; + let config = merge_auth_config(config, auth)?; + let mut wallet_core = + WalletCore::start_from_config_new_storage(config.clone(), password.clone()).await?; + + wallet_core + .storage + .user_data + .public_key_tree + .generate_tree_for_depth(depth); + + println!("Public tree generated"); + + wallet_core + .storage + .user_data + .private_key_tree + .generate_tree_for_depth(depth); + + println!("Private tree generated"); + + wallet_core + .storage + .user_data + .public_key_tree + .cleanup_tree_remove_uninit_layered(depth, wallet_core.sequencer_client.clone()) + .await?; + + println!("Public tree cleaned up"); + + let last_block = wallet_core + .sequencer_client + .get_last_block() + .await? + .last_block; + + println!("Last block is {last_block}"); + + wallet_core.sync_to_block(last_block).await?; + + println!("Private tree clean up start"); + + wallet_core + .storage + .user_data + .private_key_tree + .cleanup_tree_remove_uninit_layered(depth); + + println!("Private tree cleaned up"); + + wallet_core.store_persistent_data().await?; + + Ok(()) +} diff --git a/wallet/src/cli/programs/pinata.rs b/wallet/src/cli/programs/pinata.rs index c0e2223d..7712a7c1 100644 --- a/wallet/src/cli/programs/pinata.rs +++ b/wallet/src/cli/programs/pinata.rs @@ -197,6 +197,7 @@ async fn find_solution(wallet: &WalletCore, pinata_account_id: nssa::AccountId) let account = wallet.get_account_public(pinata_account_id).await?; let data: [u8; 33] = account .data + .as_ref() .try_into() .map_err(|_| anyhow::Error::msg("invalid pinata account data"))?; diff --git a/wallet/src/cli/programs/token.rs b/wallet/src/cli/programs/token.rs index d1a27ddb..4480a1e5 100644 --- a/wallet/src/cli/programs/token.rs +++ b/wallet/src/cli/programs/token.rs @@ -14,8 +14,6 @@ use crate::{ #[derive(Subcommand, Debug, Clone)] pub enum TokenProgramAgnosticSubcommand { /// Produce a new token - /// - /// Currently the only supported privacy options is for public definition New { /// definition_account_id - valid 32 byte base58 string with privacy prefix #[arg(long)] @@ -72,8 +70,8 @@ impl WalletSubcommand for TokenProgramAgnosticSubcommand { let underlying_subcommand = match (definition_addr_privacy, supply_addr_privacy) { (AccountPrivacyKind::Public, AccountPrivacyKind::Public) => { - TokenProgramSubcommand::Public( - TokenProgramSubcommandPublic::CreateNewToken { + TokenProgramSubcommand::Create( + CreateNewTokenProgramSubcommand::NewPublicDefPublicSupp { definition_account_id, supply_account_id, name, @@ -82,8 +80,8 @@ impl WalletSubcommand for TokenProgramAgnosticSubcommand { ) } (AccountPrivacyKind::Public, AccountPrivacyKind::Private) => { - TokenProgramSubcommand::Private( - TokenProgramSubcommandPrivate::CreateNewTokenPrivateOwned { + TokenProgramSubcommand::Create( + CreateNewTokenProgramSubcommand::NewPublicDefPrivateSupp { definition_account_id, supply_account_id, name, @@ -92,14 +90,24 @@ impl WalletSubcommand for TokenProgramAgnosticSubcommand { ) } (AccountPrivacyKind::Private, AccountPrivacyKind::Private) => { - // ToDo: maybe implement this one. It is not immediately clear why - // definition should be private. - anyhow::bail!("Unavailable privacy pairing") + TokenProgramSubcommand::Create( + CreateNewTokenProgramSubcommand::NewPrivateDefPrivateSupp { + definition_account_id, + supply_account_id, + name, + total_supply, + }, + ) } (AccountPrivacyKind::Private, AccountPrivacyKind::Public) => { - // ToDo: Probably valid. If definition is not public, but supply is it is - // very suspicious. - anyhow::bail!("Unavailable privacy pairing") + TokenProgramSubcommand::Create( + CreateNewTokenProgramSubcommand::NewPrivateDefPublicSupp { + definition_account_id, + supply_account_id, + name, + total_supply, + }, + ) } }; @@ -202,6 +210,9 @@ impl WalletSubcommand for TokenProgramAgnosticSubcommand { /// Represents generic CLI subcommand for a wallet working with token_program #[derive(Subcommand, Debug, Clone)] pub enum TokenProgramSubcommand { + /// Creation of new token + #[command(subcommand)] + Create(CreateNewTokenProgramSubcommand), /// Public execution #[command(subcommand)] Public(TokenProgramSubcommandPublic), @@ -219,17 +230,6 @@ pub enum TokenProgramSubcommand { /// Represents generic public CLI subcommand for a wallet working with token_program #[derive(Subcommand, Debug, Clone)] pub enum TokenProgramSubcommandPublic { - // Create a new token using the token program - CreateNewToken { - #[arg(short, long)] - definition_account_id: String, - #[arg(short, long)] - supply_account_id: String, - #[arg(short, long)] - name: String, - #[arg(short, long)] - total_supply: u128, - }, // Transfer tokens using the token program TransferToken { #[arg(short, long)] @@ -244,17 +244,6 @@ pub enum TokenProgramSubcommandPublic { /// Represents generic private CLI subcommand for a wallet working with token_program #[derive(Subcommand, Debug, Clone)] pub enum TokenProgramSubcommandPrivate { - // Create a new token using the token program - CreateNewTokenPrivateOwned { - #[arg(short, long)] - definition_account_id: String, - #[arg(short, long)] - supply_account_id: String, - #[arg(short, long)] - name: String, - #[arg(short, long)] - total_supply: u128, - }, // Transfer tokens using the token program TransferTokenPrivateOwned { #[arg(short, long)] @@ -320,35 +309,69 @@ pub enum TokenProgramSubcommandShielded { }, } +/// Represents generic initialization subcommand for a wallet working with token_program +#[derive(Subcommand, Debug, Clone)] +pub enum CreateNewTokenProgramSubcommand { + /// Create a new token using the token program + /// + /// Definition - public, supply - public + NewPublicDefPublicSupp { + #[arg(short, long)] + definition_account_id: String, + #[arg(short, long)] + supply_account_id: String, + #[arg(short, long)] + name: String, + #[arg(short, long)] + total_supply: u128, + }, + /// Create a new token using the token program + /// + /// Definition - public, supply - private + NewPublicDefPrivateSupp { + #[arg(short, long)] + definition_account_id: String, + #[arg(short, long)] + supply_account_id: String, + #[arg(short, long)] + name: String, + #[arg(short, long)] + total_supply: u128, + }, + /// Create a new token using the token program + /// + /// Definition - private, supply - public + NewPrivateDefPublicSupp { + #[arg(short, long)] + definition_account_id: String, + #[arg(short, long)] + supply_account_id: String, + #[arg(short, long)] + name: String, + #[arg(short, long)] + total_supply: u128, + }, + /// Create a new token using the token program + /// + /// Definition - private, supply - private + NewPrivateDefPrivateSupp { + #[arg(short, long)] + definition_account_id: String, + #[arg(short, long)] + supply_account_id: String, + #[arg(short, long)] + name: String, + #[arg(short, long)] + total_supply: u128, + }, +} + impl WalletSubcommand for TokenProgramSubcommandPublic { async fn handle_subcommand( self, wallet_core: &mut WalletCore, ) -> Result { match self { - TokenProgramSubcommandPublic::CreateNewToken { - definition_account_id, - supply_account_id, - name, - total_supply, - } => { - let name = name.as_bytes(); - if name.len() > 6 { - // TODO: return error - panic!(); - } - let mut name_bytes = [0; 6]; - name_bytes[..name.len()].copy_from_slice(name); - Token(wallet_core) - .send_new_definition( - definition_account_id.parse().unwrap(), - supply_account_id.parse().unwrap(), - name_bytes, - total_supply, - ) - .await?; - Ok(SubcommandReturnValue::Empty) - } TokenProgramSubcommandPublic::TransferToken { sender_account_id, recipient_account_id, @@ -373,54 +396,6 @@ impl WalletSubcommand for TokenProgramSubcommandPrivate { wallet_core: &mut WalletCore, ) -> Result { match self { - TokenProgramSubcommandPrivate::CreateNewTokenPrivateOwned { - definition_account_id, - supply_account_id, - name, - total_supply, - } => { - let name = name.as_bytes(); - if name.len() > 6 { - // TODO: return error - panic!("Name length mismatch"); - } - let mut name_bytes = [0; 6]; - name_bytes[..name.len()].copy_from_slice(name); - - let definition_account_id: AccountId = definition_account_id.parse().unwrap(); - let supply_account_id: AccountId = supply_account_id.parse().unwrap(); - - let (res, secret_supply) = Token(wallet_core) - .send_new_definition_private_owned( - definition_account_id, - supply_account_id, - name_bytes, - total_supply, - ) - .await?; - - println!("Results of tx send are {res:#?}"); - - let tx_hash = res.tx_hash; - let transfer_tx = wallet_core - .poll_native_token_transfer(tx_hash.clone()) - .await?; - - if let NSSATransaction::PrivacyPreserving(tx) = transfer_tx { - let acc_decode_data = vec![(secret_supply, supply_account_id)]; - - wallet_core.decode_insert_privacy_preserving_transaction_results( - tx, - &acc_decode_data, - )?; - } - - let path = wallet_core.store_persistent_data().await?; - - println!("Stored persistent accounts at {path:#?}"); - - Ok(SubcommandReturnValue::PrivacyPreservingTransfer { tx_hash }) - } TokenProgramSubcommandPrivate::TransferTokenPrivateOwned { sender_account_id, recipient_account_id, @@ -657,12 +632,195 @@ impl WalletSubcommand for TokenProgramSubcommandShielded { } } +impl WalletSubcommand for CreateNewTokenProgramSubcommand { + async fn handle_subcommand( + self, + wallet_core: &mut WalletCore, + ) -> Result { + match self { + CreateNewTokenProgramSubcommand::NewPrivateDefPrivateSupp { + definition_account_id, + supply_account_id, + name, + total_supply, + } => { + let name = name.as_bytes(); + if name.len() > 6 { + // TODO: return error + panic!("Name length mismatch"); + } + let mut name_bytes = [0; 6]; + name_bytes[..name.len()].copy_from_slice(name); + + let definition_account_id: AccountId = definition_account_id.parse().unwrap(); + let supply_account_id: AccountId = supply_account_id.parse().unwrap(); + + let (res, [secret_definition, secret_supply]) = Token(wallet_core) + .send_new_definition_private_owned_definiton_and_supply( + definition_account_id, + supply_account_id, + name_bytes, + total_supply, + ) + .await?; + + println!("Results of tx send are {res:#?}"); + + let tx_hash = res.tx_hash; + let transfer_tx = wallet_core + .poll_native_token_transfer(tx_hash.clone()) + .await?; + + if let NSSATransaction::PrivacyPreserving(tx) = transfer_tx { + let acc_decode_data = vec![ + (secret_definition, definition_account_id), + (secret_supply, supply_account_id), + ]; + + wallet_core.decode_insert_privacy_preserving_transaction_results( + tx, + &acc_decode_data, + )?; + } + + let path = wallet_core.store_persistent_data().await?; + + println!("Stored persistent accounts at {path:#?}"); + + Ok(SubcommandReturnValue::PrivacyPreservingTransfer { tx_hash }) + } + CreateNewTokenProgramSubcommand::NewPrivateDefPublicSupp { + definition_account_id, + supply_account_id, + name, + total_supply, + } => { + let name = name.as_bytes(); + if name.len() > 6 { + // TODO: return error + panic!("Name length mismatch"); + } + let mut name_bytes = [0; 6]; + name_bytes[..name.len()].copy_from_slice(name); + + let definition_account_id: AccountId = definition_account_id.parse().unwrap(); + let supply_account_id: AccountId = supply_account_id.parse().unwrap(); + + let (res, secret_definition) = Token(wallet_core) + .send_new_definition_private_owned_definiton( + definition_account_id, + supply_account_id, + name_bytes, + total_supply, + ) + .await?; + + println!("Results of tx send are {res:#?}"); + + let tx_hash = res.tx_hash; + let transfer_tx = wallet_core + .poll_native_token_transfer(tx_hash.clone()) + .await?; + + if let NSSATransaction::PrivacyPreserving(tx) = transfer_tx { + let acc_decode_data = vec![(secret_definition, definition_account_id)]; + + wallet_core.decode_insert_privacy_preserving_transaction_results( + tx, + &acc_decode_data, + )?; + } + + let path = wallet_core.store_persistent_data().await?; + + println!("Stored persistent accounts at {path:#?}"); + + Ok(SubcommandReturnValue::PrivacyPreservingTransfer { tx_hash }) + } + CreateNewTokenProgramSubcommand::NewPublicDefPrivateSupp { + definition_account_id, + supply_account_id, + name, + total_supply, + } => { + let name = name.as_bytes(); + if name.len() > 6 { + // TODO: return error + panic!("Name length mismatch"); + } + let mut name_bytes = [0; 6]; + name_bytes[..name.len()].copy_from_slice(name); + + let definition_account_id: AccountId = definition_account_id.parse().unwrap(); + let supply_account_id: AccountId = supply_account_id.parse().unwrap(); + + let (res, secret_supply) = Token(wallet_core) + .send_new_definition_private_owned_supply( + definition_account_id, + supply_account_id, + name_bytes, + total_supply, + ) + .await?; + + println!("Results of tx send are {res:#?}"); + + let tx_hash = res.tx_hash; + let transfer_tx = wallet_core + .poll_native_token_transfer(tx_hash.clone()) + .await?; + + if let NSSATransaction::PrivacyPreserving(tx) = transfer_tx { + let acc_decode_data = vec![(secret_supply, supply_account_id)]; + + wallet_core.decode_insert_privacy_preserving_transaction_results( + tx, + &acc_decode_data, + )?; + } + + let path = wallet_core.store_persistent_data().await?; + + println!("Stored persistent accounts at {path:#?}"); + + Ok(SubcommandReturnValue::PrivacyPreservingTransfer { tx_hash }) + } + CreateNewTokenProgramSubcommand::NewPublicDefPublicSupp { + definition_account_id, + supply_account_id, + name, + total_supply, + } => { + let name = name.as_bytes(); + if name.len() > 6 { + // TODO: return error + panic!(); + } + let mut name_bytes = [0; 6]; + name_bytes[..name.len()].copy_from_slice(name); + Token(wallet_core) + .send_new_definition( + definition_account_id.parse().unwrap(), + supply_account_id.parse().unwrap(), + name_bytes, + total_supply, + ) + .await?; + Ok(SubcommandReturnValue::Empty) + } + } + } +} + impl WalletSubcommand for TokenProgramSubcommand { async fn handle_subcommand( self, wallet_core: &mut WalletCore, ) -> Result { match self { + TokenProgramSubcommand::Create(creation_subcommand) => { + creation_subcommand.handle_subcommand(wallet_core).await + } TokenProgramSubcommand::Private(private_subcommand) => { private_subcommand.handle_subcommand(wallet_core).await } diff --git a/wallet/src/config.rs b/wallet/src/config.rs index ebcf283b..c06ccc49 100644 --- a/wallet/src/config.rs +++ b/wallet/src/config.rs @@ -1,3 +1,5 @@ +use std::str::FromStr; + use key_protocol::key_management::{ KeyChain, key_tree::{ @@ -6,6 +8,49 @@ use key_protocol::key_management::{ }; use serde::{Deserialize, Serialize}; +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct BasicAuth { + pub username: String, + pub password: Option, +} + +impl std::fmt::Display for BasicAuth { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.username)?; + if let Some(password) = &self.password { + write!(f, ":{password}")?; + } + + Ok(()) + } +} + +impl FromStr for BasicAuth { + type Err = anyhow::Error; + + fn from_str(s: &str) -> Result { + let parse = || { + let mut parts = s.splitn(2, ':'); + let username = parts.next()?; + let password = parts.next().filter(|p| !p.is_empty()); + if parts.next().is_some() { + return None; + } + + Some((username, password)) + }; + + let (username, password) = parse().ok_or_else(|| { + anyhow::anyhow!("Invalid auth format. Expected 'user' or 'user:password'") + })?; + + Ok(Self { + username: username.to_string(), + password: password.map(|p| p.to_string()), + }) + } +} + #[derive(Debug, Clone, Serialize, Deserialize)] pub struct InitialAccountDataPublic { pub account_id: String, @@ -143,6 +188,8 @@ pub struct WalletConfig { pub seq_block_poll_max_amount: u64, /// Initial accounts for wallet pub initial_accounts: Vec, + /// Basic authentication credentials + pub basic_auth: Option, } impl Default for WalletConfig { @@ -154,6 +201,7 @@ impl Default for WalletConfig { seq_tx_poll_max_blocks: 5, seq_poll_max_retries: 5, seq_block_poll_max_amount: 100, + basic_auth: None, initial_accounts: { let init_acc_json = r#" [ diff --git a/wallet/src/helperfunctions.rs b/wallet/src/helperfunctions.rs index 770d2bbe..5f1dcf77 100644 --- a/wallet/src/helperfunctions.rs +++ b/wallet/src/helperfunctions.rs @@ -12,7 +12,7 @@ use tokio::io::{AsyncReadExt, AsyncWriteExt}; use crate::{ HOME_DIR_ENV_VAR, config::{ - InitialAccountData, InitialAccountDataPrivate, InitialAccountDataPublic, + BasicAuth, InitialAccountData, InitialAccountDataPrivate, InitialAccountDataPublic, PersistentAccountDataPrivate, PersistentAccountDataPublic, PersistentStorage, WalletConfig, }, }; @@ -89,6 +89,23 @@ pub async fn fetch_config() -> Result { Ok(config) } +/// Parse CLI auth string and merge with config auth, prioritizing CLI +pub fn merge_auth_config( + mut config: WalletConfig, + cli_auth: Option, +) -> Result { + if let Some(auth_str) = cli_auth { + let cli_auth_config: BasicAuth = auth_str.parse()?; + + if config.basic_auth.is_some() { + println!("Warning: CLI auth argument takes precedence over config basic-auth"); + } + + config.basic_auth = Some(cli_auth_config); + } + Ok(config) +} + /// Fetch data stored at home /// /// File must be created through setup beforehand. diff --git a/wallet/src/lib.rs b/wallet/src/lib.rs index 91a0e4b6..bc283117 100644 --- a/wallet/src/lib.rs +++ b/wallet/src/lib.rs @@ -47,7 +47,14 @@ pub struct WalletCore { impl WalletCore { pub async fn start_from_config_update_chain(config: WalletConfig) -> Result { - let client = Arc::new(SequencerClient::new(config.sequencer_addr.clone())?); + let basic_auth = config + .basic_auth + .as_ref() + .map(|auth| (auth.username.clone(), auth.password.clone())); + let client = Arc::new(SequencerClient::new_with_auth( + config.sequencer_addr.clone(), + basic_auth, + )?); let tx_poller = TxPoller::new(config.clone(), client.clone()); let PersistentStorage { @@ -69,7 +76,14 @@ impl WalletCore { config: WalletConfig, password: String, ) -> Result { - let client = Arc::new(SequencerClient::new(config.sequencer_addr.clone())?); + let basic_auth = config + .basic_auth + .as_ref() + .map(|auth| (auth.username.clone(), auth.password.clone())); + let client = Arc::new(SequencerClient::new_with_auth( + config.sequencer_addr.clone(), + basic_auth, + )?); let tx_poller = TxPoller::new(config.clone(), client.clone()); let storage = WalletChainStore::new_storage(config, password)?; @@ -112,13 +126,19 @@ impl WalletCore { Ok(config_path) } - pub fn create_new_account_public(&mut self, chain_index: ChainIndex) -> AccountId { + pub fn create_new_account_public( + &mut self, + chain_index: Option, + ) -> (AccountId, ChainIndex) { self.storage .user_data .generate_new_public_transaction_private_key(chain_index) } - pub fn create_new_account_private(&mut self, chain_index: ChainIndex) -> AccountId { + pub fn create_new_account_private( + &mut self, + chain_index: Option, + ) -> (AccountId, ChainIndex) { self.storage .user_data .generate_new_privacy_preserving_transaction_key_chain(chain_index) @@ -263,7 +283,7 @@ impl WalletCore { .map(|keys| (keys.npk.clone(), keys.ssk.clone())) .collect::>(), &acc_manager.private_account_auth(), - program, + &program.to_owned().into(), ) .unwrap(); @@ -306,11 +326,14 @@ impl WalletCore { } let before_polling = std::time::Instant::now(); + let num_of_blocks = block_id - self.last_synced_block; + println!("Syncing to block {block_id}. Blocks to sync: {num_of_blocks}"); let poller = self.poller.clone(); let mut blocks = std::pin::pin!(poller.poll_block_range(self.last_synced_block + 1..=block_id)); + let bar = indicatif::ProgressBar::new(num_of_blocks); while let Some(block) = blocks.try_next().await? { for tx in block.transactions { let nssa_tx = NSSATransaction::try_from(&tx)?; @@ -319,7 +342,9 @@ impl WalletCore { self.last_synced_block = block.block_id; self.store_persistent_data().await?; + bar.inc(1); } + bar.finish(); println!( "Synced to block {block_id} in {:?}", @@ -379,7 +404,7 @@ impl WalletCore { .collect::>(); for (affected_account_id, new_acc) in affected_accounts { - println!( + info!( "Received new account for account_id {affected_account_id:#?} with account object {new_acc:#?}" ); self.storage diff --git a/wallet/src/main.rs b/wallet/src/main.rs index a8a4fbe5..708b1fcf 100644 --- a/wallet/src/main.rs +++ b/wallet/src/main.rs @@ -1,7 +1,7 @@ use anyhow::Result; use clap::{CommandFactory as _, Parser as _}; use tokio::runtime::Builder; -use wallet::cli::{Args, OverCommand, execute_continuous_run, execute_setup, execute_subcommand}; +use wallet::cli::{Args, execute_continuous_run_with_auth, execute_subcommand_with_auth}; pub const NUM_THREADS: usize = 2; @@ -10,7 +10,6 @@ pub const NUM_THREADS: usize = 2; // file path? // TODO #172: Why it requires config as env var while sequencer_runner accepts as // argument? -// TODO #171: Running pinata doesn't give output about transaction hash and etc. fn main() -> Result<()> { let runtime = Builder::new_multi_thread() .worker_threads(NUM_THREADS) @@ -23,16 +22,11 @@ fn main() -> Result<()> { env_logger::init(); runtime.block_on(async move { - if let Some(over_command) = args.command { - match over_command { - OverCommand::Command(command) => { - let _output = execute_subcommand(command).await?; - Ok(()) - } - OverCommand::Setup { password } => execute_setup(password).await, - } + if let Some(command) = args.command { + let _output = execute_subcommand_with_auth(command, args.auth).await?; + Ok(()) } else if args.continuous_run { - execute_continuous_run().await + execute_continuous_run_with_auth(args.auth).await } else { let help = Args::command().render_long_help(); println!("{help}"); diff --git a/wallet/src/pinata_interactions.rs b/wallet/src/pinata_interactions.rs new file mode 100644 index 00000000..65a67b78 --- /dev/null +++ b/wallet/src/pinata_interactions.rs @@ -0,0 +1,161 @@ +use common::{error::ExecutionFailureKind, sequencer_client::json::SendTxResponse}; +use key_protocol::key_management::ephemeral_key_holder::EphemeralKeyHolder; +use nssa::{AccountId, privacy_preserving_transaction::circuit}; +use nssa_core::{MembershipProof, SharedSecretKey, account::AccountWithMetadata}; + +use crate::{ + WalletCore, helperfunctions::produce_random_nonces, transaction_utils::AccountPreparedData, +}; + +impl WalletCore { + pub async fn claim_pinata( + &self, + pinata_account_id: AccountId, + winner_account_id: AccountId, + solution: u128, + ) -> Result { + let account_ids = vec![pinata_account_id, winner_account_id]; + let program_id = nssa::program::Program::pinata().id(); + let message = + nssa::public_transaction::Message::try_new(program_id, account_ids, vec![], solution) + .unwrap(); + + let witness_set = nssa::public_transaction::WitnessSet::for_message(&message, &[]); + let tx = nssa::PublicTransaction::new(message, witness_set); + + Ok(self.sequencer_client.send_tx_public(tx).await?) + } + + pub async fn claim_pinata_private_owned_account_already_initialized( + &self, + pinata_account_id: AccountId, + winner_account_id: AccountId, + solution: u128, + winner_proof: MembershipProof, + ) -> Result<(SendTxResponse, [SharedSecretKey; 1]), ExecutionFailureKind> { + let AccountPreparedData { + nsk: winner_nsk, + npk: winner_npk, + ipk: winner_ipk, + auth_acc: winner_pre, + proof: _, + } = self + .private_acc_preparation(winner_account_id, true, false) + .await?; + + let pinata_acc = self.get_account_public(pinata_account_id).await.unwrap(); + + let program = nssa::program::Program::pinata(); + + let pinata_pre = AccountWithMetadata::new(pinata_acc.clone(), false, pinata_account_id); + + let eph_holder_winner = EphemeralKeyHolder::new(&winner_npk); + let shared_secret_winner = eph_holder_winner.calculate_shared_secret_sender(&winner_ipk); + + let (output, proof) = circuit::execute_and_prove( + &[pinata_pre, winner_pre], + &nssa::program::Program::serialize_instruction(solution).unwrap(), + &[0, 1], + &produce_random_nonces(1), + &[(winner_npk.clone(), shared_secret_winner.clone())], + &[(winner_nsk.unwrap(), winner_proof)], + &program.into(), + ) + .unwrap(); + + let message = + nssa::privacy_preserving_transaction::message::Message::try_from_circuit_output( + vec![pinata_account_id], + vec![], + vec![( + winner_npk.clone(), + winner_ipk.clone(), + eph_holder_winner.generate_ephemeral_public_key(), + )], + output, + ) + .unwrap(); + + let witness_set = + nssa::privacy_preserving_transaction::witness_set::WitnessSet::for_message( + &message, + proof, + &[], + ); + let tx = nssa::privacy_preserving_transaction::PrivacyPreservingTransaction::new( + message, + witness_set, + ); + + Ok(( + self.sequencer_client.send_tx_private(tx).await?, + [shared_secret_winner], + )) + } + + pub async fn claim_pinata_private_owned_account_not_initialized( + &self, + pinata_account_id: AccountId, + winner_account_id: AccountId, + solution: u128, + ) -> Result<(SendTxResponse, [SharedSecretKey; 1]), ExecutionFailureKind> { + let AccountPreparedData { + nsk: _, + npk: winner_npk, + ipk: winner_ipk, + auth_acc: winner_pre, + proof: _, + } = self + .private_acc_preparation(winner_account_id, false, false) + .await?; + + let pinata_acc = self.get_account_public(pinata_account_id).await.unwrap(); + + let program = nssa::program::Program::pinata(); + + let pinata_pre = AccountWithMetadata::new(pinata_acc.clone(), false, pinata_account_id); + + let eph_holder_winner = EphemeralKeyHolder::new(&winner_npk); + let shared_secret_winner = eph_holder_winner.calculate_shared_secret_sender(&winner_ipk); + + let (output, proof) = circuit::execute_and_prove( + &[pinata_pre, winner_pre], + &nssa::program::Program::serialize_instruction(solution).unwrap(), + &[0, 2], + &produce_random_nonces(1), + &[(winner_npk.clone(), shared_secret_winner.clone())], + &[], + &program.into(), + ) + .unwrap(); + + let message = + nssa::privacy_preserving_transaction::message::Message::try_from_circuit_output( + vec![pinata_account_id], + vec![], + vec![( + winner_npk.clone(), + winner_ipk.clone(), + eph_holder_winner.generate_ephemeral_public_key(), + )], + output, + ) + .unwrap(); + + let witness_set = + nssa::privacy_preserving_transaction::witness_set::WitnessSet::for_message( + &message, + proof, + &[], + ); + let tx = nssa::privacy_preserving_transaction::PrivacyPreservingTransaction::new( + message, + witness_set, + ); + + Ok(( + self.sequencer_client.send_tx_private(tx).await?, + [shared_secret_winner], + )) + } +} diff --git a/wallet/src/program_facades/token.rs b/wallet/src/program_facades/token.rs index 7c971551..4ec9c127 100644 --- a/wallet/src/program_facades/token.rs +++ b/wallet/src/program_facades/token.rs @@ -38,7 +38,7 @@ impl Token<'_> { Ok(self.0.sequencer_client.send_tx_public(tx).await?) } - pub async fn send_new_definition_private_owned( + pub async fn send_new_definition_private_owned_supply( &self, definition_account_id: AccountId, supply_account_id: AccountId, @@ -61,11 +61,66 @@ impl Token<'_> { let first = secrets .into_iter() .next() - .expect("expected recipient's secret"); + .expect("expected supply's secret"); (resp, first) }) } + pub async fn send_new_definition_private_owned_definiton( + &self, + definition_account_id: AccountId, + supply_account_id: AccountId, + name: [u8; 6], + total_supply: u128, + ) -> Result<(SendTxResponse, SharedSecretKey), ExecutionFailureKind> { + let (instruction_data, program) = token_program_preparation_definition(name, total_supply); + + self.0 + .send_privacy_preserving_tx( + vec![ + PrivacyPreservingAccount::PrivateOwned(definition_account_id), + PrivacyPreservingAccount::Public(supply_account_id), + ], + &instruction_data, + &program, + ) + .await + .map(|(resp, secrets)| { + let first = secrets + .into_iter() + .next() + .expect("expected definition's secret"); + (resp, first) + }) + } + + pub async fn send_new_definition_private_owned_definiton_and_supply( + &self, + definition_account_id: AccountId, + supply_account_id: AccountId, + name: [u8; 6], + total_supply: u128, + ) -> Result<(SendTxResponse, [SharedSecretKey; 2]), ExecutionFailureKind> { + let (instruction_data, program) = token_program_preparation_definition(name, total_supply); + + self.0 + .send_privacy_preserving_tx( + vec![ + PrivacyPreservingAccount::PrivateOwned(definition_account_id), + PrivacyPreservingAccount::PrivateOwned(supply_account_id), + ], + &instruction_data, + &program, + ) + .await + .map(|(resp, secrets)| { + let mut iter = secrets.into_iter(); + let first = iter.next().expect("expected definition's secret"); + let second = iter.next().expect("expected supply's secret"); + (resp, [first, second]) + }) + } + pub async fn send_transfer_transaction( &self, sender_account_id: AccountId, diff --git a/wallet/src/transaction_utils.rs b/wallet/src/transaction_utils.rs new file mode 100644 index 00000000..123e49d3 --- /dev/null +++ b/wallet/src/transaction_utils.rs @@ -0,0 +1,592 @@ +use common::{error::ExecutionFailureKind, sequencer_client::json::SendTxResponse}; +use key_protocol::key_management::ephemeral_key_holder::EphemeralKeyHolder; +use nssa::{ + Account, AccountId, PrivacyPreservingTransaction, + privacy_preserving_transaction::{circuit, message::Message, witness_set::WitnessSet}, + program::Program, +}; +use nssa_core::{ + Commitment, MembershipProof, NullifierPublicKey, NullifierSecretKey, SharedSecretKey, + account::AccountWithMetadata, encryption::IncomingViewingPublicKey, program::InstructionData, +}; + +use crate::{WalletCore, helperfunctions::produce_random_nonces}; + +pub(crate) struct AccountPreparedData { + pub nsk: Option, + pub npk: NullifierPublicKey, + pub ipk: IncomingViewingPublicKey, + pub auth_acc: AccountWithMetadata, + pub proof: Option, +} + +impl WalletCore { + pub(crate) async fn private_acc_preparation( + &self, + account_id: AccountId, + is_authorized: bool, + needs_proof: bool, + ) -> Result { + let Some((from_keys, from_acc)) = self + .storage + .user_data + .get_private_account(&account_id) + .cloned() + else { + return Err(ExecutionFailureKind::KeyNotFoundError); + }; + + let mut nsk = None; + let mut proof = None; + + let from_npk = from_keys.nullifer_public_key; + let from_ipk = from_keys.incoming_viewing_public_key; + + let sender_commitment = Commitment::new(&from_npk, &from_acc); + + let sender_pre = AccountWithMetadata::new(from_acc.clone(), is_authorized, &from_npk); + + if is_authorized { + nsk = Some(from_keys.private_key_holder.nullifier_secret_key); + } + + if needs_proof { + proof = self + .sequencer_client + .get_proof_for_commitment(sender_commitment) + .await + .unwrap(); + } + + Ok(AccountPreparedData { + nsk, + npk: from_npk, + ipk: from_ipk, + auth_acc: sender_pre, + proof, + }) + } + + pub(crate) async fn private_tx_two_accs_all_init( + &self, + from: AccountId, + to: AccountId, + instruction_data: InstructionData, + tx_pre_check: impl FnOnce(&Account, &Account) -> Result<(), ExecutionFailureKind>, + program: Program, + to_proof: MembershipProof, + ) -> Result<(SendTxResponse, [SharedSecretKey; 2]), ExecutionFailureKind> { + let AccountPreparedData { + nsk: from_nsk, + npk: from_npk, + ipk: from_ipk, + auth_acc: sender_pre, + proof: from_proof, + } = self.private_acc_preparation(from, true, true).await?; + + let AccountPreparedData { + nsk: to_nsk, + npk: to_npk, + ipk: to_ipk, + auth_acc: recipient_pre, + proof: _, + } = self.private_acc_preparation(to, true, false).await?; + + tx_pre_check(&sender_pre.account, &recipient_pre.account)?; + + let eph_holder_from = EphemeralKeyHolder::new(&from_npk); + let shared_secret_from = eph_holder_from.calculate_shared_secret_sender(&from_ipk); + + let eph_holder_to = EphemeralKeyHolder::new(&to_npk); + let shared_secret_to = eph_holder_to.calculate_shared_secret_sender(&to_ipk); + + let (output, proof) = circuit::execute_and_prove( + &[sender_pre, recipient_pre], + &instruction_data, + &[1, 1], + &produce_random_nonces(2), + &[ + (from_npk.clone(), shared_secret_from.clone()), + (to_npk.clone(), shared_secret_to.clone()), + ], + &[ + (from_nsk.unwrap(), from_proof.unwrap()), + (to_nsk.unwrap(), to_proof), + ], + &program.into(), + ) + .unwrap(); + + let message = Message::try_from_circuit_output( + vec![], + vec![], + vec![ + ( + from_npk.clone(), + from_ipk.clone(), + eph_holder_from.generate_ephemeral_public_key(), + ), + ( + to_npk.clone(), + to_ipk.clone(), + eph_holder_to.generate_ephemeral_public_key(), + ), + ], + output, + ) + .unwrap(); + + let witness_set = WitnessSet::for_message(&message, proof, &[]); + let tx = PrivacyPreservingTransaction::new(message, witness_set); + + Ok(( + self.sequencer_client.send_tx_private(tx).await?, + [shared_secret_from, shared_secret_to], + )) + } + + pub(crate) async fn private_tx_two_accs_receiver_uninit( + &self, + from: AccountId, + to: AccountId, + instruction_data: InstructionData, + tx_pre_check: impl FnOnce(&Account, &Account) -> Result<(), ExecutionFailureKind>, + program: Program, + ) -> Result<(SendTxResponse, [SharedSecretKey; 2]), ExecutionFailureKind> { + let AccountPreparedData { + nsk: from_nsk, + npk: from_npk, + ipk: from_ipk, + auth_acc: sender_pre, + proof: from_proof, + } = self.private_acc_preparation(from, true, true).await?; + + let AccountPreparedData { + nsk: _, + npk: to_npk, + ipk: to_ipk, + auth_acc: recipient_pre, + proof: _, + } = self.private_acc_preparation(to, false, false).await?; + + tx_pre_check(&sender_pre.account, &recipient_pre.account)?; + + let eph_holder_from = EphemeralKeyHolder::new(&from_npk); + let shared_secret_from = eph_holder_from.calculate_shared_secret_sender(&from_ipk); + + let eph_holder_to = EphemeralKeyHolder::new(&to_npk); + let shared_secret_to = eph_holder_to.calculate_shared_secret_sender(&to_ipk); + + let (output, proof) = circuit::execute_and_prove( + &[sender_pre, recipient_pre], + &instruction_data, + &[1, 2], + &produce_random_nonces(2), + &[ + (from_npk.clone(), shared_secret_from.clone()), + (to_npk.clone(), shared_secret_to.clone()), + ], + &[(from_nsk.unwrap(), from_proof.unwrap())], + &program.into(), + ) + .unwrap(); + + let message = Message::try_from_circuit_output( + vec![], + vec![], + vec![ + ( + from_npk.clone(), + from_ipk.clone(), + eph_holder_from.generate_ephemeral_public_key(), + ), + ( + to_npk.clone(), + to_ipk.clone(), + eph_holder_to.generate_ephemeral_public_key(), + ), + ], + output, + ) + .unwrap(); + + let witness_set = WitnessSet::for_message(&message, proof, &[]); + let tx = PrivacyPreservingTransaction::new(message, witness_set); + + Ok(( + self.sequencer_client.send_tx_private(tx).await?, + [shared_secret_from, shared_secret_to], + )) + } + + pub(crate) async fn private_tx_two_accs_receiver_outer( + &self, + from: AccountId, + to_npk: NullifierPublicKey, + to_ipk: IncomingViewingPublicKey, + instruction_data: InstructionData, + tx_pre_check: impl FnOnce(&Account, &Account) -> Result<(), ExecutionFailureKind>, + program: Program, + ) -> Result<(SendTxResponse, [SharedSecretKey; 2]), ExecutionFailureKind> { + let AccountPreparedData { + nsk: from_nsk, + npk: from_npk, + ipk: from_ipk, + auth_acc: sender_pre, + proof: from_proof, + } = self.private_acc_preparation(from, true, true).await?; + + let to_acc = nssa_core::account::Account::default(); + + tx_pre_check(&sender_pre.account, &to_acc)?; + + let recipient_pre = AccountWithMetadata::new(to_acc.clone(), false, &to_npk); + + let eph_holder = EphemeralKeyHolder::new(&to_npk); + + let shared_secret_from = eph_holder.calculate_shared_secret_sender(&from_ipk); + let shared_secret_to = eph_holder.calculate_shared_secret_sender(&to_ipk); + + let (output, proof) = circuit::execute_and_prove( + &[sender_pre, recipient_pre], + &instruction_data, + &[1, 2], + &produce_random_nonces(2), + &[ + (from_npk.clone(), shared_secret_from.clone()), + (to_npk.clone(), shared_secret_to.clone()), + ], + &[(from_nsk.unwrap(), from_proof.unwrap())], + &program.into(), + ) + .unwrap(); + + let message = Message::try_from_circuit_output( + vec![], + vec![], + vec![ + ( + from_npk.clone(), + from_ipk.clone(), + eph_holder.generate_ephemeral_public_key(), + ), + ( + to_npk.clone(), + to_ipk.clone(), + eph_holder.generate_ephemeral_public_key(), + ), + ], + output, + ) + .unwrap(); + + let witness_set = WitnessSet::for_message(&message, proof, &[]); + + let tx = PrivacyPreservingTransaction::new(message, witness_set); + + Ok(( + self.sequencer_client.send_tx_private(tx).await?, + [shared_secret_from, shared_secret_to], + )) + } + + pub(crate) async fn deshielded_tx_two_accs( + &self, + from: AccountId, + to: AccountId, + instruction_data: InstructionData, + tx_pre_check: impl FnOnce(&Account, &Account) -> Result<(), ExecutionFailureKind>, + program: Program, + ) -> Result<(SendTxResponse, [nssa_core::SharedSecretKey; 1]), ExecutionFailureKind> { + let AccountPreparedData { + nsk: from_nsk, + npk: from_npk, + ipk: from_ipk, + auth_acc: sender_pre, + proof: from_proof, + } = self.private_acc_preparation(from, true, true).await?; + + let Ok(to_acc) = self.get_account_public(to).await else { + return Err(ExecutionFailureKind::KeyNotFoundError); + }; + + tx_pre_check(&sender_pre.account, &to_acc)?; + + let recipient_pre = AccountWithMetadata::new(to_acc.clone(), false, to); + + let eph_holder = EphemeralKeyHolder::new(&from_npk); + let shared_secret = eph_holder.calculate_shared_secret_sender(&from_ipk); + + let (output, proof) = circuit::execute_and_prove( + &[sender_pre, recipient_pre], + &instruction_data, + &[1, 0], + &produce_random_nonces(1), + &[(from_npk.clone(), shared_secret.clone())], + &[(from_nsk.unwrap(), from_proof.unwrap())], + &program.into(), + ) + .unwrap(); + + let message = Message::try_from_circuit_output( + vec![to], + vec![], + vec![( + from_npk.clone(), + from_ipk.clone(), + eph_holder.generate_ephemeral_public_key(), + )], + output, + ) + .unwrap(); + + let witness_set = WitnessSet::for_message(&message, proof, &[]); + + let tx = PrivacyPreservingTransaction::new(message, witness_set); + + Ok(( + self.sequencer_client.send_tx_private(tx).await?, + [shared_secret], + )) + } + + pub(crate) async fn shielded_two_accs_all_init( + &self, + from: AccountId, + to: AccountId, + instruction_data: InstructionData, + tx_pre_check: impl FnOnce(&Account, &Account) -> Result<(), ExecutionFailureKind>, + program: Program, + to_proof: MembershipProof, + ) -> Result<(SendTxResponse, [SharedSecretKey; 1]), ExecutionFailureKind> { + let Ok(from_acc) = self.get_account_public(from).await else { + return Err(ExecutionFailureKind::KeyNotFoundError); + }; + + let AccountPreparedData { + nsk: to_nsk, + npk: to_npk, + ipk: to_ipk, + auth_acc: recipient_pre, + proof: _, + } = self.private_acc_preparation(to, true, false).await?; + + tx_pre_check(&from_acc, &recipient_pre.account)?; + + let sender_pre = AccountWithMetadata::new(from_acc.clone(), true, from); + + let eph_holder = EphemeralKeyHolder::new(&to_npk); + let shared_secret = eph_holder.calculate_shared_secret_sender(&to_ipk); + + let (output, proof) = circuit::execute_and_prove( + &[sender_pre, recipient_pre], + &instruction_data, + &[0, 1], + &produce_random_nonces(1), + &[(to_npk.clone(), shared_secret.clone())], + &[(to_nsk.unwrap(), to_proof)], + &program.into(), + ) + .unwrap(); + + let message = Message::try_from_circuit_output( + vec![from], + vec![from_acc.nonce], + vec![( + to_npk.clone(), + to_ipk.clone(), + eph_holder.generate_ephemeral_public_key(), + )], + output, + ) + .unwrap(); + + let signing_key = self.storage.user_data.get_pub_account_signing_key(&from); + + let Some(signing_key) = signing_key else { + return Err(ExecutionFailureKind::KeyNotFoundError); + }; + + let witness_set = WitnessSet::for_message(&message, proof, &[signing_key]); + + let tx = PrivacyPreservingTransaction::new(message, witness_set); + + Ok(( + self.sequencer_client.send_tx_private(tx).await?, + [shared_secret], + )) + } + + pub(crate) async fn shielded_two_accs_receiver_uninit( + &self, + from: AccountId, + to: AccountId, + instruction_data: InstructionData, + tx_pre_check: impl FnOnce(&Account, &Account) -> Result<(), ExecutionFailureKind>, + program: Program, + ) -> Result<(SendTxResponse, [SharedSecretKey; 1]), ExecutionFailureKind> { + let Ok(from_acc) = self.get_account_public(from).await else { + return Err(ExecutionFailureKind::KeyNotFoundError); + }; + + let AccountPreparedData { + nsk: _, + npk: to_npk, + ipk: to_ipk, + auth_acc: recipient_pre, + proof: _, + } = self.private_acc_preparation(to, false, false).await?; + + tx_pre_check(&from_acc, &recipient_pre.account)?; + + let sender_pre = AccountWithMetadata::new(from_acc.clone(), true, from); + + let eph_holder = EphemeralKeyHolder::new(&to_npk); + let shared_secret = eph_holder.calculate_shared_secret_sender(&to_ipk); + + let (output, proof) = circuit::execute_and_prove( + &[sender_pre, recipient_pre], + &instruction_data, + &[0, 2], + &produce_random_nonces(1), + &[(to_npk.clone(), shared_secret.clone())], + &[], + &program.into(), + ) + .unwrap(); + + let message = Message::try_from_circuit_output( + vec![from], + vec![from_acc.nonce], + vec![( + to_npk.clone(), + to_ipk.clone(), + eph_holder.generate_ephemeral_public_key(), + )], + output, + ) + .unwrap(); + + let signing_key = self.storage.user_data.get_pub_account_signing_key(&from); + + let Some(signing_key) = signing_key else { + return Err(ExecutionFailureKind::KeyNotFoundError); + }; + + let witness_set = WitnessSet::for_message(&message, proof, &[signing_key]); + + let tx = PrivacyPreservingTransaction::new(message, witness_set); + + Ok(( + self.sequencer_client.send_tx_private(tx).await?, + [shared_secret], + )) + } + + pub(crate) async fn shielded_two_accs_receiver_outer( + &self, + from: AccountId, + to_npk: NullifierPublicKey, + to_ipk: IncomingViewingPublicKey, + instruction_data: InstructionData, + tx_pre_check: impl FnOnce(&Account, &Account) -> Result<(), ExecutionFailureKind>, + program: Program, + ) -> Result { + let Ok(from_acc) = self.get_account_public(from).await else { + return Err(ExecutionFailureKind::KeyNotFoundError); + }; + + let to_acc = Account::default(); + + tx_pre_check(&from_acc, &to_acc)?; + + let sender_pre = AccountWithMetadata::new(from_acc.clone(), true, from); + let recipient_pre = AccountWithMetadata::new(to_acc.clone(), false, &to_npk); + + let eph_holder = EphemeralKeyHolder::new(&to_npk); + let shared_secret = eph_holder.calculate_shared_secret_sender(&to_ipk); + + let (output, proof) = circuit::execute_and_prove( + &[sender_pre, recipient_pre], + &instruction_data, + &[0, 2], + &produce_random_nonces(1), + &[(to_npk.clone(), shared_secret.clone())], + &[], + &program.into(), + ) + .unwrap(); + + let message = Message::try_from_circuit_output( + vec![from], + vec![from_acc.nonce], + vec![( + to_npk.clone(), + to_ipk.clone(), + eph_holder.generate_ephemeral_public_key(), + )], + output, + ) + .unwrap(); + + let signing_key = self.storage.user_data.get_pub_account_signing_key(&from); + + let Some(signing_key) = signing_key else { + return Err(ExecutionFailureKind::KeyNotFoundError); + }; + + let witness_set = WitnessSet::for_message(&message, proof, &[signing_key]); + let tx = PrivacyPreservingTransaction::new(message, witness_set); + + Ok(self.sequencer_client.send_tx_private(tx).await?) + } + + pub async fn register_account_under_authenticated_transfers_programs_private( + &self, + from: AccountId, + ) -> Result<(SendTxResponse, [SharedSecretKey; 1]), ExecutionFailureKind> { + let AccountPreparedData { + nsk: _, + npk: from_npk, + ipk: from_ipk, + auth_acc: sender_pre, + proof: _, + } = self.private_acc_preparation(from, false, false).await?; + + let eph_holder_from = EphemeralKeyHolder::new(&from_npk); + let shared_secret_from = eph_holder_from.calculate_shared_secret_sender(&from_ipk); + + let instruction: u128 = 0; + + let (output, proof) = circuit::execute_and_prove( + &[sender_pre], + &Program::serialize_instruction(instruction).unwrap(), + &[2], + &produce_random_nonces(1), + &[(from_npk.clone(), shared_secret_from.clone())], + &[], + &Program::authenticated_transfer_program().into(), + ) + .unwrap(); + + let message = Message::try_from_circuit_output( + vec![], + vec![], + vec![( + from_npk.clone(), + from_ipk.clone(), + eph_holder_from.generate_ephemeral_public_key(), + )], + output, + ) + .unwrap(); + + let witness_set = WitnessSet::for_message(&message, proof, &[]); + let tx = PrivacyPreservingTransaction::new(message, witness_set); + + Ok(( + self.sequencer_client.send_tx_private(tx).await?, + [shared_secret_from], + )) + } +}