add identifier to ciphertext and use it on sync mechanism

nssa/core/src/encryption/mod.rs

@@ -8,7 +8,7 @@ use serde::{Deserialize, Serialize};
 #[cfg(feature = "host")]
 pub use shared_key_derivation::{EphemeralPublicKey, EphemeralSecretKey, ViewingPublicKey};
 
-use crate::{Commitment, account::Account};
+use crate::{Commitment, Identifier, account::Account};
 #[cfg(feature = "host")]
 pub mod shared_key_derivation;
 
@@ -40,11 +40,14 @@ impl EncryptionScheme {
     #[must_use]
     pub fn encrypt(
         account: &Account,
+        identifier: Identifier,
         shared_secret: &SharedSecretKey,
         commitment: &Commitment,
         output_index: u32,
     ) -> Ciphertext {
-        let mut buffer = account.to_bytes();
+        // Plaintext: identifier (16 bytes, little-endian) || account bytes
+        let mut buffer = identifier.to_le_bytes().to_vec();
+        buffer.extend_from_slice(&account.to_bytes());
         Self::symmetric_transform(&mut buffer, shared_secret, commitment, output_index);
         Ciphertext(buffer)
     }
@@ -86,12 +89,17 @@ impl EncryptionScheme {
         shared_secret: &SharedSecretKey,
         commitment: &Commitment,
         output_index: u32,
-    ) -> Option<Account> {
+    ) -> Option<(Identifier, Account)> {
         use std::io::Cursor;
         let mut buffer = ciphertext.0.clone();
         Self::symmetric_transform(&mut buffer, shared_secret, commitment, output_index);
 
-        let mut cursor = Cursor::new(buffer.as_slice());
+        if buffer.len() < 16 {
+            return None;
+        }
+        let identifier = Identifier::from_le_bytes(buffer[..16].try_into().unwrap());
+
+        let mut cursor = Cursor::new(&buffer[16..]);
         Account::from_cursor(&mut cursor)
             .inspect_err(|err| {
                 println!(
@@ -104,5 +112,6 @@ impl EncryptionScheme {
                 );
             })
             .ok()
+            .map(|account| (identifier, account))
     }
 }

nssa/src/privacy_preserving_transaction/circuit.rs

@@ -259,7 +259,7 @@ mod tests {
         assert_eq!(output.new_nullifiers.len(), 1);
         assert_eq!(output.ciphertexts.len(), 1);
 
-        let recipient_post = EncryptionScheme::decrypt(
+        let (_identifier, recipient_post) = EncryptionScheme::decrypt(
             &output.ciphertexts[0],
             &shared_secret,
             &output.new_commitments[0],
@@ -357,7 +357,7 @@ mod tests {
         assert_eq!(output.new_nullifiers, expected_new_nullifiers);
         assert_eq!(output.ciphertexts.len(), 2);
 
-        let sender_post = EncryptionScheme::decrypt(
+        let (_identifier, sender_post) = EncryptionScheme::decrypt(
             &output.ciphertexts[0],
             &shared_secret_1,
             &expected_new_commitments[0],
@@ -366,7 +366,7 @@ mod tests {
         .unwrap();
         assert_eq!(sender_post, expected_private_account_1);
 
-        let recipient_post = EncryptionScheme::decrypt(
+        let (_identifier, recipient_post) = EncryptionScheme::decrypt(
             &output.ciphertexts[1],
             &shared_secret_2,
             &expected_new_commitments[1],

nssa/src/privacy_preserving_transaction/message.rs

@@ -186,7 +186,7 @@ pub mod tests {
         let esk = [3; 32];
         let shared_secret = SharedSecretKey::new(&esk, &vpk);
         let epk = EphemeralPublicKey::from_scalar(esk);
-        let ciphertext = EncryptionScheme::encrypt(&account, &shared_secret, &commitment, 2);
+        let ciphertext = EncryptionScheme::encrypt(&account, 0, &shared_secret, &commitment, 2);
         let encrypted_account_data =
             EncryptedAccountData::new(ciphertext.clone(), &npk, &vpk, epk.clone());
 

program_methods/guest/src/bin/privacy_preserving_circuit.rs

@@ -425,6 +425,7 @@ fn compute_circuit_output(
                 // Encrypt and push post state
                 let encrypted_account = EncryptionScheme::encrypt(
                     &post_with_updated_nonce,
+                    *identifier,
                     shared_secret,
                     &commitment_post,
                     output_index,

wallet/src/chain_storage.rs

@@ -175,6 +175,7 @@ impl WalletChainStore {
     pub fn insert_private_account_data(
         &mut self,
         account_id: nssa::AccountId,
+        identifier: nssa_core::Identifier,
         account: nssa_core::account::Account,
     ) {
         debug!("inserting at address {account_id}, this account {account:?}");
@@ -204,8 +205,6 @@ impl WalletChainStore {
         }
 
         // Otherwise update the private key tree
-        // Identifier is hardcoded to 0 until ciphertexts carry the identifier
-        let identifier: nssa_core::Identifier = 0;
 
         // Find the node by iterating all tree nodes for this account_id
         let chain_index = self

wallet/src/lib.rs

@@ -335,7 +335,7 @@ impl WalletCore {
                     let acc_ead = tx.message.encrypted_private_post_states[output_index].clone();
                     let acc_comm = tx.message.new_commitments[output_index].clone();
 
-                    let res_acc = nssa_core::EncryptionScheme::decrypt(
+                    let (identifier, res_acc) = nssa_core::EncryptionScheme::decrypt(
                         &acc_ead.ciphertext,
                         secret,
                         &acc_comm,
@@ -348,7 +348,7 @@ impl WalletCore {
                     println!("Received new acc {res_acc:#?}");
 
                     self.storage
-                        .insert_private_account_data(*acc_account_id, res_acc);
+                        .insert_private_account_data(*acc_account_id, identifier, res_acc);
                 }
                 AccDecodeData::Skip => {}
             }
@@ -484,35 +484,28 @@ impl WalletCore {
             .storage
             .user_data
             .default_user_private_accounts
-            .iter()
-            .map(|(acc_account_id, (key_chain, _))| (*acc_account_id, key_chain, None))
+            .values()
+            .map(|(key_chain, _)| (key_chain, None))
             .chain(self.storage.user_data.private_key_tree.key_map.iter().map(
-                |(chain_index, keys_node)| {
-                    // Use identifier=0 as the expected first account for this node.
-                    // The actual identifier will be confirmed once the account is synced.
-                    let account_id = nssa::AccountId::from((
-                        &keys_node.value.0.nullifier_public_key,
-                        0_u128,
-                    ));
-                    (account_id, &keys_node.value.0, chain_index.index())
-                },
+                |(chain_index, keys_node)| (&keys_node.value.0, chain_index.index()),
             ));
 
         let affected_accounts = private_account_key_chains
-            .flat_map(|(acc_account_id, key_chain, index)| {
+            .flat_map(|(key_chain, index)| {
                 let view_tag = EncryptedAccountData::compute_view_tag(
                     &key_chain.nullifier_public_key,
                     &key_chain.viewing_public_key,
                 );
+                let new_commitments = &tx.message.new_commitments;
 
                 tx.message()
                     .encrypted_private_post_states
                     .iter()
                     .enumerate()
                     .filter(move |(_, encrypted_data)| encrypted_data.view_tag == view_tag)
-                    .filter_map(|(ciph_id, encrypted_data)| {
+                    .filter_map(move |(ciph_id, encrypted_data)| {
                         let ciphertext = &encrypted_data.ciphertext;
-                        let commitment = &tx.message.new_commitments[ciph_id];
+                        let commitment = &new_commitments[ciph_id];
                         let shared_secret =
                             key_chain.calculate_shared_secret_receiver(&encrypted_data.epk, index);
 
@@ -524,18 +517,24 @@ impl WalletCore {
                                 .try_into()
                                 .expect("Ciphertext ID is expected to fit in u32"),
                         )
+                        .map(|(identifier, res_acc)| {
+                            let account_id = nssa::AccountId::from((
+                                &key_chain.nullifier_public_key,
+                                identifier,
+                            ));
+                            (account_id, identifier, res_acc)
+                        })
                     })
-                    .map(move |res_acc| (acc_account_id, res_acc))
                     .collect::<Vec<_>>()
             })
             .collect::<Vec<_>>();
 
-        for (affected_account_id, new_acc) in affected_accounts {
+        for (affected_account_id, identifier, new_acc) in affected_accounts {
             info!(
                 "Received new account for account_id {affected_account_id:#?} with account object {new_acc:#?}"
             );
             self.storage
-                .insert_private_account_data(affected_account_id, new_acc);
+                .insert_private_account_data(affected_account_id, identifier, new_acc);
         }
     }