Merge a3593e36b7a921462228d2be36eb4febea12b619 into 1b4d8fbcf4caa1448473fd99b4ac9d386abf3bc3

lee/key_protocol/src/key_management/group_key_holder.rs

@@ -1,6 +1,6 @@
 use aes_gcm::{Aes256Gcm, KeyInit as _, aead::Aead as _};
 use lee_core::{
-    SharedSecretKey,
+    Identifier, SharedSecretKey,
     encryption::{EphemeralPublicKey, ML_KEM_768_CIPHERTEXT_LEN, ViewingPublicKey},
     program::{PdaSeed, ProgramId},
 };
@@ -146,6 +146,23 @@ impl GroupKeyHolder {
         SecretSpendingKey(hasher.finalize_fixed().into()).produce_private_key_holder(None)
     }
 
+    /// Derive keys for a shared regular account from its `identifier`.
+    ///
+    /// Computes the derivation seed via the `SharedAccountTag` domain separator, then delegates
+    /// to [`Self::derive_keys_for_shared_account`].
+    #[must_use]
+    pub fn derive_regular_shared_account_keys_from_identifier(
+        &self,
+        identifier: Identifier,
+    ) -> PrivateKeyHolder {
+        const PREFIX: &[u8; 32] = b"/LEE/v0.3/SharedAccountTag/\x00\x00\x00\x00\x00";
+        let mut hasher = sha2::Sha256::new();
+        hasher.update(PREFIX);
+        hasher.update(identifier.to_le_bytes());
+        let derivation_seed: [u8; 32] = hasher.finalize().into();
+        self.derive_keys_for_shared_account(&derivation_seed)
+    }
+
     /// Encrypts this holder's GMS under the recipient's [`SealingPublicKey`].
     ///
     /// Uses ML-KEM-768 encapsulation to derive a shared secret, then AES-256-GCM to encrypt

lee/state_machine/core/src/program/mod.rs

@@ -324,7 +324,7 @@ impl AccountPostState {
 pub type BlockValidityWindow = ValidityWindow<BlockId>;
 pub type TimestampValidityWindow = ValidityWindow<Timestamp>;
 
-#[derive(Clone, Copy, Serialize, Deserialize)]
+#[derive(Clone, Copy, Default, Serialize, Deserialize)]
 #[cfg_attr(
     any(feature = "host", test),
     derive(Debug, PartialEq, Eq, BorshSerialize, BorshDeserialize)

lee/state_machine/src/privacy_preserving_transaction/message.rs

@@ -11,7 +11,7 @@ use crate::{AccountId, error::LeeError};
 
 const PREFIX: &[u8; 32] = b"/LEE/v0.3/Message/Privacy/\x00\x00\x00\x00\x00\x00";
 
-#[derive(Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)]
+#[derive(Clone, Default, PartialEq, Eq, BorshSerialize, BorshDeserialize)]
 pub struct Message {
     pub public_account_ids: Vec<AccountId>,
     pub nonces: Vec<Nonce>,
@@ -84,6 +84,20 @@ impl Message {
 
         Sha256::digest(bytes).into()
     }
+
+    /// Ensure that the commitments, nullifiers, and ciphertexts agree.
+    pub fn validate_note_lengths(&self) -> Result<usize, LeeError> {
+        let count = self.new_nullifiers.len();
+        if self.new_commitments.len() != count || self.encrypted_private_post_states.len() != count
+        {
+            return Err(LeeError::InvalidInput(format!(
+                "Note vectors disagree in length with {count} nullifiers, {} commitments, and {} ciphertexts",
+                self.new_commitments.len(),
+                self.encrypted_private_post_states.len(),
+            )));
+        }
+        Ok(count)
+    }
 }
 
 #[cfg(test)]
@@ -141,6 +155,20 @@ pub mod tests {
         }
     }
 
+    #[test]
+    fn validate_note_lengths_accepts_matching_and_rejects_mismatched() {
+        assert_eq!(Message::default().validate_note_lengths().unwrap(), 0);
+
+        let mismatched = Message {
+            new_commitments: vec![Commitment::new(
+                &AccountId::new([0; 32]),
+                &Account::default(),
+            )],
+            ..Default::default()
+        };
+        assert!(mismatched.validate_note_lengths().is_err());
+    }
+
     #[test]
     fn hash_privacy_pinned() {
         let msg = Message {

lez/wallet/src/lib.rs

@@ -7,7 +7,7 @@
     reason = "Most of the shadows come from args parsing which is ok"
 )]
 
-use std::path::PathBuf;
+use std::{collections::HashSet, path::PathBuf};
 
 pub use account_manager::AccountIdentity;
 use anyhow::{Context as _, Result};
@@ -18,7 +18,8 @@ use key_protocol::key_management::key_tree::chain_index::ChainIndex;
 use lee::{
     Account, AccountId, PrivacyPreservingTransaction, ProgramId,
     privacy_preserving_transaction::{
-        circuit::ProgramWithDependencies, message::EncryptedAccountData,
+        circuit::ProgramWithDependencies,
+        message::{EncryptedAccountData, Message},
     },
 };
 use lee_core::{
@@ -34,7 +35,7 @@ use crate::{
     account::{AccountIdWithPrivacy, Label},
     config::WalletConfigOverrides,
     poller::TxPoller,
-    storage::key_chain::SharedAccountEntry,
+    storage::key_chain::{NullifierIndex, SharedAccountEntry},
 };
 
 pub mod account;
@@ -298,35 +299,26 @@ impl WalletCore {
             .storage
             .key_chain()
             .shared_private_account(account_id)?;
-        let holder = self
-            .storage
-            .key_chain()
-            .group_key_holder(&entry.group_label)?;
+        let keys = self.storage.key_chain().derive_shared_account_keys(entry)?;
+        let nsk = keys.nullifier_secret_key;
+        let npk = keys.generate_nullifier_public_key();
+        let vpk = keys.generate_viewing_public_key();
+        let identifier = entry.identifier;
 
-        if let (Some(pda_seed), Some(program_id)) = (entry.pda_seed, entry.authority_program_id) {
-            let keys = holder.derive_keys_for_pda(&program_id, &pda_seed);
+        if entry.pda_seed.is_some() {
             Some(AccountIdentity::PrivatePdaShared {
                 account_id,
-                nsk: keys.nullifier_secret_key,
-                npk: keys.generate_nullifier_public_key(),
-                vpk: keys.generate_viewing_public_key(),
-                identifier: entry.identifier,
+                nsk,
+                npk,
+                vpk,
+                identifier,
             })
         } else {
-            let derivation_seed = {
-                use sha2::Digest as _;
-                let mut hasher = sha2::Sha256::new();
-                hasher.update(b"/LEE/v0.3/SharedAccountTag/\x00\x00\x00\x00\x00");
-                hasher.update(entry.identifier.to_le_bytes());
-                let result: [u8; 32] = hasher.finalize().into();
-                result
-            };
-            let keys = holder.derive_keys_for_shared_account(&derivation_seed);
             Some(AccountIdentity::PrivateShared {
-                nsk: keys.nullifier_secret_key,
-                npk: keys.generate_nullifier_public_key(),
-                vpk: keys.generate_viewing_public_key(),
-                identifier: entry.identifier,
+                nsk,
+                npk,
+                vpk,
+                identifier,
             })
         }
     }
@@ -401,14 +393,6 @@ impl WalletCore {
         group_name: Label,
     ) -> Result<SharedAccountInfo> {
         let identifier: lee_core::Identifier = rand::random();
-        let derivation_seed = {
-            use sha2::Digest as _;
-            let mut hasher = sha2::Sha256::new();
-            hasher.update(b"/LEE/v0.3/SharedAccountTag/\x00\x00\x00\x00\x00");
-            hasher.update(identifier.to_le_bytes());
-            let result: [u8; 32] = hasher.finalize().into();
-            result
-        };
 
         let holder = self
             .storage
@@ -416,7 +400,7 @@ impl WalletCore {
             .group_key_holder(&group_name)
             .context(format!("Group '{group_name}' not found"))?;
 
-        let keys = holder.derive_keys_for_shared_account(&derivation_seed);
+        let keys = holder.derive_regular_shared_account_keys_from_identifier(identifier);
         let npk = keys.generate_nullifier_public_key();
         let vpk = keys.generate_viewing_public_key();
         let account_id = AccountId::from((&npk, &vpk, identifier));
@@ -522,6 +506,12 @@ impl WalletCore {
         tx: &lee::privacy_preserving_transaction::PrivacyPreservingTransaction,
         acc_decode_mask: &[AccDecodeData],
     ) -> Result<()> {
+        let note_count = tx.message.validate_note_lengths()?;
+        anyhow::ensure!(
+            note_count >= acc_decode_mask.len(),
+            "Decode mask has {} entries but the transaction has {note_count} notes",
+            acc_decode_mask.len(),
+        );
         for (output_index, acc_decode_data) in acc_decode_mask.iter().enumerate() {
             match acc_decode_data {
                 AccDecodeData::Decode(secret, acc_account_id) => {
@@ -739,10 +729,21 @@ impl WalletCore {
         let mut blocks =
             std::pin::pin!(poller.poll_block_range(last_synced_block.saturating_add(1)..=block_id));
 
+        // Get the latest nullifiers for all owned accounts.
+        let mut index = self.storage.key_chain().build_latest_nullifier_index();
         let bar = indicatif::ProgressBar::new(num_of_blocks);
         while let Some(block) = blocks.try_next().await? {
             for tx in block.body.transactions {
-                self.sync_private_accounts_with_tx(tx);
+                let LeeTransaction::PrivacyPreserving(pp_tx) = &tx else {
+                    continue;
+                };
+                pp_tx.message.validate_note_lengths()?;
+                // Eagerly decrypt note updates using expected nullifiers.
+                let handled = self
+                    .storage
+                    .key_chain_mut()
+                    .sync_updates_via_nullifiers(&pp_tx.message, &mut index);
+                self.sync_private_accounts_with_message(&pp_tx.message, &mut index, &handled);
             }
 
             self.storage.set_last_synced_block(block.header.block_id);
@@ -759,11 +760,12 @@ impl WalletCore {
         Ok(())
     }
 
-    fn sync_private_accounts_with_tx(&mut self, tx: LeeTransaction) {
-        let LeeTransaction::PrivacyPreserving(tx) = tx else {
-            return;
-        };
-
+    fn sync_private_accounts_with_message(
+        &mut self,
+        message: &Message,
+        index: &mut NullifierIndex,
+        handled: &HashSet<usize>,
+    ) {
         let affected_accounts = self
             .storage
             .key_chain()
@@ -773,13 +775,18 @@ impl WalletCore {
                     &key_chain.nullifier_public_key,
                     &key_chain.viewing_public_key,
                 );
-                let new_commitments = &tx.message.new_commitments;
+                let new_commitments = &message.new_commitments;
 
-                tx.message()
+                message
                     .encrypted_private_post_states
                     .iter()
                     .enumerate()
-                    .filter(move |(_, encrypted_data)| encrypted_data.view_tag == view_tag)
+                    .filter(move |(ciph_id, encrypted_data)| {
+                        // If we have not decrypted the update using the nullifiers,
+                        // the note may be an initialized one, for which we should
+                        // scan.
+                        !handled.contains(ciph_id) && encrypted_data.view_tag == view_tag
+                    })
                     .filter_map(move |(ciph_id, encrypted_data)| {
                         let ciphertext = &encrypted_data.ciphertext;
                         let commitment = &new_commitments[ciph_id];
@@ -801,17 +808,21 @@ impl WalletCore {
                                 &key_chain.viewing_public_key,
                                 &kind,
                             );
-                            (account_id, kind, res_acc)
+                            let nsk = key_chain.private_key_holder.nullifier_secret_key;
+                            (account_id, kind, res_acc, nsk)
                         })
                     })
                     .collect::<Vec<_>>()
             })
             .collect::<Vec<_>>();
 
-        for (affected_account_id, kind, new_acc) in affected_accounts {
+        for (affected_account_id, kind, new_acc, nsk) in affected_accounts {
             info!(
                 "Received new account for account_id {affected_account_id:#?} with account object {new_acc:#?}"
             );
+            // Await the account's next update by its nullifier, so later updates
+            // to it are caught without tag matching.
+            index.track(affected_account_id, &new_acc, &nsk);
             self.storage
                 .key_chain_mut()
                 .insert_private_account(affected_account_id, kind, new_acc)
@@ -819,54 +830,37 @@ impl WalletCore {
         }
 
         // Scan for updates to shared accounts (GMS-derived).
-        self.sync_shared_private_accounts_with_tx(&tx);
+        self.sync_shared_private_accounts_with_tx(message, index, handled);
     }
 
-    fn sync_shared_private_accounts_with_tx(&mut self, tx: &PrivacyPreservingTransaction) {
+    fn sync_shared_private_accounts_with_tx(
+        &mut self,
+        message: &Message,
+        index: &mut NullifierIndex,
+        handled: &HashSet<usize>,
+    ) {
         let shared_keys: Vec<_> = self
             .storage
             .key_chain()
             .shared_private_accounts_iter()
             .filter_map(|(&account_id, entry)| {
-                let holder = self
-                    .storage
-                    .key_chain()
-                    .group_key_holder(&entry.group_label)?;
-
-                let keys = match (&entry.pda_seed, &entry.authority_program_id) {
-                    (Some(pda_seed), Some(program_id)) => {
-                        holder.derive_keys_for_pda(program_id, pda_seed)
-                    }
-                    (Some(_), None) => return None, // PDA without program_id, skip
-                    _ => {
-                        let derivation_seed = {
-                            use sha2::Digest as _;
-                            let mut hasher = sha2::Sha256::new();
-                            hasher.update(b"/LEE/v0.3/SharedAccountTag/\x00\x00\x00\x00\x00");
-                            hasher.update(entry.identifier.to_le_bytes());
-                            let result: [u8; 32] = hasher.finalize().into();
-                            result
-                        };
-                        holder.derive_keys_for_shared_account(&derivation_seed)
-                    }
-                };
+                let keys = self.storage.key_chain().derive_shared_account_keys(entry)?;
                 let npk = keys.generate_nullifier_public_key();
                 let vpk = keys.generate_viewing_public_key();
+                let nsk = keys.nullifier_secret_key;
                 let vsk = keys.viewing_secret_key;
-                Some((account_id, npk, vpk, vsk))
+                Some((account_id, npk, vpk, vsk, nsk))
             })
             .collect();
 
-        for (account_id, npk, vpk, vsk) in shared_keys {
+        for (account_id, npk, vpk, vsk, nsk) in shared_keys {
             let view_tag = EncryptedAccountData::compute_view_tag(&npk, &vpk);
 
-            for (ciph_id, encrypted_data) in tx
-                .message()
-                .encrypted_private_post_states
-                .iter()
-                .enumerate()
+            for (ciph_id, encrypted_data) in
+                message.encrypted_private_post_states.iter().enumerate()
             {
-                if encrypted_data.view_tag != view_tag {
+                // If already decrypted or the tag does not match, skip.
+                if handled.contains(&ciph_id) || encrypted_data.view_tag != view_tag {
                     continue;
                 }
 
@@ -875,7 +869,7 @@ impl WalletCore {
                 else {
                     continue;
                 };
-                let commitment = &tx.message.new_commitments[ciph_id];
+                let commitment = &message.new_commitments[ciph_id];
 
                 if let Some((_kind, new_acc)) = lee_core::EncryptionScheme::decrypt(
                     &encrypted_data.ciphertext,
@@ -886,6 +880,7 @@ impl WalletCore {
                         .expect("Ciphertext ID is expected to fit in u32"),
                 ) {
                     info!("Synced shared account {account_id:#?} with new state {new_acc:#?}");
+                    index.track(account_id, &new_acc, &nsk);
                     self.storage
                         .key_chain_mut()
                         .update_shared_private_account_state(&account_id, new_acc);

lez/wallet/src/storage/key_chain.rs

@@ -1,15 +1,18 @@
 use core::panic;
-use std::collections::{BTreeMap, btree_map::Entry};
+use std::collections::{BTreeMap, HashMap, HashSet, btree_map::Entry};
 
 use anyhow::{Context as _, Result, anyhow};
 use key_protocol::key_management::{
     KeyChain,
     group_key_holder::GroupKeyHolder,
     key_tree::{KeyTreePrivate, KeyTreePublic, chain_index::ChainIndex, traits::KeyTreeNode as _},
-    secret_holders::{SeedHolder, ViewingSecretKey},
+    secret_holders::{PrivateKeyHolder, SeedHolder, ViewingSecretKey},
+};
+use lee::{Account, AccountId, privacy_preserving_transaction::message::Message};
+use lee_core::{
+    Commitment, EncryptionScheme, Identifier, Nullifier, NullifierSecretKey, PrivateAccountKind,
+    SharedSecretKey,
 };
-use lee::{Account, AccountId};
-use lee_core::{Identifier, PrivateAccountKind};
 use log::{debug, warn};
 use serde::{Deserialize, Serialize};
 use testnet_initial_state::{PrivateAccountPrivateInitialData, PublicAccountPrivateInitialData};
@@ -59,6 +62,41 @@ pub struct SharedAccountEntry {
     pub account: Account,
 }
 
+/// Maps each owned or shared private account to the nullifier its next update will publish,
+/// so sync can spot updates by nullifier rather than view tag.
+#[derive(Default)]
+pub struct NullifierIndex(HashMap<Nullifier, AccountId>);
+
+impl NullifierIndex {
+    fn next_update_nullifier(
+        account_id: AccountId,
+        account: &Account,
+        nsk: &NullifierSecretKey,
+    ) -> Nullifier {
+        Nullifier::for_account_update(&Commitment::new(&account_id, account), nsk)
+    }
+
+    /// Returns the account whose next update would publish `nullifier`.
+    #[must_use]
+    pub fn account_for(&self, nullifier: &Nullifier) -> Option<AccountId> {
+        self.0.get(nullifier).copied()
+    }
+
+    /// Indexes `account_id` by the nullifier its next update will publish.
+    pub fn track(&mut self, account_id: AccountId, account: &Account, nsk: &NullifierSecretKey) {
+        self.0.insert(
+            Self::next_update_nullifier(account_id, account, nsk),
+            account_id,
+        );
+    }
+
+    /// Replaces a spent nullifier with the account's `next` one.
+    pub fn update(&mut self, spent: &Nullifier, next: Nullifier, account_id: AccountId) {
+        self.0.remove(spent);
+        self.0.insert(next, account_id);
+    }
+}
+
 #[derive(Clone, Debug)]
 #[cfg_attr(test, derive(PartialEq, Eq))]
 pub struct UserKeyChain {
@@ -210,6 +248,19 @@ impl UserKeyChain {
     /// those.
     #[must_use]
     pub fn private_account(&self, account_id: AccountId) -> Option<FoundPrivateAccount<'_>> {
+        self.private_accounts().find_map(|found| {
+            let expected_id = AccountId::for_private_account(
+                &found.key_chain.nullifier_public_key,
+                &found.key_chain.viewing_public_key,
+                found.kind,
+            );
+            (expected_id == account_id).then_some(found)
+        })
+    }
+
+    /// Iterates every owned private account (imported and generated), one
+    /// [`FoundPrivateAccount`] per identity. Excludes shared accounts.
+    pub fn private_accounts(&self) -> impl Iterator<Item = FoundPrivateAccount<'_>> {
         self.imported_private_accounts
             .iter()
             .flat_map(|(key, data)| {
@@ -238,14 +289,6 @@ impl UserKeyChain {
                             })
                     }),
             )
-            .find_map(|found| {
-                let expected_id = AccountId::for_private_account(
-                    &found.key_chain.nullifier_public_key,
-                    &found.key_chain.viewing_public_key,
-                    found.kind,
-                );
-                (expected_id == account_id).then_some(found)
-            })
     }
 
     #[must_use]
@@ -286,6 +329,124 @@ impl UserKeyChain {
             )
     }
 
+    /// Re-derives the [`PrivateKeyHolder`] for a shared account `entry`, dispatching on PDA vs
+    /// regular. `None` if the group key holder is absent or a PDA entry lacks its program id.
+    #[must_use]
+    pub fn derive_shared_account_keys(
+        &self,
+        entry: &SharedAccountEntry,
+    ) -> Option<PrivateKeyHolder> {
+        let holder = self.group_key_holder(&entry.group_label)?;
+        Some(match (&entry.pda_seed, &entry.authority_program_id) {
+            (Some(pda_seed), Some(program_id)) => holder.derive_keys_for_pda(program_id, pda_seed),
+            (Some(_), None) => return None,
+            _ => holder.derive_regular_shared_account_keys_from_identifier(entry.identifier),
+        })
+    }
+
+    /// Maps each owned and shared account's current-state update nullifier to its `account_id`,
+    /// so co-owner updates are found during sync by nullifier rather than view tag.
+    #[must_use]
+    pub fn build_latest_nullifier_index(&self) -> NullifierIndex {
+        let mut index = NullifierIndex::default();
+
+        // For each (regular) found account the user owns, compute its nullifier and put
+        // into the map. This is the next nullifier it will look for.
+        for found in self.private_accounts() {
+            let account_id = AccountId::for_private_account(
+                &found.key_chain.nullifier_public_key,
+                &found.key_chain.viewing_public_key,
+                found.kind,
+            );
+            let nsk = found.key_chain.private_key_holder.nullifier_secret_key;
+            index.track(account_id, found.account, &nsk);
+        }
+
+        // Same for the shared accounts.
+        for (&account_id, entry) in self.shared_private_accounts_iter() {
+            let Some(keys) = self.derive_shared_account_keys(entry) else {
+                continue;
+            };
+            let nsk = keys.nullifier_secret_key;
+            index.track(account_id, &entry.account, &nsk);
+        }
+
+        index
+    }
+
+    /// Applies every watched nullifier the `message` publishes: decrypts the position-aligned
+    /// note, stores the new state, and rolls the index to the account's next nullifier. Returns
+    /// the output slots handled, so the view-tag pass can skip them.
+    pub fn sync_updates_via_nullifiers(
+        &mut self,
+        message: &Message,
+        index: &mut NullifierIndex,
+    ) -> HashSet<usize> {
+        let mut handled = HashSet::new();
+        for (i, (old_nullifier, _)) in message.new_nullifiers.iter().enumerate() {
+            // Get the nullifier information if awaiting the nullifier.
+            let Some(account_id) = index.account_for(old_nullifier) else {
+                continue;
+            };
+            // Try decrypting the commitment connected to the nullifier and get the next
+            // nullifier to await.
+            if let Some(new_nullifier) = self.apply_nullifier_update(account_id, message, i) {
+                // Update the index to await for the new state of the account, i.e.
+                // the new nullifier.
+                index.update(old_nullifier, new_nullifier, account_id);
+                // Record that this nullifier's position can be skipped for scanning.
+                handled.insert(i);
+            }
+        }
+        handled
+    }
+
+    /// Decrypts the note at slot `i` for `account_id`, stores the new state, and returns the
+    /// account's next update nullifier. `None` if keys or decryption fail.
+    fn apply_nullifier_update(
+        &mut self,
+        account_id: AccountId,
+        message: &Message,
+        i: usize,
+    ) -> Option<Nullifier> {
+        let encrypted = &message.encrypted_private_post_states[i];
+        let commitment = &message.new_commitments[i];
+        let ciph_id = u32::try_from(i).ok()?;
+
+        let (nsk, secret, is_shared) = if let Some(entry) = self.shared_private_account(account_id)
+        {
+            let keys = self.derive_shared_account_keys(entry)?;
+            let secret = SharedSecretKey::decapsulate(
+                &encrypted.epk,
+                &keys.viewing_secret_key.d,
+                &keys.viewing_secret_key.z,
+            )?;
+            (keys.nullifier_secret_key, secret, true)
+        } else {
+            let found = self.private_account(account_id)?;
+            let secret = found
+                .key_chain
+                .calculate_shared_secret_receiver(&encrypted.epk)?;
+            (
+                found.key_chain.private_key_holder.nullifier_secret_key,
+                secret,
+                false,
+            )
+        };
+
+        let (kind, new_account) =
+            EncryptionScheme::decrypt(&encrypted.ciphertext, &secret, commitment, ciph_id)?;
+        let new_nullifier = NullifierIndex::next_update_nullifier(account_id, &new_account, &nsk);
+
+        if is_shared {
+            self.update_shared_private_account_state(&account_id, new_account);
+        } else {
+            self.insert_private_account(account_id, kind, new_account)
+                .ok()?;
+        }
+        Some(new_nullifier)
+    }
+
     pub fn add_imported_public_account(&mut self, private_key: lee::PrivateKey) {
         let account_id = AccountId::from(&lee::PublicKey::new_from_private_key(&private_key));
 
@@ -696,8 +857,166 @@ impl Default for UserKeyChain {
 
 #[cfg(test)]
 mod tests {
+    use lee_core::encryption::EncryptedAccountData;
+
     use super::*;
 
+    #[test]
+    fn nullifier_sync_updates_sole_owned_account() {
+        let mut kc = UserKeyChain::default();
+
+        let key_chain = KeyChain::new_os_random();
+        let nsk = key_chain.private_key_holder.nullifier_secret_key;
+        let identifier = 0;
+        let account_id = AccountId::for_private_account(
+            &key_chain.nullifier_public_key,
+            &key_chain.viewing_public_key,
+            &PrivateAccountKind::Regular(identifier),
+        );
+
+        let old_account = Account::default();
+        kc.add_imported_private_account(key_chain.clone(), None, identifier, old_account.clone());
+
+        let old_nullifier =
+            Nullifier::for_account_update(&Commitment::new(&account_id, &old_account), &nsk);
+        let mut index = kc.build_latest_nullifier_index();
+        assert_eq!(index.account_for(&old_nullifier), Some(account_id));
+
+        let new_account = Account {
+            balance: 150,
+            ..Account::default()
+        };
+        let new_commitment = Commitment::new(&account_id, &new_account);
+        let (sender_ss, epk) = SharedSecretKey::encapsulate(&key_chain.viewing_public_key);
+        let ciphertext = EncryptionScheme::encrypt(
+            &new_account,
+            &PrivateAccountKind::Regular(identifier),
+            &sender_ss,
+            &new_commitment,
+            0,
+        );
+        let note = EncryptedAccountData::new(
+            ciphertext,
+            &key_chain.nullifier_public_key,
+            &key_chain.viewing_public_key,
+            epk,
+        );
+
+        let message = Message {
+            encrypted_private_post_states: vec![note],
+            new_commitments: vec![new_commitment],
+            new_nullifiers: vec![(old_nullifier, [0; 32])],
+            ..Default::default()
+        };
+
+        let handled = kc.sync_updates_via_nullifiers(&message, &mut index);
+
+        assert_eq!(handled, HashSet::from([0]));
+        assert_eq!(
+            kc.private_account(account_id).unwrap().account,
+            &new_account
+        );
+        let new_nullifier =
+            Nullifier::for_account_update(&Commitment::new(&account_id, &new_account), &nsk);
+        assert_eq!(index.account_for(&new_nullifier), Some(account_id));
+        assert!(index.account_for(&old_nullifier).is_none());
+    }
+
+    #[test]
+    fn nullifier_sync_updates_shared_account() {
+        let mut kc = UserKeyChain::default();
+
+        let label = Label::new("group");
+        let holder = GroupKeyHolder::new();
+        let identifier = 0;
+        let keys = holder.derive_regular_shared_account_keys_from_identifier(identifier);
+        let npk = keys.generate_nullifier_public_key();
+        let vpk = keys.generate_viewing_public_key();
+        let nsk = keys.nullifier_secret_key;
+        let account_id = AccountId::from((&npk, &vpk, identifier));
+
+        kc.insert_group_key_holder(label.clone(), holder);
+        let old_account = Account::default();
+        kc.insert_shared_private_account(
+            account_id,
+            SharedAccountEntry {
+                group_label: label,
+                identifier,
+                pda_seed: None,
+                authority_program_id: None,
+                account: old_account.clone(),
+            },
+        );
+
+        let old_nullifier =
+            Nullifier::for_account_update(&Commitment::new(&account_id, &old_account), &nsk);
+        let mut index = kc.build_latest_nullifier_index();
+        assert_eq!(index.account_for(&old_nullifier), Some(account_id));
+
+        let new_account = Account {
+            balance: 250,
+            ..Account::default()
+        };
+        let new_commitment = Commitment::new(&account_id, &new_account);
+        let (sender_ss, epk) = SharedSecretKey::encapsulate(&vpk);
+        let ciphertext = EncryptionScheme::encrypt(
+            &new_account,
+            &PrivateAccountKind::Regular(identifier),
+            &sender_ss,
+            &new_commitment,
+            0,
+        );
+        let note = EncryptedAccountData::new(ciphertext, &npk, &vpk, epk);
+        let message = Message {
+            encrypted_private_post_states: vec![note],
+            new_commitments: vec![new_commitment],
+            new_nullifiers: vec![(old_nullifier, [0; 32])],
+            ..Default::default()
+        };
+
+        let handled = kc.sync_updates_via_nullifiers(&message, &mut index);
+
+        assert_eq!(handled, HashSet::from([0]));
+        assert_eq!(
+            kc.shared_private_account(account_id).unwrap().account,
+            new_account
+        );
+        let new_nullifier =
+            Nullifier::for_account_update(&Commitment::new(&account_id, &new_account), &nsk);
+        assert_eq!(index.account_for(&new_nullifier), Some(account_id));
+        assert!(index.account_for(&old_nullifier).is_none());
+    }
+
+    #[test]
+    fn nullifier_sync_ignores_unindexed_nullifier() {
+        let mut kc = UserKeyChain::default();
+
+        let key_chain = KeyChain::new_os_random();
+        let identifier = 0;
+        let account_id = AccountId::for_private_account(
+            &key_chain.nullifier_public_key,
+            &key_chain.viewing_public_key,
+            &PrivateAccountKind::Regular(identifier),
+        );
+        let account = Account::default();
+        kc.add_imported_private_account(key_chain, None, identifier, account.clone());
+
+        let mut index = kc.build_latest_nullifier_index();
+        let unindexed = Nullifier::for_account_update(
+            &Commitment::new(&AccountId::new([9; 32]), &Account::default()),
+            &[9; 32],
+        );
+        let message = Message {
+            new_nullifiers: vec![(unindexed, [0; 32])],
+            ..Default::default()
+        };
+
+        let handled = kc.sync_updates_via_nullifiers(&message, &mut index);
+
+        assert!(handled.is_empty());
+        assert_eq!(kc.private_account(account_id).unwrap().account, &account);
+    }
+
     #[test]
     fn new_account() {
         let mut user_data = UserKeyChain::default();