From 50a197278879e8d9416e1c3f091e8fac44bb4513 Mon Sep 17 00:00:00 2001
From: Artem Gureev <artguril@proton.me>
Date: Fri, 10 Jul 2026 14:43:51 +0000
Subject: [PATCH] fix(wallet): catch up a shared account from genesis on
 registration

---
 lez/wallet/src/cli/account.rs       | 41 +++++++------
 lez/wallet/src/lib.rs               | 58 +++++++++++++++++-
 lez/wallet/src/storage/key_chain.rs | 95 +++++++++++++++++++++++++++++
 3 files changed, 175 insertions(+), 19 deletions(-)

diff --git a/lez/wallet/src/cli/account.rs b/lez/wallet/src/cli/account.rs
index c700d2e9..4aa3b4c8 100644
--- a/lez/wallet/src/cli/account.rs
+++ b/lez/wallet/src/cli/account.rs
@@ -197,7 +197,7 @@ impl NewSubcommand {
         Ok(SubcommandReturnValue::RegisterAccount { account_id })
     }
 
-    fn handle_private_gms(
+    async fn handle_private_gms(
         group: &Label,
         label: Option<Label>,
         pda: bool,
@@ -229,14 +229,18 @@ impl NewSubcommand {
                 pid[i] = u32::from_le_bytes(chunk.try_into().unwrap());
             }
 
-            wallet_core.create_shared_pda_account(
-                group.clone(),
-                pda_seed,
-                pid,
-                identifier.unwrap_or_else(rand::random),
-            )?
+            wallet_core
+                .create_shared_pda_account(
+                    group.clone(),
+                    pda_seed,
+                    pid,
+                    identifier.unwrap_or_else(rand::random),
+                )
+                .await?
         } else {
-            wallet_core.create_shared_regular_account(group.clone())?
+            wallet_core
+                .create_shared_regular_account(group.clone())
+                .await?
         };
 
         if let Some(label) = label {
@@ -295,15 +299,18 @@ impl WalletSubcommand for NewSubcommand {
                 seed,
                 program_id,
                 identifier,
-            } => Self::handle_private_gms(
-                &group,
-                label,
-                pda,
-                seed,
-                program_id,
-                identifier,
-                wallet_core,
-            ),
+            } => {
+                Self::handle_private_gms(
+                    &group,
+                    label,
+                    pda,
+                    seed,
+                    program_id,
+                    identifier,
+                    wallet_core,
+                )
+                .await
+            }
             Self::PrivateAccountsKey { cci } => Self::handle_private_accounts_key(cci, wallet_core),
         }
     }
diff --git a/lez/wallet/src/lib.rs b/lez/wallet/src/lib.rs
index 4984f784..accd6fdf 100644
--- a/lez/wallet/src/lib.rs
+++ b/lez/wallet/src/lib.rs
@@ -415,8 +415,60 @@ impl WalletCore {
         );
     }
 
+    /// Re-derive a shared account's state by scanning its keypair from genesis to the current
+    /// synced block. The init note's nullifier is deterministic on ID, so we await it and let
+    /// the nullifier pass decode the init and every subsequent update.
+    ///
+    /// If no initialization is found, will return `Ok` and default to usual hot-sync.
+    async fn catch_up_shared_account(&mut self, account_id: AccountId) -> Result<()> {
+        use futures::TryStreamExt as _;
+
+        let cursor = self.storage.last_synced_block();
+        if cursor == 0
+            || self
+                .storage
+                .key_chain()
+                .shared_private_account(account_id)
+                .is_none()
+        {
+            return Ok(());
+        }
+
+        info!("Scanning shared account {account_id:#?} from genesis to block {cursor}");
+
+        let mut index = NullifierIndex::default();
+        index.track_initialization(account_id);
+
+        let poller = self.optimal_poller();
+        let mut blocks = std::pin::pin!(poller.poll_block_range(1..=cursor));
+        while let Some(block) = blocks.try_next().await? {
+            for tx in block.body.transactions {
+                let LeeTransaction::PrivacyPreserving(pp_tx) = &tx else {
+                    continue;
+                };
+                pp_tx.message.validate_note_lengths()?;
+                // Sync updates while watching only the init nullifier.
+                self.storage
+                    .key_chain_mut()
+                    .sync_updates_via_nullifiers(&pp_tx.message, &mut index);
+            }
+        }
+
+        let now = self.storage.last_synced_block();
+        // This is a defence-in-depth. Currently during the async update the cursor
+        // cannot advance. However, de-sync can be possible later. This hard error
+        // will signal this.
+        if now != cursor {
+            return Err(anyhow::anyhow!(
+                "Shared-account catched-up to {cursor} with a cursor de-sync advancing to {now}"
+            ));
+        }
+
+        Ok(())
+    }
+
     /// Create a shared PDA account from a group's GMS. Returns the `AccountId` and derived keys.
-    pub fn create_shared_pda_account(
+    pub async fn create_shared_pda_account(
         &mut self,
         group_name: Label,
         pda_seed: lee_core::program::PdaSeed,
@@ -441,6 +493,7 @@ impl WalletCore {
             Some(pda_seed),
             Some(program_id),
         );
+        self.catch_up_shared_account(account_id).await?;
 
         Ok(SharedAccountInfo {
             account_id,
@@ -451,7 +504,7 @@ impl WalletCore {
 
     /// Create a shared regular private account from a group's GMS. Returns the `AccountId` and
     /// derived keys. The derivation seed is computed deterministically from a random identifier.
-    pub fn create_shared_regular_account(
+    pub async fn create_shared_regular_account(
         &mut self,
         group_name: Label,
     ) -> Result<SharedAccountInfo> {
@@ -469,6 +522,7 @@ impl WalletCore {
         let account_id = AccountId::from((&npk, &vpk, identifier));
 
         self.register_shared_account(account_id, group_name, identifier, None, None);
+        self.catch_up_shared_account(account_id).await?;
 
         Ok(SharedAccountInfo {
             account_id,
diff --git a/lez/wallet/src/storage/key_chain.rs b/lez/wallet/src/storage/key_chain.rs
index 72dcb78f..bb1971d7 100644
--- a/lez/wallet/src/storage/key_chain.rs
+++ b/lez/wallet/src/storage/key_chain.rs
@@ -90,6 +90,14 @@ impl NullifierIndex {
         );
     }
 
+    /// Indexes `account_id` by the nullifier its initialization publishes.
+    pub fn track_initialization(&mut self, account_id: AccountId) {
+        self.0.insert(
+            Nullifier::for_account_initialization(&account_id),
+            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);
@@ -987,6 +995,93 @@ mod tests {
         assert!(index.account_for(&old_nullifier).is_none());
     }
 
+    // The genesis catch-up seeds only the init nullifier and lets the nullifier pass decode the
+    // init note and every subsequent (randomly-tagged) update. Verify a shared account rolls from
+    // default through its init to a later update purely by nullifier — the path the catch-up runs.
+    #[test]
+    fn nullifier_sync_catches_up_shared_account_from_init() {
+        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);
+        kc.insert_shared_private_account(
+            account_id,
+            SharedAccountEntry {
+                group_label: label,
+                identifier,
+                pda_seed: None,
+                authority_program_id: None,
+                account: Account::default(),
+            },
+        );
+
+        let mut index = NullifierIndex::default();
+        index.track_initialization(account_id);
+
+        // A note publishing `spent` and carrying the state `next`.
+        let make_message = |spent: Nullifier, next: &Account| {
+            let commitment = Commitment::new(&account_id, next);
+            let (sender_ss, epk) = SharedSecretKey::encapsulate(&vpk);
+            let ciphertext = EncryptionScheme::encrypt(
+                next,
+                &PrivateAccountKind::Regular(identifier),
+                &sender_ss,
+                &commitment,
+                0,
+            );
+            let note = EncryptedAccountData::new(ciphertext, &npk, &vpk, epk);
+            Message {
+                encrypted_private_post_states: vec![note],
+                new_commitments: vec![commitment],
+                new_nullifiers: vec![(spent, [0; 32])],
+                ..Default::default()
+            }
+        };
+
+        // Init: default -> initialized, discovered via the seeded init nullifier.
+        let initialized = Account {
+            balance: 250,
+            ..Account::default()
+        };
+        let init_msg = make_message(
+            Nullifier::for_account_initialization(&account_id),
+            &initialized,
+        );
+        assert_eq!(
+            kc.sync_updates_via_nullifiers(&init_msg, &mut index),
+            HashSet::from([0])
+        );
+        assert_eq!(
+            kc.shared_private_account(account_id).unwrap().account,
+            initialized
+        );
+
+        // Update: initialized -> updated, discovered via the now-tracked update nullifier.
+        let updated = Account {
+            balance: 500,
+            ..Account::default()
+        };
+        let update_spent =
+            Nullifier::for_account_update(&Commitment::new(&account_id, &initialized), &nsk);
+        let update_msg = make_message(update_spent, &updated);
+        assert_eq!(
+            kc.sync_updates_via_nullifiers(&update_msg, &mut index),
+            HashSet::from([0])
+        );
+        assert_eq!(
+            kc.shared_private_account(account_id).unwrap().account,
+            updated
+        );
+    }
+
     #[test]
     fn nullifier_sync_ignores_unindexed_nullifier() {
         let mut kc = UserKeyChain::default();