From 69b81ea6217a5bc03875c76cdb7fd8ddfc325a60 Mon Sep 17 00:00:00 2001 From: Moudy Date: Thu, 7 May 2026 17:35:51 +0200 Subject: [PATCH] fix: address review feedback, persist group data in wallet storage --- integration_tests/tests/shared_accounts.rs | 210 ++++++++++++++++++ .../src/key_management/group_key_holder.rs | 89 ++++---- key_protocol/src/key_protocol_core/mod.rs | 19 +- wallet/src/cli/account.rs | 92 ++------ wallet/src/config.rs | 12 + wallet/src/helperfunctions.rs | 2 + wallet/src/lib.rs | 117 ++++++++-- wallet/src/privacy_preserving_tx.rs | 4 +- 8 files changed, 411 insertions(+), 134 deletions(-) create mode 100644 integration_tests/tests/shared_accounts.rs diff --git a/integration_tests/tests/shared_accounts.rs b/integration_tests/tests/shared_accounts.rs new file mode 100644 index 00000000..b91502d1 --- /dev/null +++ b/integration_tests/tests/shared_accounts.rs @@ -0,0 +1,210 @@ +#![expect( + clippy::tests_outside_test_module, + reason = "Integration test file, not inside a #[cfg(test)] module" +)] + +//! Shared account integration tests. +//! +//! Demonstrates: +//! 1. Group creation and GMS distribution via seal/unseal. +//! 2. Shared regular private account creation via `--for-gms`. +//! 3. Funding a shared account from a public account. +//! 4. Syncing discovers the funded shared account state. + +use std::time::Duration; + +use anyhow::{Context as _, Result}; +use integration_tests::{TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, format_public_account_id}; +use log::info; +use tokio::test; +use wallet::cli::{ + Command, SubcommandReturnValue, + account::{AccountSubcommand, NewSubcommand}, + group::GroupSubcommand, + programs::native_token_transfer::AuthTransferSubcommand, +}; + +/// Create a group, create a shared account from it, and verify registration. +#[test] +async fn group_create_and_shared_account_registration() -> Result<()> { + let mut ctx = TestContext::new().await?; + + // Create a group + let command = Command::Group(GroupSubcommand::New { + name: "test-group".to_string(), + }); + wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; + + // Verify group exists + assert!( + ctx.wallet() + .storage() + .user_data + .group_key_holder("test-group") + .is_some() + ); + + // Create a shared regular private account from the group + let command = Command::Account(AccountSubcommand::New(NewSubcommand::Private { + cci: None, + label: Some("shared-acc".to_string()), + for_gms: Some("test-group".to_string()), + pda: false, + seed: None, + program_id: None, + })); + + let result = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; + let SubcommandReturnValue::RegisterAccount { + account_id: shared_account_id, + } = result + else { + anyhow::bail!("Expected RegisterAccount return value"); + }; + + // Verify shared account is registered in storage + let shared_private_accounts = &ctx.wallet().storage().user_data.shared_private_accounts; + let entry = shared_private_accounts + .get(&shared_account_id) + .context("Shared account not found in storage")?; + assert_eq!(entry.group_label, "test-group"); + assert!(entry.pda_seed.is_none()); + + info!("Shared account registered: {shared_account_id}"); + Ok(()) +} + +/// GMS seal/unseal round-trip: export GMS, re-import under a new name, verify key agreement. +#[test] +async fn group_export_import_key_agreement() -> Result<()> { + let mut ctx = TestContext::new().await?; + + // Create a group + let command = Command::Group(GroupSubcommand::New { + name: "alice-group".to_string(), + }); + wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; + + // Export the GMS + let holder = ctx + .wallet() + .storage() + .user_data + .group_key_holder("alice-group") + .context("Group not found")?; + let gms_hex = hex::encode(holder.dangerous_raw_gms()); + + // Import under a different name (simulating Bob receiving the GMS) + let command = Command::Group(GroupSubcommand::Import { + name: "bob-copy".to_string(), + gms: gms_hex, + }); + wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; + + // Both derive the same keys for the same tag + let alice_holder = ctx + .wallet() + .storage() + .user_data + .group_key_holder("alice-group") + .unwrap(); + let bob_holder = ctx + .wallet() + .storage() + .user_data + .group_key_holder("bob-copy") + .unwrap(); + + let tag = [42_u8; 32]; + let alice_npk = alice_holder + .derive_keys_for_shared_account(&tag) + .generate_nullifier_public_key(); + let bob_npk = bob_holder + .derive_keys_for_shared_account(&tag) + .generate_nullifier_public_key(); + + assert_eq!( + alice_npk, bob_npk, + "Key agreement: same GMS produces same keys" + ); + + info!("Key agreement verified"); + Ok(()) +} + +/// Fund a shared account from a public account via auth-transfer, then sync. +#[test] +async fn fund_shared_account_from_public() -> Result<()> { + let mut ctx = TestContext::new().await?; + + // Create group and shared account + let command = Command::Group(GroupSubcommand::New { + name: "fund-group".to_string(), + }); + wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; + + let command = Command::Account(AccountSubcommand::New(NewSubcommand::Private { + cci: None, + label: None, + for_gms: Some("fund-group".to_string()), + pda: false, + seed: None, + program_id: None, + })); + let result = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; + let SubcommandReturnValue::RegisterAccount { + account_id: shared_id, + } = result + else { + anyhow::bail!("Expected RegisterAccount return value"); + }; + + // Initialize the shared account under auth-transfer + let command = Command::AuthTransfer(AuthTransferSubcommand::Init { + account_id: Some(format!("Private/{shared_id}")), + account_label: None, + }); + wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; + + tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; + + // Fund from a public account + let from_public = ctx.existing_public_accounts()[0]; + let command = Command::AuthTransfer(AuthTransferSubcommand::Send { + from: Some(format_public_account_id(from_public)), + from_label: None, + to: Some(format!("Private/{shared_id}")), + to_label: None, + to_npk: None, + to_vpk: None, + to_identifier: None, + amount: 100, + }); + wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; + + tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; + + // Sync private accounts + let command = Command::Account(AccountSubcommand::SyncPrivate); + wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; + + // Verify the shared account was updated + let entry = ctx + .wallet() + .storage() + .user_data + .shared_private_accounts + .get(&shared_id) + .context("Shared account not found after sync")?; + + info!( + "Shared account balance after funding: {}", + entry.account.balance + ); + assert_eq!( + entry.account.balance, 100, + "Shared account should have received 100" + ); + + Ok(()) +} diff --git a/key_protocol/src/key_management/group_key_holder.rs b/key_protocol/src/key_management/group_key_holder.rs index e6634f88..533906a1 100644 --- a/key_protocol/src/key_management/group_key_holder.rs +++ b/key_protocol/src/key_management/group_key_holder.rs @@ -2,7 +2,7 @@ use aes_gcm::{Aes256Gcm, KeyInit as _, aead::Aead as _}; use nssa_core::{ SharedSecretKey, encryption::{Scalar, shared_key_derivation::Secp256k1Point}, - program::PdaSeed, + program::{PdaSeed, ProgramId}, }; use rand::{RngCore as _, rngs::OsRng}; use serde::{Deserialize, Serialize}; @@ -83,40 +83,51 @@ impl GroupKeyHolder { /// Derive a per-PDA [`SecretSpendingKey`] by mixing the seed into the SHA-256 input. /// - /// Each distinct `pda_seed` produces a distinct SSK in the full 256-bit space, so - /// adversarial seed-grinding cannot collide two PDAs' derived keys under the same + /// Each distinct `(program_id, pda_seed)` pair produces a distinct SSK in the full 256-bit + /// space, so adversarial seed-grinding cannot collide two PDAs' derived keys under the same /// group. Uses the codebase's 32-byte protocol-versioned domain-separation convention. - fn secret_spending_key_for_pda(&self, pda_seed: &PdaSeed) -> SecretSpendingKey { + fn secret_spending_key_for_pda( + &self, + program_id: &ProgramId, + pda_seed: &PdaSeed, + ) -> SecretSpendingKey { const PREFIX: &[u8; 32] = b"/LEE/v0.3/GroupKeyDerivation/SSK"; let mut hasher = sha2::Sha256::new(); hasher.update(PREFIX); hasher.update(self.gms); + for word in program_id { + hasher.update(word.to_le_bytes()); + } hasher.update(pda_seed.as_ref()); SecretSpendingKey(hasher.finalize_fixed().into()) } - /// Derive keys for a specific PDA. + /// Derive keys for a specific PDA under a given program. /// /// All controllers holding the same GMS independently derive the same keys for the - /// same PDA because the derivation is deterministic in (GMS, seed). + /// same `(program_id, seed)` because the derivation is deterministic. #[must_use] - pub fn derive_keys_for_pda(&self, pda_seed: &PdaSeed) -> PrivateKeyHolder { - self.secret_spending_key_for_pda(pda_seed) + pub fn derive_keys_for_pda( + &self, + program_id: &ProgramId, + pda_seed: &PdaSeed, + ) -> PrivateKeyHolder { + self.secret_spending_key_for_pda(program_id, pda_seed) .produce_private_key_holder(None) } /// Derive keys for a shared regular (non-PDA) private account. /// /// Uses a distinct domain separator from `derive_keys_for_pda` to prevent cross-domain - /// key collisions. The `tag` should be a stable, unique 32-byte value (e.g. derived from - /// a random identifier at account creation time). + /// key collisions. The `derivation_seed` should be a stable, unique 32-byte value + /// (e.g. derived deterministically from the account's identifier). #[must_use] - pub fn derive_keys_for_shared_account(&self, tag: &[u8; 32]) -> PrivateKeyHolder { + pub fn derive_keys_for_shared_account(&self, derivation_seed: &[u8; 32]) -> PrivateKeyHolder { const PREFIX: &[u8; 32] = b"/LEE/v0.3/GroupKeyDerivation/SHA"; let mut hasher = sha2::Sha256::new(); hasher.update(PREFIX); hasher.update(self.gms); - hasher.update(tag); + hasher.update(derivation_seed); SecretSpendingKey(hasher.finalize_fixed().into()).produce_private_key_holder(None) } @@ -210,6 +221,8 @@ mod tests { use super::*; + const TEST_PROGRAM_ID: ProgramId = [9; 8]; + /// Two holders from the same GMS derive identical keys for the same PDA seed. #[test] fn same_gms_same_seed_produces_same_keys() { @@ -218,8 +231,8 @@ mod tests { let holder_b = GroupKeyHolder::from_gms(gms); let seed = PdaSeed::new([1; 32]); - let keys_a = holder_a.derive_keys_for_pda(&seed); - let keys_b = holder_b.derive_keys_for_pda(&seed); + let keys_a = holder_a.derive_keys_for_pda(&TEST_PROGRAM_ID, &seed); + let keys_b = holder_b.derive_keys_for_pda(&TEST_PROGRAM_ID, &seed); assert_eq!( keys_a.generate_nullifier_public_key().to_byte_array(), @@ -235,10 +248,10 @@ mod tests { let seed_b = PdaSeed::new([2; 32]); let npk_a = holder - .derive_keys_for_pda(&seed_a) + .derive_keys_for_pda(&TEST_PROGRAM_ID, &seed_a) .generate_nullifier_public_key(); let npk_b = holder - .derive_keys_for_pda(&seed_b) + .derive_keys_for_pda(&TEST_PROGRAM_ID, &seed_b) .generate_nullifier_public_key(); assert_ne!(npk_a.to_byte_array(), npk_b.to_byte_array()); @@ -252,10 +265,10 @@ mod tests { let seed = PdaSeed::new([1; 32]); let npk_a = holder_a - .derive_keys_for_pda(&seed) + .derive_keys_for_pda(&TEST_PROGRAM_ID, &seed) .generate_nullifier_public_key(); let npk_b = holder_b - .derive_keys_for_pda(&seed) + .derive_keys_for_pda(&TEST_PROGRAM_ID, &seed) .generate_nullifier_public_key(); assert_ne!(npk_a.to_byte_array(), npk_b.to_byte_array()); @@ -269,10 +282,10 @@ mod tests { let seed = PdaSeed::new([1; 32]); let npk_original = original - .derive_keys_for_pda(&seed) + .derive_keys_for_pda(&TEST_PROGRAM_ID, &seed) .generate_nullifier_public_key(); let npk_restored = restored - .derive_keys_for_pda(&seed) + .derive_keys_for_pda(&TEST_PROGRAM_ID, &seed) .generate_nullifier_public_key(); assert_eq!(npk_original.to_byte_array(), npk_restored.to_byte_array()); @@ -284,7 +297,7 @@ mod tests { let holder = GroupKeyHolder::from_gms([42_u8; 32]); let seed = PdaSeed::new([1; 32]); let npk = holder - .derive_keys_for_pda(&seed) + .derive_keys_for_pda(&TEST_PROGRAM_ID, &seed) .generate_nullifier_public_key(); assert_ne!(npk, NullifierPublicKey([0; 32])); @@ -304,7 +317,7 @@ mod tests { let holder = GroupKeyHolder::from_gms(gms); let npk = holder - .derive_keys_for_pda(&seed) + .derive_keys_for_pda(&TEST_PROGRAM_ID, &seed) .generate_nullifier_public_key(); let account_id = AccountId::for_private_pda(&program_id, &seed, &npk); @@ -333,10 +346,10 @@ mod tests { let seed = PdaSeed::new([1; 32]); let npk_original = original - .derive_keys_for_pda(&seed) + .derive_keys_for_pda(&TEST_PROGRAM_ID, &seed) .generate_nullifier_public_key(); let npk_restored = restored - .derive_keys_for_pda(&seed) + .derive_keys_for_pda(&TEST_PROGRAM_ID, &seed) .generate_nullifier_public_key(); assert_eq!(npk_original, npk_restored); @@ -354,7 +367,7 @@ mod tests { let seed = PdaSeed::new([5; 32]); let group_npk = GroupKeyHolder::from_gms(shared_bytes) - .derive_keys_for_pda(&seed) + .derive_keys_for_pda(&TEST_PROGRAM_ID, &seed) .generate_nullifier_public_key(); let personal_npk = SecretSpendingKey(shared_bytes) @@ -382,10 +395,10 @@ mod tests { let seed = PdaSeed::new([1; 32]); assert_eq!( holder - .derive_keys_for_pda(&seed) + .derive_keys_for_pda(&TEST_PROGRAM_ID, &seed) .generate_nullifier_public_key(), restored - .derive_keys_for_pda(&seed) + .derive_keys_for_pda(&TEST_PROGRAM_ID, &seed) .generate_nullifier_public_key(), ); } @@ -468,7 +481,7 @@ mod tests { .iter() .map(|gms| { GroupKeyHolder::from_gms(*gms) - .derive_keys_for_pda(&seed) + .derive_keys_for_pda(&TEST_PROGRAM_ID, &seed) .generate_nullifier_public_key() }) .collect(); @@ -493,7 +506,7 @@ mod tests { let program_id: nssa_core::program::ProgramId = [1; 8]; // Derive Alice's keys - let alice_keys = alice_holder.derive_keys_for_pda(&pda_seed); + let alice_keys = alice_holder.derive_keys_for_pda(&TEST_PROGRAM_ID, &pda_seed); let alice_npk = alice_keys.generate_nullifier_public_key(); // Seal GMS for Bob using Bob's viewing key, Bob unseals @@ -508,7 +521,7 @@ mod tests { // Key agreement: both derive identical NPK and AccountId let bob_npk = bob_holder - .derive_keys_for_pda(&pda_seed) + .derive_keys_for_pda(&TEST_PROGRAM_ID, &pda_seed) .generate_nullifier_public_key(); assert_eq!(alice_npk, bob_npk); @@ -517,27 +530,27 @@ mod tests { assert_eq!(alice_account_id, bob_account_id); } - /// Same GMS + same tag produces same keys for shared accounts. + /// Same GMS + same derivation seed produces same keys for shared accounts. #[test] - fn shared_account_same_gms_same_tag_produces_same_keys() { + fn shared_account_same_gms_same_seed_produces_same_keys() { let gms = [42_u8; 32]; - let tag = [1_u8; 32]; + let derivation_seed = [1_u8; 32]; let holder_a = GroupKeyHolder::from_gms(gms); let holder_b = GroupKeyHolder::from_gms(gms); let npk_a = holder_a - .derive_keys_for_shared_account(&tag) + .derive_keys_for_shared_account(&derivation_seed) .generate_nullifier_public_key(); let npk_b = holder_b - .derive_keys_for_shared_account(&tag) + .derive_keys_for_shared_account(&derivation_seed) .generate_nullifier_public_key(); assert_eq!(npk_a, npk_b); } - /// Different tags produce different keys for shared accounts. + /// Different derivation seeds produce different keys for shared accounts. #[test] - fn shared_account_different_tags_produce_different_keys() { + fn shared_account_different_seeds_produce_different_keys() { let holder = GroupKeyHolder::from_gms([42_u8; 32]); let npk_a = holder .derive_keys_for_shared_account(&[1_u8; 32]) @@ -556,7 +569,7 @@ mod tests { let bytes = [1_u8; 32]; let pda_npk = holder - .derive_keys_for_pda(&PdaSeed::new(bytes)) + .derive_keys_for_pda(&TEST_PROGRAM_ID, &PdaSeed::new(bytes)) .generate_nullifier_public_key(); let shared_npk = holder .derive_keys_for_shared_account(&bytes) diff --git a/key_protocol/src/key_protocol_core/mod.rs b/key_protocol/src/key_protocol_core/mod.rs index ea8d8405..a18c1d3a 100644 --- a/key_protocol/src/key_protocol_core/mod.rs +++ b/key_protocol/src/key_protocol_core/mod.rs @@ -27,10 +27,12 @@ pub struct UserPrivateAccountData { pub struct SharedAccountEntry { pub group_label: String, pub identifier: Identifier, - /// For PDA accounts, the seed used to derive keys via `derive_keys_for_pda`. - /// `None` for regular shared accounts (keys derived from identifier via tag). + /// For PDA accounts, the seed and program ID used to derive keys via `derive_keys_for_pda`. + /// `None` for regular shared accounts (keys derived from identifier via derivation seed). #[serde(default)] pub pda_seed: Option, + #[serde(default)] + pub pda_program_id: Option, pub account: Account, } @@ -55,7 +57,7 @@ pub struct NSSAUserData { /// Old wallet files with `pda_accounts` (plain Account values) are incompatible with /// this type. The `default` attribute ensures they deserialize as empty rather than failing. #[serde(default)] - pub shared_accounts: BTreeMap, + pub shared_private_accounts: BTreeMap, } impl NSSAUserData { @@ -115,7 +117,7 @@ impl NSSAUserData { public_key_tree, private_key_tree, group_key_holders: BTreeMap::new(), - shared_accounts: BTreeMap::new(), + shared_private_accounts: BTreeMap::new(), }) } @@ -274,7 +276,7 @@ mod tests { fn group_key_holders_default_empty() { let user_data = NSSAUserData::default(); assert!(user_data.group_key_holders.is_empty()); - assert!(user_data.shared_accounts.is_empty()); + assert!(user_data.shared_private_accounts.is_empty()); } #[test] @@ -285,6 +287,7 @@ mod tests { group_label: String::from("test-group"), identifier: 42, pda_seed: None, + pda_program_id: None, account: nssa_core::account::Account::default(), }; let encoded = bincode::serialize(&entry).expect("serialize"); @@ -297,6 +300,7 @@ mod tests { group_label: String::from("pda-group"), identifier: u128::MAX, pda_seed: Some(PdaSeed::new([7_u8; 32])), + pda_program_id: Some([9; 8]), account: nssa_core::account::Account::default(), }; let pda_encoded = bincode::serialize(&pda_entry).expect("serialize pda"); @@ -315,6 +319,7 @@ mod tests { group_label: String::from("old"), identifier: 1, pda_seed: None, + pda_program_id: None, account: nssa_core::account::Account::default(), }; let encoded = bincode::serialize(&entry).expect("serialize"); @@ -345,8 +350,8 @@ mod tests { // PDA shared account: derive via seed let seed = PdaSeed::new([2_u8; 32]); - let pda_keys_a = holder.derive_keys_for_pda(&seed); - let pda_keys_b = holder.derive_keys_for_pda(&seed); + let pda_keys_a = holder.derive_keys_for_pda(&[9; 8], &seed); + let pda_keys_b = holder.derive_keys_for_pda(&[9; 8], &seed); assert_eq!( pda_keys_a.generate_nullifier_public_key(), pda_keys_b.generate_nullifier_public_key(), diff --git a/wallet/src/cli/account.rs b/wallet/src/cli/account.rs index 3bb7310b..cfb1c8e4 100644 --- a/wallet/src/cli/account.rs +++ b/wallet/src/cli/account.rs @@ -176,15 +176,7 @@ impl WalletSubcommand for NewSubcommand { } if let Some(group_name) = for_gms { - // GMS-derived account - let holder = wallet_core - .storage() - .user_data - .group_key_holder(&group_name) - .context(format!("Group '{group_name}' not found"))?; - - if pda { - // PDA shared account + let info = if pda { let seed_hex = seed.context("--seed is required for PDA accounts")?; let pid_hex = program_id.context("--program-id is required for PDA accounts")?; @@ -204,73 +196,27 @@ impl WalletSubcommand for NewSubcommand { pid[i] = u32::from_le_bytes(chunk.try_into().unwrap()); } - let keys = holder.derive_keys_for_pda(&pda_seed); - let npk = keys.generate_nullifier_public_key(); - let vpk = keys.generate_viewing_public_key(); - let account_id = nssa::AccountId::for_private_pda(&pid, &pda_seed, &npk); - - if let Some(label) = label { - wallet_core - .storage - .labels - .insert(account_id.to_string(), Label::new(label)); - } - - wallet_core.register_shared_account( - account_id, - group_name.clone(), - u128::MAX, - Some(pda_seed), - ); - - println!("PDA shared account from group '{group_name}'"); - println!("AccountId: {account_id}"); - println!("NPK: {}", hex::encode(npk.0)); - println!("VPK: {}", hex::encode(&vpk.0)); - - wallet_core.store_persistent_data().await?; - Ok(SubcommandReturnValue::RegisterAccount { account_id }) + wallet_core.create_shared_pda_account(&group_name, pda_seed, pid)? } else { - // Regular shared account. The tag is derived deterministically - // from the identifier so that keys can be re-derived without - // storing the tag separately. - let identifier: nssa_core::Identifier = rand::random(); - let tag = { - 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 - }; + wallet_core.create_shared_regular_account(&group_name)? + }; - let keys = holder.derive_keys_for_shared_account(&tag); - let npk = keys.generate_nullifier_public_key(); - let vpk = keys.generate_viewing_public_key(); - let account_id = nssa::AccountId::from((&npk, identifier)); - - if let Some(label) = label { - wallet_core - .storage - .labels - .insert(account_id.to_string(), Label::new(label)); - } - - wallet_core.register_shared_account( - account_id, - group_name.clone(), - identifier, - None, - ); - - println!("Shared account from group '{group_name}'"); - println!("AccountId: Private/{account_id}"); - println!("NPK: {}", hex::encode(npk.0)); - println!("VPK: {}", hex::encode(&vpk.0)); - - wallet_core.store_persistent_data().await?; - Ok(SubcommandReturnValue::RegisterAccount { account_id }) + if let Some(label) = label { + wallet_core + .storage + .labels + .insert(info.account_id.to_string(), Label::new(label)); } + + println!("Shared account from group '{group_name}'"); + println!("AccountId: Private/{}", info.account_id); + println!("NPK: {}", hex::encode(info.npk.0)); + println!("VPK: {}", hex::encode(&info.vpk.0)); + + wallet_core.store_persistent_data().await?; + Ok(SubcommandReturnValue::RegisterAccount { + account_id: info.account_id, + }) } else { // Standard wallet-tree-derived account let (account_id, chain_index) = wallet_core.create_new_account_private(cci); diff --git a/wallet/src/config.rs b/wallet/src/config.rs index bbd98ac7..d8e186bd 100644 --- a/wallet/src/config.rs +++ b/wallet/src/config.rs @@ -98,6 +98,18 @@ pub struct PersistentStorage { /// "2rnKprXqWGWJTkDZKsQbFXa4ctKRbapsdoTKQFnaVGG8"). #[serde(default)] pub labels: HashMap, + /// Group key holders for shared account management. + #[serde(default)] + pub group_key_holders: std::collections::BTreeMap< + String, + key_protocol::key_management::group_key_holder::GroupKeyHolder, + >, + /// Cached state of shared private accounts (PDA and regular). + #[serde(default)] + pub shared_private_accounts: std::collections::BTreeMap< + nssa::AccountId, + key_protocol::key_protocol_core::SharedAccountEntry, + >, } impl PersistentStorage { diff --git a/wallet/src/helperfunctions.rs b/wallet/src/helperfunctions.rs index 94755f6e..57416c55 100644 --- a/wallet/src/helperfunctions.rs +++ b/wallet/src/helperfunctions.rs @@ -204,6 +204,8 @@ pub fn produce_data_for_storage( accounts: vec_for_storage, last_synced_block, labels, + group_key_holders: user_data.group_key_holders.clone(), + shared_private_accounts: user_data.shared_private_accounts.clone(), } } diff --git a/wallet/src/lib.rs b/wallet/src/lib.rs index f179ec44..545b704b 100644 --- a/wallet/src/lib.rs +++ b/wallet/src/lib.rs @@ -51,6 +51,13 @@ pub enum AccDecodeData { Decode(nssa_core::SharedSecretKey, AccountId), } +/// Info returned when creating a shared account. +pub struct SharedAccountInfo { + pub account_id: AccountId, + pub npk: nssa_core::NullifierPublicKey, + pub vpk: nssa_core::encryption::ViewingPublicKey, +} + #[derive(Debug, thiserror::Error)] pub enum ExecutionFailureKind { #[error("Failed to get data from sequencer")] @@ -98,6 +105,8 @@ impl WalletCore { accounts: persistent_accounts, last_synced_block, labels, + group_key_holders, + shared_private_accounts, } = PersistentStorage::from_path(&storage_path).with_context(|| { format!( "Failed to read persistent storage at {}", @@ -109,7 +118,12 @@ impl WalletCore { config_path, storage_path, config_overrides, - |config| WalletChainStore::new(config, persistent_accounts, labels), + |config| { + let mut store = WalletChainStore::new(config, persistent_accounts, labels)?; + store.user_data.group_key_holders = group_key_holders; + store.user_data.shared_private_accounts = shared_private_accounts; + Ok(store) + }, last_synced_block, ) } @@ -305,25 +319,93 @@ impl WalletCore { } /// Register a shared account in storage for sync tracking. - pub fn register_shared_account( + fn register_shared_account( &mut self, account_id: AccountId, group_label: String, identifier: nssa_core::Identifier, pda_seed: Option, + pda_program_id: Option, ) { use key_protocol::key_protocol_core::SharedAccountEntry; - self.storage.user_data.shared_accounts.insert( + self.storage.user_data.shared_private_accounts.insert( account_id, SharedAccountEntry { group_label, identifier, pda_seed, + pda_program_id, account: Account::default(), }, ); } + /// Create a shared PDA account from a group's GMS. Returns the `AccountId` and derived keys. + pub fn create_shared_pda_account( + &mut self, + group_name: &str, + pda_seed: nssa_core::program::PdaSeed, + program_id: nssa_core::program::ProgramId, + ) -> Result { + let holder = self + .storage + .user_data + .group_key_holder(group_name) + .context(format!("Group '{group_name}' not found"))?; + + 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); + + self.register_shared_account( + account_id, + String::from(group_name), + u128::MAX, + Some(pda_seed), + Some(program_id), + ); + + Ok(SharedAccountInfo { + account_id, + npk, + vpk, + }) + } + + /// 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(&mut self, group_name: &str) -> Result { + let identifier: nssa_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 + .user_data + .group_key_holder(group_name) + .context(format!("Group '{group_name}' not found"))?; + + 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, identifier)); + + self.register_shared_account(account_id, String::from(group_name), identifier, None, None); + + Ok(SharedAccountInfo { + account_id, + npk, + vpk, + }) + } + /// Get account balance. pub async fn get_account_balance(&self, acc: AccountId) -> Result { Ok(self.sequencer_client.get_account_balance(acc).await?) @@ -596,14 +678,14 @@ impl WalletCore { } // Scan for updates to shared accounts (GMS-derived). - self.sync_shared_accounts_with_tx(&tx); + self.sync_shared_private_accounts_with_tx(&tx); } - fn sync_shared_accounts_with_tx(&mut self, tx: &PrivacyPreservingTransaction) { + fn sync_shared_private_accounts_with_tx(&mut self, tx: &PrivacyPreservingTransaction) { let shared_keys: Vec<_> = self .storage .user_data - .shared_accounts + .shared_private_accounts .iter() .filter_map(|(&account_id, entry)| { let holder = self @@ -612,9 +694,13 @@ impl WalletCore { .group_key_holders .get(&entry.group_label)?; - let keys = entry.pda_seed.as_ref().map_or_else( - || { - let tag = { + let keys = match (&entry.pda_seed, &entry.pda_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"); @@ -622,10 +708,9 @@ impl WalletCore { let result: [u8; 32] = hasher.finalize().into(); result }; - holder.derive_keys_for_shared_account(&tag) - }, - |pda_seed| holder.derive_keys_for_pda(pda_seed), - ); + holder.derive_keys_for_shared_account(&derivation_seed) + } + }; let npk = keys.generate_nullifier_public_key(); let vpk = keys.generate_viewing_public_key(); let vsk = keys.viewing_secret_key; @@ -658,7 +743,11 @@ impl WalletCore { .expect("Ciphertext ID is expected to fit in u32"), ) { info!("Synced shared account {account_id:#?} with new state {new_acc:#?}"); - if let Some(entry) = self.storage.user_data.shared_accounts.get_mut(&account_id) + if let Some(entry) = self + .storage + .user_data + .shared_private_accounts + .get_mut(&account_id) { entry.account = new_acc; } diff --git a/wallet/src/privacy_preserving_tx.rs b/wallet/src/privacy_preserving_tx.rs index dfac8180..cdf0bed7 100644 --- a/wallet/src/privacy_preserving_tx.rs +++ b/wallet/src/privacy_preserving_tx.rs @@ -335,7 +335,7 @@ async fn private_pda_preparation( let acc = wallet .storage .user_data - .shared_accounts + .shared_private_accounts .get(&account_id) .map(|e| e.account.clone()) .unwrap_or_default(); @@ -386,7 +386,7 @@ async fn private_shared_preparation( let acc = wallet .storage .user_data - .shared_accounts + .shared_private_accounts .get(&account_id) .map(|e| e.account.clone()) .unwrap_or_default();