wip

integration_tests/src/config.rs

@@ -144,6 +144,7 @@ impl InitialData {
                     account_id,
                     account: account.clone(),
                     key_chain: key_chain.clone(),
+                    identifier: 0,
                 }))
             }))
             .collect()

integration_tests/tests/auth_transfer/private.rs

@@ -170,7 +170,7 @@ async fn private_transfer_to_owned_account_using_claiming_path() -> Result<()> {
     };
 
     // Get the keys for the newly created account
-    let (to_keys, _) = ctx
+    let (to_keys, _, _) = ctx
         .wallet()
         .storage()
         .user_data
@@ -336,7 +336,7 @@ async fn private_transfer_to_owned_account_continuous_run_path() -> Result<()> {
     };
 
     // Get the newly created account's keys
-    let (to_keys, _) = ctx
+    let (to_keys, _, _) = ctx
         .wallet()
         .storage()
         .user_data

integration_tests/tests/keys_restoration.rs

@@ -59,7 +59,7 @@ async fn sync_private_account_with_non_zero_chain_index() -> Result<()> {
     };
 
     // Get the keys for the newly created account
-    let (to_keys, _) = ctx
+    let (to_keys, _, _) = ctx
         .wallet()
         .storage()
         .user_data

integration_tests/tests/token.rs

@@ -1131,7 +1131,7 @@ async fn token_claiming_path_with_private_accounts() -> Result<()> {
     };
 
     // Get keys for foreign mint (claiming path)
-    let (holder_keys, _) = ctx
+    let (holder_keys, _, _) = ctx
         .wallet()
         .storage()
         .user_data

integration_tests/tests/tps.rs

@@ -250,8 +250,8 @@ fn build_privacy_transaction() -> PrivacyPreservingTransaction {
         Program::serialize_instruction(balance_to_move).unwrap(),
         vec![1, 2],
         vec![
-            (sender_npk.clone(), sender_ss),
+            (sender_npk.clone(), 0, sender_ss),
-            (recipient_npk.clone(), recipient_ss),
+            (recipient_npk.clone(), 0, recipient_ss),
         ],
         vec![sender_nsk],
         vec![Some(proof)],

key_protocol/src/key_management/key_tree/keys_private.rs

@@ -1,5 +1,5 @@
 use k256::{Scalar, elliptic_curve::PrimeField as _};
-use nssa_core::{NullifierPublicKey, encryption::ViewingPublicKey};
+use nssa_core::{Identifier, NullifierPublicKey, encryption::ViewingPublicKey};
 use serde::{Deserialize, Serialize};
 
 use crate::key_management::{
@@ -10,8 +10,7 @@ use crate::key_management::{
 
 #[derive(Debug, Serialize, Deserialize, Clone)]
 pub struct ChildKeysPrivate {
-    // this should store value: (keychain, vec<identifier, account>) 
+    pub value: (KeyChain, nssa::Account, Identifier),
-    pub value: (KeyChain, nssa::Account),
     pub ccc: [u8; 32],
     /// Can be [`None`] if root.
     pub cci: Option<u32>,
@@ -48,6 +47,7 @@ impl KeyNode for ChildKeysPrivate {
                     },
                 },
                 nssa::Account::default(),
+                0,
             ),
             ccc,
             cci: None,
@@ -97,6 +97,7 @@ impl KeyNode for ChildKeysPrivate {
                     },
                 },
                 nssa::Account::default(),
+                0,
             ),
             ccc,
             cci: Some(cci),
@@ -120,7 +121,7 @@ impl KeyNode for ChildKeysPrivate {
     clippy::single_char_lifetime_names,
     reason = "TODO add meaningful name"
 )]
-impl<'a> From<&'a ChildKeysPrivate> for &'a (KeyChain, nssa::Account) {
+impl<'a> From<&'a ChildKeysPrivate> for &'a (KeyChain, nssa::Account, Identifier) {
     fn from(value: &'a ChildKeysPrivate) -> Self {
         &value.value
     }
@@ -130,7 +131,7 @@ impl<'a> From<&'a ChildKeysPrivate> for &'a (KeyChain, nssa::Account) {
     clippy::single_char_lifetime_names,
     reason = "TODO add meaningful name"
 )]
-impl<'a> From<&'a mut ChildKeysPrivate> for &'a mut (KeyChain, nssa::Account) {
+impl<'a> From<&'a mut ChildKeysPrivate> for &'a mut (KeyChain, nssa::Account, Identifier) {
     fn from(value: &'a mut ChildKeysPrivate) -> Self {
         &mut value.value
     }

key_protocol/src/key_protocol_core/mod.rs

@@ -4,6 +4,8 @@ use anyhow::Result;
 use k256::AffinePoint;
 use serde::{Deserialize, Serialize};
 
+use nssa_core::Identifier;
+
 use crate::key_management::{
     KeyChain,
     key_tree::{KeyTreePrivate, KeyTreePublic, chain_index::ChainIndex},
@@ -17,9 +19,8 @@ pub struct NSSAUserData {
     /// Default public accounts.
     pub default_pub_account_signing_keys: BTreeMap<nssa::AccountId, nssa::PrivateKey>,
     /// Default private accounts.
-    // TODO: this should store (keychain, account, identifier) in the values of the map
     pub default_user_private_accounts:
-        BTreeMap<nssa::AccountId, (KeyChain, nssa_core::account::Account)>,
+        BTreeMap<nssa::AccountId, (KeyChain, nssa_core::account::Account, Identifier)>,
     /// Tree of public keys.
     pub public_key_tree: KeyTreePublic,
     /// Tree of private keys.
@@ -43,12 +44,15 @@ impl NSSAUserData {
     }
 
     fn valid_private_key_transaction_pairing_check(
-        accounts_keys_map: &BTreeMap<nssa::AccountId, (KeyChain, nssa_core::account::Account)>,
+        accounts_keys_map: &BTreeMap<
+            nssa::AccountId,
+            (KeyChain, nssa_core::account::Account, Identifier),
+        >,
     ) -> bool {
         let mut check_res = true;
-        for (account_id, (key, _)) in accounts_keys_map {
+        for (account_id, (key, _, identifier)) in accounts_keys_map {
-            // TODO: Generalize to other identifiers
+            let expected_account_id =
-            let expected_account_id = nssa::AccountId::from((&key.nullifier_public_key, 0));
+                nssa::AccountId::from((&key.nullifier_public_key, *identifier));
             if expected_account_id != *account_id {
                 println!("{expected_account_id}, {account_id}");
                 check_res = false;
@@ -61,7 +65,7 @@ impl NSSAUserData {
         default_accounts_keys: BTreeMap<nssa::AccountId, nssa::PrivateKey>,
         default_accounts_key_chains: BTreeMap<
             nssa::AccountId,
-            (KeyChain, nssa_core::account::Account),
+            (KeyChain, nssa_core::account::Account, Identifier),
         >,
         public_key_tree: KeyTreePublic,
         private_key_tree: KeyTreePrivate,
@@ -140,7 +144,7 @@ impl NSSAUserData {
     pub fn get_private_account(
         &self,
         account_id: nssa::AccountId,
-    ) -> Option<&(KeyChain, nssa_core::account::Account)> {
+    ) -> Option<&(KeyChain, nssa_core::account::Account, Identifier)> {
         self.default_user_private_accounts
             .get(&account_id)
             .or_else(|| self.private_key_tree.get_node(account_id).map(Into::into))
@@ -150,7 +154,7 @@ impl NSSAUserData {
     pub fn get_private_account_mut(
         &mut self,
         account_id: &nssa::AccountId,
-    ) -> Option<&mut (KeyChain, nssa_core::account::Account)> {
+    ) -> Option<&mut (KeyChain, nssa_core::account::Account, Identifier)> {
         // First seek in defaults
         if let Some(key) = self.default_user_private_accounts.get_mut(account_id) {
             Some(key)

testnet_initial_state/src/lib.rs

@@ -98,6 +98,7 @@ pub struct PrivateAccountPrivateInitialData {
     pub account_id: nssa::AccountId,
     pub account: nssa_core::account::Account,
     pub key_chain: KeyChain,
+    pub identifier: nssa_core::Identifier,
 }
 
 #[must_use]
@@ -150,6 +151,7 @@ pub fn initial_priv_accounts_private_keys() -> Vec<PrivateAccountPrivateInitialD
                 nonce: 0.into(),
             },
             key_chain: key_chain_1,
+            identifier: 0,
         },
         PrivateAccountPrivateInitialData {
             account_id: AccountId::from((&key_chain_2.nullifier_public_key, 0)),
@@ -160,6 +162,7 @@ pub fn initial_priv_accounts_private_keys() -> Vec<PrivateAccountPrivateInitialD
                 nonce: 0.into(),
             },
             key_chain: key_chain_2,
+            identifier: 0,
         },
     ]
 }

wallet-ffi/src/keys.rs

@@ -116,7 +116,8 @@ pub unsafe extern "C" fn wallet_ffi_get_private_account_keys(
 
     let account_id = AccountId::new(unsafe { (*account_id).data });
 
-    let Some((key_chain, _account)) = wallet.storage().user_data.get_private_account(account_id)
+    let Some((key_chain, _account, _identifier)) =
+        wallet.storage().user_data.get_private_account(account_id)
     else {
         print_error("Private account not found in wallet");
         return WalletFfiError::AccountNotFound;

wallet/src/chain_storage.rs

@@ -78,7 +78,7 @@ impl WalletChainStore {
                     }
                     InitialAccountData::Private(data) => {
                         private_init_acc_map
-                            .insert(data.account_id, (data.key_chain, data.account));
+                            .insert(data.account_id, (data.key_chain, data.account, data.identifier));
                     }
                 },
             }
@@ -117,7 +117,8 @@ impl WalletChainStore {
                     // startup. Fix this when program id can be fetched
                     // from the node and queried from the wallet.
                     account.program_owner = Program::authenticated_transfer_program().id();
-                    private_init_acc_map.insert(data.account_id, (data.key_chain, account));
+                    private_init_acc_map
+                        .insert(data.account_id, (data.key_chain, account, data.identifier));
                 }
             }
         }

wallet/src/cli/account.rs

@@ -156,7 +156,7 @@ impl WalletSubcommand for NewSubcommand {
                         .insert(account_id.to_string(), Label::new(label));
                 }
 
-                let (key, _) = wallet_core
+                let (key, _, _) = wallet_core
                     .storage
                     .user_data
                     .get_private_account(account_id)
@@ -229,7 +229,7 @@ impl WalletSubcommand for AccountSubcommand {
                             println!("pk {}", hex::encode(public_key.value()));
                         }
                         AccountPrivacyKind::Private => {
-                            let (key, _) = wallet_core
+                            let (key, _, _) = wallet_core
                                 .storage
                                 .user_data
                                 .get_private_account(account_id)

wallet/src/helperfunctions.rs

@@ -188,12 +188,13 @@ pub fn produce_data_for_storage(
         );
     }
 
-    for (account_id, (key_chain, account)) in &user_data.default_user_private_accounts {
+    for (account_id, (key_chain, account, identifier)) in &user_data.default_user_private_accounts {
         vec_for_storage.push(
             InitialAccountData::Private(Box::new(PrivateAccountPrivateInitialData {
                 account_id: *account_id,
                 account: account.clone(),
                 key_chain: key_chain.clone(),
+                identifier: *identifier,
             }))
             .into(),
         );

wallet/src/lib.rs

@@ -300,7 +300,8 @@ impl WalletCore {
 
     #[must_use]
     pub fn get_private_account_commitment(&self, account_id: AccountId) -> Option<Commitment> {
-        let (_keys, account) = self.storage.user_data.get_private_account(account_id)?;
+        let (_keys, account, _identifier) =
+            self.storage.user_data.get_private_account(account_id)?;
         Some(Commitment::new(&account_id, account))
     }
 
@@ -393,7 +394,7 @@ impl WalletCore {
             acc_manager.visibility_mask().to_vec(),
             private_account_keys
                 .iter()
-                .map(|keys| (keys.npk.clone(), 0, keys.ssk))
+                .map(|keys| (keys.npk.clone(), keys.identifier, keys.ssk))
                 .collect::<Vec<_>>(),
             acc_manager.private_account_auth(),
             acc_manager.private_account_membership_proofs(),
@@ -484,7 +485,7 @@ impl WalletCore {
             .user_data
             .default_user_private_accounts
             .iter()
-            .map(|(acc_account_id, (key_chain, _))| (*acc_account_id, key_chain, None))
+            .map(|(acc_account_id, (key_chain, _, _))| (*acc_account_id, key_chain, None))
             .chain(self.storage.user_data.private_key_tree.key_map.iter().map(
                 |(chain_index, keys_node)| {
                     (

wallet/src/privacy_preserving_tx.rs

@@ -2,7 +2,7 @@ use anyhow::Result;
 use key_protocol::key_management::ephemeral_key_holder::EphemeralKeyHolder;
 use nssa::{AccountId, PrivateKey};
 use nssa_core::{
-    MembershipProof, NullifierPublicKey, NullifierSecretKey, SharedSecretKey,
+    Identifier, MembershipProof, NullifierPublicKey, NullifierSecretKey, SharedSecretKey,
     account::{AccountWithMetadata, Nonce},
     encryption::{EphemeralPublicKey, ViewingPublicKey},
 };
@@ -36,6 +36,7 @@ impl PrivacyPreservingAccount {
 
 pub struct PrivateAccountKeys {
     pub npk: NullifierPublicKey,
+    pub identifier: Identifier,
     pub ssk: SharedSecretKey,
     pub vpk: ViewingPublicKey,
     pub epk: EphemeralPublicKey,
@@ -87,6 +88,7 @@ impl AccountManager {
                     let pre = AccountPreparedData {
                         nsk: None,
                         npk,
+                        identifier: 0,
                         vpk,
                         pre_state: auth_acc,
                         proof: None,
@@ -139,6 +141,7 @@ impl AccountManager {
 
                     Some(PrivateAccountKeys {
                         npk: pre.npk.clone(),
+                        identifier: pre.identifier,
                         ssk: eph_holder.calculate_shared_secret_sender(&pre.vpk),
                         vpk: pre.vpk.clone(),
                         epk: eph_holder.generate_ephemeral_public_key(),
@@ -193,6 +196,7 @@ impl AccountManager {
 struct AccountPreparedData {
     nsk: Option<NullifierSecretKey>,
     npk: NullifierPublicKey,
+    identifier: Identifier,
     vpk: ViewingPublicKey,
     pre_state: AccountWithMetadata,
     proof: Option<MembershipProof>,
@@ -202,7 +206,7 @@ async fn private_acc_preparation(
     wallet: &WalletCore,
     account_id: AccountId,
 ) -> Result<AccountPreparedData, ExecutionFailureKind> {
-    let Some((from_keys, from_acc)) = wallet
+    let Some((from_keys, from_acc, from_identifier)) = wallet
         .storage
         .user_data
         .get_private_account(account_id)
@@ -224,11 +228,13 @@ async fn private_acc_preparation(
 
     // TODO: Technically we could allow unauthorized owned accounts, but currently we don't have
     // support from that in the wallet.
-    let sender_pre = AccountWithMetadata::new(from_acc.clone(), true, (&from_npk, 0));
+    let sender_pre =
+        AccountWithMetadata::new(from_acc.clone(), true, (&from_npk, from_identifier));
 
     Ok(AccountPreparedData {
         nsk: Some(nsk),
         npk: from_npk,
+        identifier: from_identifier,
         vpk: from_vpk,
         pre_state: sender_pre,
         proof,