From 5fbe0ce9c194c6913bc56648b9abe0603746b6e0 Mon Sep 17 00:00:00 2001 From: Artem Gureev Date: Tue, 30 Jun 2026 10:24:27 +0000 Subject: [PATCH] refactor(wallet): account-id computred via PrivateAddressPlaintext --- .../src/key_management/group_key_holder.rs | 21 +++++----- .../src/key_management/key_tree/mod.rs | 9 +++-- lez/wallet/src/account_manager.rs | 16 +++++--- lez/wallet/src/cli/account.rs | 9 +++-- lez/wallet/src/lib.rs | 7 ++-- lez/wallet/src/storage/key_chain.rs | 38 +++++++++++-------- 6 files changed, 58 insertions(+), 42 deletions(-) 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 7bb94792..849d7b8c 100644 --- a/lee/key_protocol/src/key_management/group_key_holder.rs +++ b/lee/key_protocol/src/key_management/group_key_holder.rs @@ -323,11 +323,11 @@ mod tests { /// Pins the end-to-end derivation for a fixed (GMS, `ProgramId`, `PdaSeed`). Any change /// to `secret_spending_key_for_pda`, the `PrivateKeyHolder` nsk/npk chain, or the - /// `AccountId::for_private_pda` formula breaks this test. Mirrors the pinned-value + /// `PrivateAddressPlaintext::pda_account_id` formula breaks this test. Mirrors the pinned-value /// pattern from `for_private_pda_matches_pinned_value` in `lee_core`. #[test] fn pinned_end_to_end_derivation_for_private_pda() { - use lee_core::{account::AccountId, program::ProgramId}; + use lee_core::{account::PrivateAddressPlaintext, program::ProgramId}; let gms = [42_u8; 32]; let seed = PdaSeed::new([1; 32]); @@ -337,7 +337,8 @@ mod tests { let keys = holder.derive_keys_for_pda(&TEST_PROGRAM_ID, &seed); let npk = keys.generate_nullifier_public_key(); let vpk = keys.generate_viewing_public_key(); - let account_id = AccountId::for_private_pda(&program_id, &seed, &npk, &vpk, u128::MAX); + let account_id = PrivateAddressPlaintext::new(npk, vpk.clone(), u128::MAX) + .pda_account_id(&program_id, &seed); let expected_npk = NullifierPublicKey([ 136, 176, 234, 71, 208, 8, 143, 142, 126, 155, 132, 18, 71, 27, 88, 56, 100, 90, 79, @@ -345,8 +346,8 @@ mod tests { ]); // AccountId is derived from (program_id, seed, npk), so it changes when npk changes. // We verify npk is pinned, and AccountId is deterministically derived from it. - let expected_account_id = - AccountId::for_private_pda(&program_id, &seed, &expected_npk, &vpk, u128::MAX); + let expected_account_id = PrivateAddressPlaintext::new(expected_npk, vpk, u128::MAX) + .pda_account_id(&program_id, &seed); assert_eq!(npk, expected_npk); assert_eq!(account_id, expected_account_id); @@ -524,7 +525,7 @@ mod tests { /// Full lifecycle: create group, distribute GMS via seal/unseal, verify key agreement. #[test] fn group_pda_lifecycle() { - use lee_core::account::AccountId; + use lee_core::account::PrivateAddressPlaintext; let alice_holder = GroupKeyHolder::new(); let pda_seed = PdaSeed::new([42_u8; 32]); @@ -549,10 +550,10 @@ mod tests { let alice_vpk = alice_keys.generate_viewing_public_key(); let bob_group_vpk = bob_group_keys.generate_viewing_public_key(); - let alice_account_id = - AccountId::for_private_pda(&program_id, &pda_seed, &alice_npk, &alice_vpk, 0); - let bob_account_id = - AccountId::for_private_pda(&program_id, &pda_seed, &bob_npk, &bob_group_vpk, 0); + let alice_account_id = PrivateAddressPlaintext::new(alice_npk, alice_vpk, 0) + .pda_account_id(&program_id, &pda_seed); + let bob_account_id = PrivateAddressPlaintext::new(bob_npk, bob_group_vpk, 0) + .pda_account_id(&program_id, &pda_seed); assert_eq!(alice_account_id, bob_account_id); } diff --git a/lee/key_protocol/src/key_management/key_tree/mod.rs b/lee/key_protocol/src/key_management/key_tree/mod.rs index e7d490d2..c592647b 100644 --- a/lee/key_protocol/src/key_management/key_tree/mod.rs +++ b/lee/key_protocol/src/key_management/key_tree/mod.rs @@ -275,11 +275,12 @@ impl KeyTree { identifier: Identifier, ) -> Option { let node = self.key_map.get(cci)?; - let account_id = lee::AccountId::for_regular_private_account( - &node.value.0.nullifier_public_key, - &node.value.0.viewing_public_key, + let account_id = lee::PrivateAddressPlaintext::new( + node.value.0.nullifier_public_key, + node.value.0.viewing_public_key.clone(), identifier, - ); + ) + .account_id(); if self.account_id_map.contains_key(&account_id) { return None; } diff --git a/lez/wallet/src/account_manager.rs b/lez/wallet/src/account_manager.rs index f77d1062..226d4b5b 100644 --- a/lez/wallet/src/account_manager.rs +++ b/lez/wallet/src/account_manager.rs @@ -2,7 +2,7 @@ use core::fmt; use anyhow::Result; use keycard_wallet::{KeycardWallet, python_path}; -use lee::{AccountId, PrivateKey, PublicKey, Signature}; +use lee::{AccountId, PrivateAddressPlaintext, PrivateKey, PublicKey, Signature}; use lee_core::{ Identifier, InputAccountIdentity, MembershipProof, NullifierPublicKey, NullifierSecretKey, SharedSecretKey, @@ -30,7 +30,7 @@ pub enum AccountIdentity { identifier: Identifier, }, /// An owned private PDA: wallet holds the nsk/npk; `account_id` was derived via - /// [`AccountId::for_private_pda`]. + /// [`PrivateAddressPlaintext::pda_account_id`]. PrivatePdaOwned(AccountId), /// A foreign private PDA: wallet knows the recipient's npk/vpk but not their nsk. /// Uses a default (uninitialised) account. @@ -50,7 +50,7 @@ pub enum AccountIdentity { identifier: Identifier, }, /// A shared private PDA with externally-provided keys (e.g. from GMS). - /// `account_id` was derived via [`AccountId::for_private_pda`]. + /// `account_id` was derived via [`PrivateAddressPlaintext::pda_account_id`]. PrivatePdaShared { account_id: AccountId, nsk: NullifierSecretKey, @@ -261,7 +261,11 @@ impl AccountManager { identifier, } => { let acc = lee_core::account::Account::default(); - let auth_acc = AccountWithMetadata::new(acc, false, (&npk, &vpk, identifier)); + let auth_acc = AccountWithMetadata::new( + acc, + false, + PrivateAddressPlaintext::new(npk, vpk.clone(), identifier).account_id(), + ); let mut random_seed: [u8; 32] = [0; 32]; OsRng.fill_bytes(&mut random_seed); let pre = AccountPreparedData { @@ -309,7 +313,9 @@ impl AccountManager { vpk, identifier, } => { - let account_id = lee::AccountId::from((&npk, &vpk, identifier)); + let account_id = + lee::PrivateAddressPlaintext::new(npk, vpk.clone(), identifier) + .account_id(); let pre = private_shared_acc_preparation( wallet, account_id, nsk, npk, vpk, identifier, false, ) diff --git a/lez/wallet/src/cli/account.rs b/lez/wallet/src/cli/account.rs index 2f3dfa9f..42e3df4c 100644 --- a/lez/wallet/src/cli/account.rs +++ b/lez/wallet/src/cli/account.rs @@ -525,11 +525,12 @@ impl WalletSubcommand for ImportSubcommand { let key_chain: KeyChain = serde_json::from_str(&key_chain_json) .map_err(|err| anyhow::anyhow!("Invalid key chain JSON: {err}"))?; let account = lee::Account::from(account_state); - let account_id = lee::AccountId::from(( - &key_chain.nullifier_public_key, - &key_chain.viewing_public_key, + let account_id = lee::PrivateAddressPlaintext::new( + key_chain.nullifier_public_key, + key_chain.viewing_public_key.clone(), identifier, - )); + ) + .account_id(); wallet_core .storage_mut() diff --git a/lez/wallet/src/lib.rs b/lez/wallet/src/lib.rs index 010e438e..d71a795b 100644 --- a/lez/wallet/src/lib.rs +++ b/lez/wallet/src/lib.rs @@ -17,7 +17,7 @@ use common::{HashType, transaction::LeeTransaction}; use config::WalletConfig; use key_protocol::key_management::key_tree::chain_index::ChainIndex; use lee::{ - Account, AccountId, PrivacyPreservingTransaction, + Account, AccountId, PrivacyPreservingTransaction, PrivateAddressPlaintext, privacy_preserving_transaction::{ circuit::ProgramWithDependencies, message::EncryptedAccountData, }, @@ -377,7 +377,8 @@ impl WalletCore { let keys = holder.derive_keys_for_pda(&program_id, &pda_seed); let npk = keys.generate_nullifier_public_key(); let vpk = keys.generate_viewing_public_key(); - let account_id = AccountId::for_private_pda(&program_id, &pda_seed, &npk, &vpk, identifier); + let account_id = PrivateAddressPlaintext::new(npk, vpk.clone(), identifier) + .pda_account_id(&program_id, &pda_seed); self.register_shared_account( account_id, @@ -419,7 +420,7 @@ impl WalletCore { let keys = holder.derive_keys_for_shared_account(&derivation_seed); let npk = keys.generate_nullifier_public_key(); let vpk = keys.generate_viewing_public_key(); - let account_id = AccountId::from((&npk, &vpk, identifier)); + let account_id = PrivateAddressPlaintext::new(npk, vpk.clone(), identifier).account_id(); self.register_shared_account(account_id, group_name, identifier, None, None); diff --git a/lez/wallet/src/storage/key_chain.rs b/lez/wallet/src/storage/key_chain.rs index 5db90119..bfecb9d6 100644 --- a/lez/wallet/src/storage/key_chain.rs +++ b/lez/wallet/src/storage/key_chain.rs @@ -696,6 +696,8 @@ impl Default for UserKeyChain { #[cfg(test)] mod tests { + use lee::PrivateAddressPlaintext; + use super::*; #[test] @@ -734,11 +736,12 @@ mod tests { let mut user_data = UserKeyChain::default(); let key_chain = KeyChain::new_os_random(); - let account_id = AccountId::from(( - &key_chain.nullifier_public_key, - &key_chain.viewing_public_key, + let account_id = PrivateAddressPlaintext::new( + key_chain.nullifier_public_key, + key_chain.viewing_public_key.clone(), 0, - )); + ) + .account_id(); let account = lee_core::account::Account::default(); user_data.add_imported_private_account(key_chain, None, 0, account); @@ -753,11 +756,12 @@ mod tests { let mut user_data = UserKeyChain::default(); let key_chain = KeyChain::new_os_random(); - let account_id = AccountId::from(( - &key_chain.nullifier_public_key, - &key_chain.viewing_public_key, + let account_id = PrivateAddressPlaintext::new( + key_chain.nullifier_public_key, + key_chain.viewing_public_key.clone(), 0, - )); + ) + .account_id(); let account = lee_core::account::Account::default(); user_data.add_imported_private_account(key_chain, None, 0, account.clone()); @@ -802,11 +806,12 @@ mod tests { let mut user_data = UserKeyChain::default(); let key_chain = KeyChain::new_os_random(); - let account_id = AccountId::from(( - &key_chain.nullifier_public_key, - &key_chain.viewing_public_key, + let account_id = PrivateAddressPlaintext::new( + key_chain.nullifier_public_key, + key_chain.viewing_public_key, 0, - )); + ) + .account_id(); let new_account = lee_core::account::Account { balance: 100, @@ -827,11 +832,12 @@ mod tests { let mut user_data = UserKeyChain::default(); let key_chain = KeyChain::new_os_random(); - let account_id1 = AccountId::from(( - &key_chain.nullifier_public_key, - &key_chain.viewing_public_key, + let account_id1 = PrivateAddressPlaintext::new( + key_chain.nullifier_public_key, + key_chain.viewing_public_key.clone(), 0, - )); + ) + .account_id(); let account = lee_core::account::Account::default(); user_data.add_imported_private_account(key_chain, None, 0, account);