From 48f95b1b7aaec7aad388a5ca5a2b5d11d16edd42 Mon Sep 17 00:00:00 2001 From: Moudy Date: Mon, 27 Apr 2026 02:43:51 +0200 Subject: [PATCH] feat: add GroupKeyHolder storage and PrivateGroupPda wallet variant --- key_protocol/src/key_protocol_core/mod.rs | 49 +++++++++++ wallet/src/lib.rs | 8 +- wallet/src/privacy_preserving_tx.rs | 102 +++++++++++++++++++++- 3 files changed, 157 insertions(+), 2 deletions(-) diff --git a/key_protocol/src/key_protocol_core/mod.rs b/key_protocol/src/key_protocol_core/mod.rs index 8186865f..b0ee18cb 100644 --- a/key_protocol/src/key_protocol_core/mod.rs +++ b/key_protocol/src/key_protocol_core/mod.rs @@ -6,6 +6,7 @@ use serde::{Deserialize, Serialize}; use crate::key_management::{ KeyChain, + group_key_holder::GroupKeyHolder, key_tree::{KeyTreePrivate, KeyTreePublic, chain_index::ChainIndex}, secret_holders::SeedHolder, }; @@ -23,6 +24,17 @@ pub struct NSSAUserData { pub public_key_tree: KeyTreePublic, /// Tree of private keys. pub private_key_tree: KeyTreePrivate, + /// Group key holders for private PDA groups, keyed by a human-readable label. + /// Defaults to empty for backward compatibility with wallets that predate group PDAs. + /// An older wallet binary that re-serializes this struct will drop the field. + #[serde(default)] + pub group_key_holders: BTreeMap, + /// Cached plaintext state of group PDA accounts, keyed by `AccountId`. + /// Updated after each group PDA transaction by decrypting the circuit output. + /// The sequencer only stores encrypted commitments, so this local cache is the + /// only source of plaintext state for group PDAs. + #[serde(default)] + pub group_pda_accounts: BTreeMap, } impl NSSAUserData { @@ -81,6 +93,8 @@ impl NSSAUserData { default_user_private_accounts: default_accounts_key_chains, public_key_tree, private_key_tree, + group_key_holders: BTreeMap::new(), + group_pda_accounts: BTreeMap::new(), }) } @@ -177,6 +191,20 @@ impl NSSAUserData { .copied() .chain(self.private_key_tree.account_id_map.keys().copied()) } + + /// Returns the `GroupKeyHolder` for the given label, if it exists. + #[must_use] + pub fn get_group_key_holder(&self, label: &str) -> Option<&GroupKeyHolder> { + self.group_key_holders.get(label) + } + + /// Inserts or replaces a `GroupKeyHolder` under the given label. + /// + /// If a holder already exists under this label, it is silently replaced and the old + /// GMS is lost. Callers must ensure label uniqueness across groups. + pub fn insert_group_key_holder(&mut self, label: String, holder: GroupKeyHolder) { + self.group_key_holders.insert(label, holder); + } } impl Default for NSSAUserData { @@ -196,6 +224,27 @@ impl Default for NSSAUserData { mod tests { use super::*; + #[test] + fn group_key_holder_storage_round_trip() { + let mut user_data = NSSAUserData::default(); + assert!(user_data.get_group_key_holder("test-group").is_none()); + + let holder = GroupKeyHolder::from_gms([42_u8; 32]); + user_data.insert_group_key_holder(String::from("test-group"), holder.clone()); + + let retrieved = user_data + .get_group_key_holder("test-group") + .expect("should exist"); + assert_eq!(retrieved.dangerous_raw_gms(), holder.dangerous_raw_gms()); + assert_eq!(retrieved.epoch(), holder.epoch()); + } + + #[test] + fn group_key_holders_default_empty() { + let user_data = NSSAUserData::default(); + assert!(user_data.group_key_holders.is_empty()); + } + #[test] fn new_account() { let mut user_data = NSSAUserData::default(); diff --git a/wallet/src/lib.rs b/wallet/src/lib.rs index 460cfcfd..e546ceec 100644 --- a/wallet/src/lib.rs +++ b/wallet/src/lib.rs @@ -41,7 +41,7 @@ pub mod cli; pub mod config; pub mod helperfunctions; pub mod poller; -mod privacy_preserving_tx; +pub mod privacy_preserving_tx; pub mod program_facades; pub const HOME_DIR_ENV_VAR: &str = "NSSA_WALLET_HOME_DIR"; @@ -201,6 +201,12 @@ impl WalletCore { &self.storage } + /// Get mutable storage (e.g. for inserting group key holders). + #[must_use] + pub const fn storage_mut(&mut self) -> &mut WalletChainStore { + &mut self.storage + } + /// Restore storage from an existing mnemonic phrase. pub fn restore_storage(&mut self, mnemonic: &Mnemonic, password: &str) -> Result<()> { self.storage = WalletChainStore::restore_storage( diff --git a/wallet/src/privacy_preserving_tx.rs b/wallet/src/privacy_preserving_tx.rs index 14a805c7..7272e33a 100644 --- a/wallet/src/privacy_preserving_tx.rs +++ b/wallet/src/privacy_preserving_tx.rs @@ -5,6 +5,7 @@ use nssa_core::{ MembershipProof, NullifierPublicKey, NullifierSecretKey, SharedSecretKey, account::{AccountWithMetadata, Nonce}, encryption::{EphemeralPublicKey, ViewingPublicKey}, + program::{PdaSeed, ProgramId}, }; use crate::{ExecutionFailureKind, WalletCore}; @@ -17,6 +18,14 @@ pub enum PrivacyPreservingAccount { npk: NullifierPublicKey, vpk: ViewingPublicKey, }, + /// A private PDA owned by a group. The wallet derives keys from the + /// `GroupKeyHolder` stored under `group_label`, then computes the + /// `AccountId` via `AccountId::for_private_pda(program_id, seed, npk)`. + PrivateGroupPda { + group_label: String, + program_id: ProgramId, + seed: PdaSeed, + }, } impl PrivacyPreservingAccount { @@ -29,7 +38,9 @@ impl PrivacyPreservingAccount { pub const fn is_private(&self) -> bool { matches!( &self, - Self::PrivateOwned(_) | Self::PrivateForeign { npk: _, vpk: _ } + Self::PrivateOwned(_) + | Self::PrivateForeign { npk: _, vpk: _ } + | Self::PrivateGroupPda { .. } ) } } @@ -94,6 +105,16 @@ impl AccountManager { (State::Private(pre), 2) } + PrivacyPreservingAccount::PrivateGroupPda { + group_label, + program_id, + seed, + } => { + let pre = + group_pda_preparation(wallet, &group_label, &program_id, &seed).await?; + + (State::Private(pre), 3) + } }; pre_states.push(state); @@ -106,6 +127,7 @@ impl AccountManager { }) } + #[must_use] pub fn pre_states(&self) -> Vec { self.states .iter() @@ -116,10 +138,12 @@ impl AccountManager { .collect() } + #[must_use] pub fn visibility_mask(&self) -> &[u8] { &self.visibility_mask } + #[must_use] pub fn public_account_nonces(&self) -> Vec { self.states .iter() @@ -130,6 +154,7 @@ impl AccountManager { .collect() } + #[must_use] pub fn private_account_keys(&self) -> Vec { self.states .iter() @@ -149,6 +174,7 @@ impl AccountManager { .collect() } + #[must_use] pub fn private_account_auth(&self) -> Vec { self.states .iter() @@ -159,6 +185,7 @@ impl AccountManager { .collect() } + #[must_use] pub fn private_account_membership_proofs(&self) -> Vec> { self.states .iter() @@ -169,6 +196,7 @@ impl AccountManager { .collect() } + #[must_use] pub fn public_account_ids(&self) -> Vec { self.states .iter() @@ -179,6 +207,7 @@ impl AccountManager { .collect() } + #[must_use] pub fn public_account_auth(&self) -> Vec<&PrivateKey> { self.states .iter() @@ -198,6 +227,61 @@ struct AccountPreparedData { proof: Option, } +async fn group_pda_preparation( + wallet: &WalletCore, + group_label: &str, + program_id: &ProgramId, + seed: &PdaSeed, +) -> Result { + let holder = wallet + .storage + .user_data + .get_group_key_holder(group_label) + .ok_or(ExecutionFailureKind::KeyNotFoundError)?; + + let keys = holder.derive_keys_for_pda(seed); + let npk = keys.generate_nullifier_public_key(); + let vpk = keys.generate_viewing_public_key(); + let nsk = keys.nullifier_secret_key; + let account_id = nssa::AccountId::for_private_pda(program_id, seed, &npk); + + // Check local cache first (private PDA state is encrypted on-chain, the sequencer + // only stores commitments). Fall back to default for new PDAs. + let acc = wallet + .storage + .user_data + .group_pda_accounts + .get(&account_id) + .cloned() + .unwrap_or_default(); + + let exists = acc != nssa_core::account::Account::default(); + + // is_authorized tracks whether the account existed on-chain before this tx. + // NSK is only provided for existing accounts: the circuit consumes NSKs sequentially + // from an iterator and asserts none are left over, so supplying an NSK for a new + // (unauthorized) account would trigger the over-supply assertion. This matches the + // PrivateForeign path (nsk: None for unauthorized accounts). + let pre_state = AccountWithMetadata::new(acc, exists, account_id); + + let proof = if exists { + wallet + .check_private_account_initialized(account_id) + .await + .unwrap_or(None) + } else { + None + }; + + Ok(AccountPreparedData { + nsk: exists.then_some(nsk), + npk, + vpk, + pre_state, + proof, + }) +} + async fn private_acc_preparation( wallet: &WalletCore, account_id: AccountId, @@ -234,3 +318,19 @@ async fn private_acc_preparation( proof, }) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn private_group_pda_is_private() { + let acc = PrivacyPreservingAccount::PrivateGroupPda { + group_label: String::from("test"), + program_id: [1; 8], + seed: PdaSeed::new([2; 32]), + }; + assert!(acc.is_private()); + assert!(!acc.is_public()); + } +}