refactor: small adjustments to privacy preserving tx sending

wallet/src/cli/programs/native_token_transfer.rs

@@ -75,10 +75,8 @@ impl WalletSubcommand for AuthTransferSubcommand {
                     AccountPrivacyKind::Private => {
                         let account_id = account_id.parse()?;
 
-                        let (res, [secret]) = wallet_core
+                        let (res, secret) = NativeTokenTransfer(wallet_core)
-                            .register_account_under_authenticated_transfers_programs_private(
+                            .register_account_private(account_id)
-                                account_id,
-                            )
                             .await?;
 
                         println!("Results of tx send is {res:#?}");

wallet/src/lib.rs

@@ -32,7 +32,6 @@ pub mod helperfunctions;
 pub mod poller;
 mod privacy_preserving_tx;
 pub mod program_facades;
-pub mod transaction_utils;
 
 pub struct WalletCore {
     pub storage: WalletChainStore,
@@ -214,12 +213,24 @@ impl WalletCore {
         &self,
         accounts: Vec<PrivacyPreservingAccount>,
         instruction_data: &InstructionData,
-        tx_pre_check: impl FnOnce(&[&Account]) -> Result<(), ExecutionFailureKind>,
         program: &Program,
     ) -> Result<(SendTxResponse, Vec<SharedSecretKey>), ExecutionFailureKind> {
-        let payload = privacy_preserving_tx::Payload::new(self, accounts).await?;
+        self.send_privacy_preserving_tx_with_pre_check(accounts, instruction_data, program, |_| {
+            Ok(())
+        })
+        .await
+    }
 
-        let pre_states = payload.pre_states();
+    pub async fn send_privacy_preserving_tx_with_pre_check(
+        &self,
+        accounts: Vec<PrivacyPreservingAccount>,
+        instruction_data: &InstructionData,
+        program: &Program,
+        tx_pre_check: impl FnOnce(&[&Account]) -> Result<(), ExecutionFailureKind>,
+    ) -> Result<(SendTxResponse, Vec<SharedSecretKey>), ExecutionFailureKind> {
+        let acc_manager = privacy_preserving_tx::AccountManager::new(self, accounts).await?;
+
+        let pre_states = acc_manager.pre_states();
         tx_pre_check(
             &pre_states
                 .iter()
@@ -227,25 +238,25 @@ impl WalletCore {
                 .collect::<Vec<_>>(),
         )?;
 
-        let private_account_keys = payload.private_account_keys();
+        let private_account_keys = acc_manager.private_account_keys();
         let (output, proof) = nssa::privacy_preserving_transaction::circuit::execute_and_prove(
             &pre_states,
             instruction_data,
-            payload.visibility_mask(),
+            acc_manager.visibility_mask(),
             &produce_random_nonces(private_account_keys.len()),
             &private_account_keys
                 .iter()
                 .map(|keys| (keys.npk.clone(), keys.ssk.clone()))
                 .collect::<Vec<_>>(),
-            &payload.private_account_auth(),
+            &acc_manager.private_account_auth(),
             program,
         )
         .unwrap();
 
         let message =
             nssa::privacy_preserving_transaction::message::Message::try_from_circuit_output(
-                payload.public_account_ids(),
+                acc_manager.public_account_ids(),
-                Vec::from_iter(payload.public_account_nonces()),
+                Vec::from_iter(acc_manager.public_account_nonces()),
                 private_account_keys
                     .iter()
                     .map(|keys| (keys.npk.clone(), keys.ipk.clone(), keys.epk.clone()))
@@ -258,7 +269,7 @@ impl WalletCore {
             nssa::privacy_preserving_transaction::witness_set::WitnessSet::for_message(
                 &message,
                 proof,
-                &payload.witness_signing_keys(),
+                &acc_manager.witness_signing_keys(),
             );
         let tx = PrivacyPreservingTransaction::new(message, witness_set);
 

wallet/src/privacy_preserving_tx.rs

@@ -7,7 +7,7 @@ use nssa_core::{
     encryption::{EphemeralPublicKey, IncomingViewingPublicKey},
 };
 
-use crate::{WalletCore, transaction_utils::AccountPreparedData};
+use crate::WalletCore;
 
 pub enum PrivacyPreservingAccount {
     Public(AccountId),
@@ -33,12 +33,12 @@ enum State {
     Private(AccountPreparedData),
 }
 
-pub struct Payload {
+pub struct AccountManager {
     states: Vec<State>,
     visibility_mask: Vec<u8>,
 }
 
-impl Payload {
+impl AccountManager {
     pub async fn new(
         wallet: &WalletCore,
         accounts: Vec<PrivacyPreservingAccount>,
@@ -60,16 +60,8 @@ impl Payload {
                     (State::Public { account, sk }, 0)
                 }
                 PrivacyPreservingAccount::PrivateOwned(account_id) => {
-                    let mut pre = wallet
+                    let pre = private_acc_preparation(wallet, account_id).await?;
-                        .private_acc_preparation(account_id, true, true)
+                    let mask = if pre.auth_acc.is_authorized { 1 } else { 2 };
-                        .await?;
-                    let mut mask = 1;
-
-                    if pre.proof.is_none() {
-                        pre.auth_acc.is_authorized = false;
-                        pre.nsk = None;
-                        mask = 2
-                    };
 
                     (State::Private(pre), mask)
                 }
@@ -116,7 +108,7 @@ impl Payload {
         self.states
             .iter()
             .filter_map(|state| match state {
-                State::Public { account, .. } => Some(account.account.nonce),
+                State::Public { account, sk } => sk.as_ref().map(|_| account.account.nonce),
                 _ => None,
             })
             .collect()
@@ -171,3 +163,50 @@ impl Payload {
             .collect()
     }
 }
+
+struct AccountPreparedData {
+    nsk: Option<NullifierSecretKey>,
+    npk: NullifierPublicKey,
+    ipk: IncomingViewingPublicKey,
+    auth_acc: AccountWithMetadata,
+    proof: Option<MembershipProof>,
+}
+
+async fn private_acc_preparation(
+    wallet: &WalletCore,
+    account_id: AccountId,
+) -> Result<AccountPreparedData, ExecutionFailureKind> {
+    let Some((from_keys, from_acc)) = wallet
+        .storage
+        .user_data
+        .get_private_account(&account_id)
+        .cloned()
+    else {
+        return Err(ExecutionFailureKind::KeyNotFoundError);
+    };
+
+    let mut nsk = Some(from_keys.private_key_holder.nullifier_secret_key);
+
+    let from_npk = from_keys.nullifer_public_key;
+    let from_ipk = from_keys.incoming_viewing_public_key;
+
+    // TODO: Remove this unwrap, error types must be compatible
+    let proof = wallet
+        .check_private_account_initialized(&account_id)
+        .await
+        .unwrap();
+
+    if proof.is_none() {
+        nsk = None;
+    }
+
+    let sender_pre = AccountWithMetadata::new(from_acc.clone(), proof.is_some(), &from_npk);
+
+    Ok(AccountPreparedData {
+        nsk,
+        npk: from_npk,
+        ipk: from_ipk,
+        auth_acc: sender_pre,
+        proof,
+    })
+}

wallet/src/program_facades/native_token_transfer/deshielded.rs

@@ -14,14 +14,14 @@ impl NativeTokenTransfer<'_> {
         let (instruction_data, program, tx_pre_check) = auth_transfer_preparation(balance_to_move);
 
         self.0
-            .send_privacy_preserving_tx(
+            .send_privacy_preserving_tx_with_pre_check(
                 vec![
                     PrivacyPreservingAccount::PrivateOwned(from),
                     PrivacyPreservingAccount::Public(to),
                 ],
                 &instruction_data,
-                tx_pre_check,
                 &program,
+                tx_pre_check,
             )
             .await
             .map(|(resp, secrets)| {

wallet/src/program_facades/native_token_transfer/private.rs

@@ -1,13 +1,34 @@
 use std::vec;
 
 use common::{error::ExecutionFailureKind, sequencer_client::json::SendTxResponse};
-use nssa::AccountId;
+use nssa::{AccountId, program::Program};
 use nssa_core::{NullifierPublicKey, SharedSecretKey, encryption::IncomingViewingPublicKey};
 
 use super::{NativeTokenTransfer, auth_transfer_preparation};
 use crate::PrivacyPreservingAccount;
 
 impl NativeTokenTransfer<'_> {
+    pub async fn register_account_private(
+        &self,
+        from: AccountId,
+    ) -> Result<(SendTxResponse, SharedSecretKey), ExecutionFailureKind> {
+        let instruction: u128 = 0;
+
+        self.0
+            .send_privacy_preserving_tx_with_pre_check(
+                vec![PrivacyPreservingAccount::PrivateOwned(from)],
+                &Program::serialize_instruction(instruction).unwrap(),
+                &Program::authenticated_transfer_program(),
+                |_| Ok(()),
+            )
+            .await
+            .map(|(resp, secrets)| {
+                let mut secrets_iter = secrets.into_iter();
+                let first = secrets_iter.next().expect("expected sender's secret");
+                (resp, first)
+            })
+    }
+
     pub async fn send_private_transfer_to_outer_account(
         &self,
         from: AccountId,
@@ -18,7 +39,7 @@ impl NativeTokenTransfer<'_> {
         let (instruction_data, program, tx_pre_check) = auth_transfer_preparation(balance_to_move);
 
         self.0
-            .send_privacy_preserving_tx(
+            .send_privacy_preserving_tx_with_pre_check(
                 vec![
                     PrivacyPreservingAccount::PrivateOwned(from),
                     PrivacyPreservingAccount::PrivateForeign {
@@ -27,8 +48,8 @@ impl NativeTokenTransfer<'_> {
                     },
                 ],
                 &instruction_data,
-                tx_pre_check,
                 &program,
+                tx_pre_check,
             )
             .await
             .map(|(resp, secrets)| {
@@ -48,14 +69,14 @@ impl NativeTokenTransfer<'_> {
         let (instruction_data, program, tx_pre_check) = auth_transfer_preparation(balance_to_move);
 
         self.0
-            .send_privacy_preserving_tx(
+            .send_privacy_preserving_tx_with_pre_check(
                 vec![
                     PrivacyPreservingAccount::PrivateOwned(from),
                     PrivacyPreservingAccount::PrivateOwned(to),
                 ],
                 &instruction_data,
-                tx_pre_check,
                 &program,
+                tx_pre_check,
             )
             .await
             .map(|(resp, secrets)| {

wallet/src/program_facades/native_token_transfer/shielded.rs

@@ -15,14 +15,14 @@ impl NativeTokenTransfer<'_> {
         let (instruction_data, program, tx_pre_check) = auth_transfer_preparation(balance_to_move);
 
         self.0
-            .send_privacy_preserving_tx(
+            .send_privacy_preserving_tx_with_pre_check(
                 vec![
                     PrivacyPreservingAccount::Public(from),
                     PrivacyPreservingAccount::PrivateOwned(to),
                 ],
                 &instruction_data,
-                tx_pre_check,
                 &program,
+                tx_pre_check,
             )
             .await
             .map(|(resp, secrets)| {
@@ -44,7 +44,7 @@ impl NativeTokenTransfer<'_> {
         let (instruction_data, program, tx_pre_check) = auth_transfer_preparation(balance_to_move);
 
         self.0
-            .send_privacy_preserving_tx(
+            .send_privacy_preserving_tx_with_pre_check(
                 vec![
                     PrivacyPreservingAccount::Public(from),
                     PrivacyPreservingAccount::PrivateForeign {
@@ -53,8 +53,8 @@ impl NativeTokenTransfer<'_> {
                     },
                 ],
                 &instruction_data,
-                tx_pre_check,
                 &program,
+                tx_pre_check,
             )
             .await
             .map(|(resp, secrets)| {

wallet/src/program_facades/pinata.rs

@@ -38,7 +38,6 @@ impl Pinata<'_> {
                     PrivacyPreservingAccount::PrivateOwned(winner_account_id),
                 ],
                 &nssa::program::Program::serialize_instruction(solution).unwrap(),
-                |_| Ok(()),
                 &nssa::program::Program::pinata(),
             )
             .await

wallet/src/program_facades/token.rs

@@ -1,5 +1,5 @@
 use common::{error::ExecutionFailureKind, sequencer_client::json::SendTxResponse};
-use nssa::{Account, AccountId, program::Program};
+use nssa::{AccountId, program::Program};
 use nssa_core::{
     NullifierPublicKey, SharedSecretKey, encryption::IncomingViewingPublicKey,
     program::InstructionData,
@@ -45,8 +45,7 @@ impl Token<'_> {
         name: [u8; 6],
         total_supply: u128,
     ) -> Result<(SendTxResponse, SharedSecretKey), ExecutionFailureKind> {
-        let (instruction_data, program, tx_pre_check) =
+        let (instruction_data, program) = token_program_preparation_definition(name, total_supply);
-            token_program_preparation_definition(name, total_supply);
 
         self.0
             .send_privacy_preserving_tx(
@@ -55,7 +54,6 @@ impl Token<'_> {
                     PrivacyPreservingAccount::PrivateOwned(supply_account_id),
                 ],
                 &instruction_data,
-                tx_pre_check,
                 &program,
             )
             .await
@@ -114,7 +112,7 @@ impl Token<'_> {
         recipient_account_id: AccountId,
         amount: u128,
     ) -> Result<(SendTxResponse, [SharedSecretKey; 2]), ExecutionFailureKind> {
-        let (instruction_data, program, tx_pre_check) = token_program_preparation_transfer(amount);
+        let (instruction_data, program) = token_program_preparation_transfer(amount);
 
         self.0
             .send_privacy_preserving_tx(
@@ -123,7 +121,6 @@ impl Token<'_> {
                     PrivacyPreservingAccount::PrivateOwned(recipient_account_id),
                 ],
                 &instruction_data,
-                tx_pre_check,
                 &program,
             )
             .await
@@ -142,7 +139,7 @@ impl Token<'_> {
         recipient_ipk: IncomingViewingPublicKey,
         amount: u128,
     ) -> Result<(SendTxResponse, [SharedSecretKey; 2]), ExecutionFailureKind> {
-        let (instruction_data, program, tx_pre_check) = token_program_preparation_transfer(amount);
+        let (instruction_data, program) = token_program_preparation_transfer(amount);
 
         self.0
             .send_privacy_preserving_tx(
@@ -154,7 +151,6 @@ impl Token<'_> {
                     },
                 ],
                 &instruction_data,
-                tx_pre_check,
                 &program,
             )
             .await
@@ -172,7 +168,7 @@ impl Token<'_> {
         recipient_account_id: AccountId,
         amount: u128,
     ) -> Result<(SendTxResponse, SharedSecretKey), ExecutionFailureKind> {
-        let (instruction_data, program, tx_pre_check) = token_program_preparation_transfer(amount);
+        let (instruction_data, program) = token_program_preparation_transfer(amount);
 
         self.0
             .send_privacy_preserving_tx(
@@ -181,7 +177,6 @@ impl Token<'_> {
                     PrivacyPreservingAccount::Public(recipient_account_id),
                 ],
                 &instruction_data,
-                tx_pre_check,
                 &program,
             )
             .await
@@ -200,7 +195,7 @@ impl Token<'_> {
         recipient_account_id: AccountId,
         amount: u128,
     ) -> Result<(SendTxResponse, SharedSecretKey), ExecutionFailureKind> {
-        let (instruction_data, program, tx_pre_check) = token_program_preparation_transfer(amount);
+        let (instruction_data, program) = token_program_preparation_transfer(amount);
 
         self.0
             .send_privacy_preserving_tx(
@@ -209,7 +204,6 @@ impl Token<'_> {
                     PrivacyPreservingAccount::PrivateOwned(recipient_account_id),
                 ],
                 &instruction_data,
-                tx_pre_check,
                 &program,
             )
             .await
@@ -229,7 +223,7 @@ impl Token<'_> {
         recipient_ipk: IncomingViewingPublicKey,
         amount: u128,
     ) -> Result<(SendTxResponse, SharedSecretKey), ExecutionFailureKind> {
-        let (instruction_data, program, tx_pre_check) = token_program_preparation_transfer(amount);
+        let (instruction_data, program) = token_program_preparation_transfer(amount);
 
         self.0
             .send_privacy_preserving_tx(
@@ -241,7 +235,6 @@ impl Token<'_> {
                     },
                 ],
                 &instruction_data,
-                tx_pre_check,
                 &program,
             )
             .await
@@ -255,13 +248,7 @@ impl Token<'_> {
     }
 }
 
-fn token_program_preparation_transfer(
+fn token_program_preparation_transfer(amount: u128) -> (InstructionData, Program) {
-    amount: u128,
-) -> (
-    InstructionData,
-    Program,
-    impl FnOnce(&[&Account]) -> Result<(), ExecutionFailureKind>,
-) {
     // Instruction must be: [0x01 || amount (little-endian 16 bytes) || 0x00 || 0x00 || 0x00 ||
     // 0x00 || 0x00 || 0x00].
     let mut instruction = [0; 23];
@@ -269,26 +256,20 @@ fn token_program_preparation_transfer(
     instruction[1..17].copy_from_slice(&amount.to_le_bytes());
     let instruction_data = Program::serialize_instruction(instruction).unwrap();
     let program = Program::token();
-    let tx_pre_check = |_: &[&Account]| Ok(());
 
-    (instruction_data, program, tx_pre_check)
+    (instruction_data, program)
 }
 
 fn token_program_preparation_definition(
     name: [u8; 6],
     total_supply: u128,
-) -> (
+) -> (InstructionData, Program) {
-    InstructionData,
-    Program,
-    impl FnOnce(&[&Account]) -> Result<(), ExecutionFailureKind>,
-) {
     // Instruction must be: [0x00 || total_supply (little-endian 16 bytes) || name (6 bytes)]
     let mut instruction = [0; 23];
     instruction[1..17].copy_from_slice(&total_supply.to_le_bytes());
     instruction[17..].copy_from_slice(&name);
     let instruction_data = Program::serialize_instruction(instruction).unwrap();
     let program = Program::token();
-    let tx_pre_check = |_: &[&Account]| Ok(());
 
-    (instruction_data, program, tx_pre_check)
+    (instruction_data, program)
 }

wallet/src/transaction_utils.rs

@@ -1,117 +0,0 @@
-use common::{error::ExecutionFailureKind, sequencer_client::json::SendTxResponse};
-use key_protocol::key_management::ephemeral_key_holder::EphemeralKeyHolder;
-use nssa::{
-    AccountId, PrivacyPreservingTransaction,
-    privacy_preserving_transaction::{circuit, message::Message, witness_set::WitnessSet},
-    program::Program,
-};
-use nssa_core::{
-    MembershipProof, NullifierPublicKey, NullifierSecretKey, SharedSecretKey,
-    account::AccountWithMetadata, encryption::IncomingViewingPublicKey,
-};
-
-use crate::{WalletCore, helperfunctions::produce_random_nonces};
-
-pub(crate) struct AccountPreparedData {
-    pub nsk: Option<NullifierSecretKey>,
-    pub npk: NullifierPublicKey,
-    pub ipk: IncomingViewingPublicKey,
-    pub auth_acc: AccountWithMetadata,
-    pub proof: Option<MembershipProof>,
-}
-
-impl WalletCore {
-    pub(crate) async fn private_acc_preparation(
-        &self,
-        account_id: AccountId,
-        is_authorized: bool,
-        needs_proof: bool,
-    ) -> Result<AccountPreparedData, ExecutionFailureKind> {
-        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_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 {
-            // TODO: Remove this unwrap, error types must be compatible
-            proof = self
-                .check_private_account_initialized(&account_id)
-                .await
-                .unwrap();
-        }
-
-        Ok(AccountPreparedData {
-            nsk,
-            npk: from_npk,
-            ipk: from_ipk,
-            auth_acc: sender_pre,
-            proof,
-        })
-    }
-
-    // TODO: Remove
-    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(),
-        )
-        .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],
-        ))
-    }
-}