diff --git a/artifacts/lee/privacy_preserving_circuit/privacy_preserving_circuit.bin b/artifacts/lee/privacy_preserving_circuit/privacy_preserving_circuit.bin index bacba91e..bb1b66fc 100644 Binary files a/artifacts/lee/privacy_preserving_circuit/privacy_preserving_circuit.bin and b/artifacts/lee/privacy_preserving_circuit/privacy_preserving_circuit.bin differ diff --git a/artifacts/lez/programs/amm.bin b/artifacts/lez/programs/amm.bin index 11382930..2bf13353 100644 Binary files a/artifacts/lez/programs/amm.bin and b/artifacts/lez/programs/amm.bin differ diff --git a/artifacts/lez/programs/associated_token_account.bin b/artifacts/lez/programs/associated_token_account.bin index 1f01ba89..31604d5a 100644 Binary files a/artifacts/lez/programs/associated_token_account.bin and b/artifacts/lez/programs/associated_token_account.bin differ diff --git a/artifacts/lez/programs/authenticated_transfer.bin b/artifacts/lez/programs/authenticated_transfer.bin index 2fc88f43..25ff2172 100644 Binary files a/artifacts/lez/programs/authenticated_transfer.bin and b/artifacts/lez/programs/authenticated_transfer.bin differ diff --git a/artifacts/lez/programs/bridge.bin b/artifacts/lez/programs/bridge.bin index 8ee361a7..0063fff9 100644 Binary files a/artifacts/lez/programs/bridge.bin and b/artifacts/lez/programs/bridge.bin differ diff --git a/artifacts/lez/programs/clock.bin b/artifacts/lez/programs/clock.bin index 45467bd3..10cc941c 100644 Binary files a/artifacts/lez/programs/clock.bin and b/artifacts/lez/programs/clock.bin differ diff --git a/artifacts/lez/programs/faucet.bin b/artifacts/lez/programs/faucet.bin index 6204eb9c..c389298b 100644 Binary files a/artifacts/lez/programs/faucet.bin and b/artifacts/lez/programs/faucet.bin differ diff --git a/artifacts/lez/programs/pinata.bin b/artifacts/lez/programs/pinata.bin index 8f16d3e1..82e7fb78 100644 Binary files a/artifacts/lez/programs/pinata.bin and b/artifacts/lez/programs/pinata.bin differ diff --git a/artifacts/lez/programs/pinata_token.bin b/artifacts/lez/programs/pinata_token.bin index 0caf5046..08900cda 100644 Binary files a/artifacts/lez/programs/pinata_token.bin and b/artifacts/lez/programs/pinata_token.bin differ diff --git a/artifacts/lez/programs/token.bin b/artifacts/lez/programs/token.bin index 1786fa01..568c6f40 100644 Binary files a/artifacts/lez/programs/token.bin and b/artifacts/lez/programs/token.bin differ diff --git a/artifacts/lez/programs/vault.bin b/artifacts/lez/programs/vault.bin index d670b6bc..91482c04 100644 Binary files a/artifacts/lez/programs/vault.bin and b/artifacts/lez/programs/vault.bin differ diff --git a/lee/key_protocol/src/key_management/group_key_holder.rs b/lee/key_protocol/src/key_management/group_key_holder.rs index 0821e4cd..1aef6c91 100644 --- a/lee/key_protocol/src/key_management/group_key_holder.rs +++ b/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 diff --git a/lee/state_machine/core/src/program/mod.rs b/lee/state_machine/core/src/program/mod.rs index 0d96642e..770bcf2d 100644 --- a/lee/state_machine/core/src/program/mod.rs +++ b/lee/state_machine/core/src/program/mod.rs @@ -324,7 +324,7 @@ impl AccountPostState { pub type BlockValidityWindow = ValidityWindow; pub type TimestampValidityWindow = ValidityWindow; -#[derive(Clone, Copy, Serialize, Deserialize)] +#[derive(Clone, Copy, Default, Serialize, Deserialize)] #[cfg_attr( any(feature = "host", test), derive(Debug, PartialEq, Eq, BorshSerialize, BorshDeserialize) diff --git a/lee/state_machine/src/privacy_preserving_transaction/message.rs b/lee/state_machine/src/privacy_preserving_transaction/message.rs index 73e62021..c1f2615c 100644 --- a/lee/state_machine/src/privacy_preserving_transaction/message.rs +++ b/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, pub nonces: Vec, @@ -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 { + 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 { diff --git a/lez/wallet/src/lib.rs b/lez/wallet/src/lib.rs index c8821b8f..4b4ba149 100644 --- a/lez/wallet/src/lib.rs +++ b/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 { 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, + ) { 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::>() }) .collect::>(); - 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, + ) { 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); diff --git a/lez/wallet/src/storage/key_chain.rs b/lez/wallet/src/storage/key_chain.rs index 5db90119..72dcb78f 100644 --- a/lez/wallet/src/storage/key_chain.rs +++ b/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); + +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 { + 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> { + 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> { 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 { + 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 { + 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 { + 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();