use privateaccountkind in storage and fix circuit

This commit is contained in:
Sergio Chouhy 2026-05-04 21:40:30 -03:00
parent 7d5e1492c4
commit 11949e9fa1
12 changed files with 114 additions and 82 deletions

View File

@ -1,5 +1,5 @@
use k256::{Scalar, elliptic_curve::PrimeField as _};
use nssa_core::{Identifier, NullifierPublicKey, encryption::ViewingPublicKey};
use nssa_core::{NullifierPublicKey, PrivateAccountKind, encryption::ViewingPublicKey};
use serde::{Deserialize, Serialize};
use crate::key_management::{
@ -10,7 +10,7 @@ use crate::key_management::{
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct ChildKeysPrivate {
pub value: (KeyChain, Vec<(Identifier, nssa::Account)>),
pub value: (KeyChain, Vec<(PrivateAccountKind, nssa::Account)>),
pub ccc: [u8; 32],
/// Can be [`None`] if root.
pub cci: Option<u32>,
@ -115,9 +115,8 @@ impl KeyTreeNode for ChildKeysPrivate {
}
fn account_ids(&self) -> impl Iterator<Item = nssa::AccountId> {
self.value.1.iter().map(|(identifier, _)| {
nssa::AccountId::from((&self.value.0.nullifier_public_key, *identifier))
})
let npk = self.value.0.nullifier_public_key;
self.value.1.iter().map(move |(kind, _)| nssa::AccountId::for_private_account(&npk, kind))
}
}

View File

@ -319,6 +319,7 @@ mod tests {
use std::{collections::HashSet, str::FromStr as _};
use nssa::AccountId;
use nssa_core::PrivateAccountKind;
use super::*;
@ -532,7 +533,7 @@ mod tests {
.get_mut(&ChainIndex::from_str("/1").unwrap())
.unwrap();
acc.value.1.push((
0,
PrivateAccountKind::Regular(0),
nssa::Account {
balance: 2,
..nssa::Account::default()
@ -544,7 +545,7 @@ mod tests {
.get_mut(&ChainIndex::from_str("/2").unwrap())
.unwrap();
acc.value.1.push((
0,
PrivateAccountKind::Regular(0),
nssa::Account {
balance: 3,
..nssa::Account::default()
@ -556,7 +557,7 @@ mod tests {
.get_mut(&ChainIndex::from_str("/0/1").unwrap())
.unwrap();
acc.value.1.push((
0,
PrivateAccountKind::Regular(0),
nssa::Account {
balance: 5,
..nssa::Account::default()
@ -568,7 +569,7 @@ mod tests {
.get_mut(&ChainIndex::from_str("/1/0").unwrap())
.unwrap();
acc.value.1.push((
0,
PrivateAccountKind::Regular(0),
nssa::Account {
balance: 6,
..nssa::Account::default()

View File

@ -3,7 +3,7 @@ use std::collections::BTreeMap;
use anyhow::Result;
use k256::AffinePoint;
use nssa::{Account, AccountId};
use nssa_core::Identifier;
use nssa_core::{Identifier, PrivateAccountKind};
use serde::{Deserialize, Serialize};
use crate::key_management::{
@ -17,7 +17,7 @@ pub type PublicKey = AffinePoint;
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct UserPrivateAccountData {
pub key_chain: KeyChain,
pub accounts: Vec<(Identifier, Account)>,
pub accounts: Vec<(PrivateAccountKind, Account)>,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
@ -53,9 +53,9 @@ impl NSSAUserData {
) -> bool {
let mut check_res = true;
for (account_id, entry) in accounts_keys_map {
let any_match = entry.accounts.iter().any(|(identifier, _)| {
nssa::AccountId::from((&entry.key_chain.nullifier_public_key, *identifier))
== *account_id
let npk = &entry.key_chain.nullifier_public_key;
let any_match = entry.accounts.iter().any(|(kind, _)| {
nssa::AccountId::for_private_account(npk, kind) == *account_id
});
if !any_match {
println!("No matching entry found for account_id {account_id}");
@ -155,24 +155,22 @@ impl NSSAUserData {
) -> Option<(KeyChain, nssa_core::account::Account, Identifier)> {
// Check default accounts
if let Some(entry) = self.default_user_private_accounts.get(&account_id) {
for (identifier, account) in &entry.accounts {
let expected_id =
nssa::AccountId::from((&entry.key_chain.nullifier_public_key, *identifier));
if expected_id == account_id {
return Some((entry.key_chain.clone(), account.clone(), *identifier));
}
let npk = &entry.key_chain.nullifier_public_key;
if let Some((kind, account)) =
entry.accounts.iter().find(|(kind, _)| nssa::AccountId::for_private_account(npk, kind) == account_id)
{
return Some((entry.key_chain.clone(), account.clone(), kind.identifier()));
}
return None;
}
// Check tree
if let Some(node) = self.private_key_tree.get_node(account_id) {
let key_chain = &node.value.0;
for (identifier, account) in &node.value.1 {
let expected_id =
nssa::AccountId::from((&key_chain.nullifier_public_key, *identifier));
if expected_id == account_id {
return Some((key_chain.clone(), account.clone(), *identifier));
}
let npk = &key_chain.nullifier_public_key;
if let Some((kind, account)) =
node.value.1.iter().find(|(kind, _)| nssa::AccountId::for_private_account(npk, kind) == account_id)
{
return Some((key_chain.clone(), account.clone(), kind.identifier()));
}
}
None

View File

@ -22,9 +22,9 @@ pub type Scalar = [u8; 32];
/// to reconstruct the account's [`AccountId`] on the receiver side.
///
/// [`AccountId`]: crate::account::AccountId
#[derive(Debug, Clone, PartialEq, Eq)]
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum PrivateAccountKind {
Account(Identifier),
Regular(Identifier),
Pda {
program_id: ProgramId,
seed: PdaSeed,
@ -33,14 +33,14 @@ pub enum PrivateAccountKind {
}
impl PrivateAccountKind {
/// Account(ident): 0x00 || ident (16 LE) || [0u8; 64]
/// Regular(ident): 0x00 || ident (16 LE) || [0u8; 64]
/// Pda { program_id, seed, ident }: 0x01 || program_id (32 LE) || seed (32) || ident (16 LE)
pub const HEADER_LEN: usize = 81;
#[must_use]
pub fn identifier(&self) -> Identifier {
match self {
Self::Account(identifier) => *identifier,
Self::Regular(identifier) => *identifier,
Self::Pda { identifier, .. } => *identifier,
}
}
@ -49,7 +49,7 @@ impl PrivateAccountKind {
pub fn to_header_bytes(&self) -> [u8; Self::HEADER_LEN] {
let mut bytes = [0u8; Self::HEADER_LEN];
match self {
Self::Account(identifier) => {
Self::Regular(identifier) => {
bytes[0] = 0x00;
bytes[1..17].copy_from_slice(&identifier.to_le_bytes());
// bytes[17..81] are zero padding
@ -72,7 +72,7 @@ impl PrivateAccountKind {
match bytes[0] {
0x00 => {
let identifier = Identifier::from_le_bytes(bytes[1..17].try_into().unwrap());
Some(Self::Account(identifier))
Some(Self::Regular(identifier))
}
0x01 => {
let mut program_id = [0u32; 8];
@ -208,7 +208,7 @@ mod tests {
let account_ct = EncryptionScheme::encrypt(
&account,
&PrivateAccountKind::Account(42),
&PrivateAccountKind::Regular(42),
&secret,
&commitment,
0,

View File

@ -8,6 +8,7 @@ use serde::{Deserialize, Serialize};
use crate::{
BlockId, Identifier, NullifierPublicKey, Timestamp,
account::{Account, AccountId, AccountWithMetadata},
encryption::PrivateAccountKind,
};
pub const DEFAULT_PROGRAM_ID: ProgramId = [0; 8];
@ -96,6 +97,17 @@ impl AccountId {
.expect("Hash output must be exactly 32 bytes long"),
)
}
/// Derives the [`AccountId`] for a private account from the nullifier public key and kind.
#[must_use]
pub fn for_private_account(npk: &NullifierPublicKey, kind: &PrivateAccountKind) -> Self {
match kind {
PrivateAccountKind::Regular(identifier) => Self::from((npk, *identifier)),
PrivateAccountKind::Pda { program_id, seed, identifier } => {
Self::for_private_pda(program_id, seed, npk, *identifier)
}
}
}
}
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq)]

View File

@ -253,7 +253,7 @@ pub mod tests {
let esk = [3; 32];
let shared_secret = SharedSecretKey::new(&esk, &vpk);
let epk = EphemeralPublicKey::from_scalar(esk);
let ciphertext = EncryptionScheme::encrypt(&account, &PrivateAccountKind::Account(0), &shared_secret, &commitment, 2);
let ciphertext = EncryptionScheme::encrypt(&account, &PrivateAccountKind::Regular(0), &shared_secret, &commitment, 2);
let encrypted_account_data =
EncryptedAccountData::new(ciphertext.clone(), &npk, &vpk, epk.clone());

View File

@ -45,11 +45,11 @@ struct ExecutionState {
/// `AccountId` entry or as an equality check against the existing one, making the rule: one
/// `(program, seed)` → one account per tx.
pda_family_binding: HashMap<(ProgramId, PdaSeed), AccountId>,
/// Map from a mask-3 `pre_state`'s position in `visibility_mask` to the npk supplied for
/// that position in `private_account_keys`. Built once in `derive_from_outputs` by walking
/// `visibility_mask` in lock-step with `private_account_keys`, used later by the claim and
/// caller-seeds authorization paths.
private_pda_npk_by_position: HashMap<usize, NullifierPublicKey>,
/// Map from a mask-3 `pre_state`'s position in `visibility_mask` to the (npk, identifier)
/// supplied for that position in `private_account_keys`. Built once in `derive_from_outputs`
/// by walking `visibility_mask` in lock-step with `private_account_keys`, used later by the
/// claim and caller-seeds authorization paths.
private_pda_npk_by_position: HashMap<usize, (NullifierPublicKey, Identifier)>,
}
impl ExecutionState {
@ -64,18 +64,18 @@ impl ExecutionState {
// pre_state order across all masks 1/2/3, so walk `visibility_mask` in lock-step. The
// downstream `compute_circuit_output` also consumes the same iterator and its trailing
// assertions catch an over-supply of keys; under-supply surfaces here.
let mut private_pda_npk_by_position: HashMap<usize, NullifierPublicKey> = HashMap::new();
let mut private_pda_npk_by_position: HashMap<usize, (NullifierPublicKey, Identifier)> = HashMap::new();
{
let mut keys_iter = private_account_keys.iter();
for (pos, &mask) in visibility_mask.iter().enumerate() {
if matches!(mask, 1..=3) {
let (npk, _, _) = keys_iter.next().unwrap_or_else(|| {
let (npk, identifier, _) = keys_iter.next().unwrap_or_else(|| {
panic!(
"private_account_keys shorter than visibility_mask demands: no key for masked position {pos} (mask {mask})"
)
});
if mask == 3 {
private_pda_npk_by_position.insert(pos, *npk);
private_pda_npk_by_position.insert(pos, (*npk, *identifier));
}
}
}
@ -363,11 +363,11 @@ impl ExecutionState {
);
}
Claim::Pda(seed) => {
let npk = self
let (npk, identifier) = self
.private_pda_npk_by_position
.get(&pre_state_position)
.expect("private PDA pre_state must have an npk in the position map");
let pda = AccountId::for_private_pda(&program_id, &seed, npk, u128::MAX);
let pda = AccountId::for_private_pda(&program_id, &seed, npk, *identifier);
assert_eq!(
pre_account_id, pda,
"Invalid private PDA claim for account {pre_account_id}"
@ -454,7 +454,7 @@ fn assert_family_binding(
fn resolve_authorization_and_record_bindings(
pda_family_binding: &mut HashMap<(ProgramId, PdaSeed), AccountId>,
private_pda_bound_positions: &mut HashMap<usize, (ProgramId, PdaSeed)>,
private_pda_npk_by_position: &HashMap<usize, NullifierPublicKey>,
private_pda_npk_by_position: &HashMap<usize, (NullifierPublicKey, Identifier)>,
pre_account_id: AccountId,
pre_state_position: usize,
caller_program_id: Option<ProgramId>,
@ -467,8 +467,8 @@ fn resolve_authorization_and_record_bindings(
if AccountId::for_public_pda(&caller, seed) == pre_account_id {
return Some((*seed, false, caller));
}
if let Some(npk) = private_pda_npk_by_position.get(&pre_state_position)
&& AccountId::for_private_pda(&caller, seed, npk, u128::MAX) == pre_account_id
if let Some((npk, identifier)) = private_pda_npk_by_position.get(&pre_state_position)
&& AccountId::for_private_pda(&caller, seed, npk, *identifier) == pre_account_id
{
return Some((*seed, true, caller));
}
@ -608,7 +608,7 @@ fn compute_circuit_output(
// Encrypt and push post state
let encrypted_account = EncryptionScheme::encrypt(
&post_with_updated_nonce,
&PrivateAccountKind::Account(*identifier),
&PrivateAccountKind::Regular(*identifier),
shared_secret,
&commitment_post,
output_index,

View File

@ -73,8 +73,8 @@ impl WalletChainStore {
PersistentAccountData::Private(data) => {
let npk = data.data.value.0.nullifier_public_key;
let chain_index = data.chain_index;
for identifier in &data.identifiers {
let account_id = nssa::AccountId::from((&npk, *identifier));
for kind in &data.kinds {
let account_id = nssa::AccountId::for_private_account(&npk, kind);
private_tree
.account_id_map
.insert(account_id, chain_index.clone());
@ -90,7 +90,7 @@ impl WalletChainStore {
data.account_id(),
UserPrivateAccountData {
key_chain: data.key_chain,
accounts: vec![(data.identifier, data.account)],
accounts: vec![(PrivateAccountKind::Regular(data.identifier), data.account)],
},
);
}
@ -136,7 +136,7 @@ impl WalletChainStore {
account_id,
UserPrivateAccountData {
key_chain: data.key_chain,
accounts: vec![(data.identifier, account)],
accounts: vec![(PrivateAccountKind::Regular(data.identifier), account)],
},
);
}
@ -195,7 +195,6 @@ impl WalletChainStore {
account: nssa_core::account::Account,
) {
debug!("inserting at address {account_id}, this account {account:?}");
let identifier = kind.identifier();
// Update default accounts if present
if let Entry::Occupied(mut entry) = self
@ -204,10 +203,10 @@ impl WalletChainStore {
.entry(account_id)
{
let entry = entry.get_mut();
if let Some((_, acc)) = entry.accounts.iter_mut().find(|(id, _)| *id == identifier) {
if let Some((_, acc)) = entry.accounts.iter_mut().find(|(k, _)| k == kind) {
*acc = account;
} else {
entry.accounts.push((identifier, account));
entry.accounts.push((kind.clone(), account));
}
return;
}
@ -230,29 +229,21 @@ impl WalletChainStore {
.key_map
.get_mut(&chain_index)
{
if let Some((_, acc)) = node.value.1.iter_mut().find(|(id, _)| *id == identifier) {
if let Some((_, acc)) = node.value.1.iter_mut().find(|(k, _)| k == kind) {
*acc = account;
} else {
node.value.1.push((identifier, account));
node.value.1.push((kind.clone(), account));
}
}
} else {
// Node not yet in account_id_map — find it by checking all nodes
for (ci, node) in &mut self.user_data.private_key_tree.key_map {
let npk = &node.value.0.nullifier_public_key;
let expected_id = match kind {
PrivateAccountKind::Account(id) => nssa::AccountId::from((npk, *id)),
PrivateAccountKind::Pda { program_id, seed, identifier: id } => {
nssa::AccountId::for_private_pda(program_id, seed, npk, *id)
}
};
if expected_id == account_id {
if let Some((_, acc)) =
node.value.1.iter_mut().find(|(id, _)| *id == identifier)
{
if nssa::AccountId::for_private_account(npk, kind) == account_id {
if let Some((_, acc)) = node.value.1.iter_mut().find(|(k, _)| k == kind) {
*acc = account;
} else {
node.value.1.push((identifier, account));
node.value.1.push((kind.clone(), account));
}
// Register in account_id_map
self.user_data
@ -298,7 +289,7 @@ mod tests {
data: public_data,
}),
PersistentAccountData::Private(Box::new(PersistentAccountDataPrivate {
identifiers: vec![],
kinds: vec![],
chain_index: ChainIndex::root(),
data: private_data,
})),

View File

@ -28,7 +28,7 @@ pub struct PersistentAccountDataPublic {
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PersistentAccountDataPrivate {
pub identifiers: Vec<nssa_core::Identifier>,
pub kinds: Vec<nssa_core::PrivateAccountKind>,
pub chain_index: ChainIndex,
pub data: ChildKeysPrivate,
}

View File

@ -166,10 +166,10 @@ pub fn produce_data_for_storage(
}
for (chain_index, node) in &user_data.private_key_tree.key_map {
let identifiers = node.value.1.iter().map(|(id, _)| *id).collect();
let kinds = node.value.1.iter().map(|(kind, _)| kind.clone()).collect();
vec_for_storage.push(
PersistentAccountDataPrivate {
identifiers,
kinds,
chain_index: chain_index.clone(),
data: node.clone(),
}
@ -188,12 +188,12 @@ pub fn produce_data_for_storage(
}
for entry in user_data.default_user_private_accounts.values() {
for (identifier, account) in &entry.accounts {
for (kind, account) in &entry.accounts {
vec_for_storage.push(
InitialAccountData::Private(Box::new(PrivateAccountPrivateInitialData {
account: account.clone(),
key_chain: entry.key_chain.clone(),
identifier: *identifier,
identifier: kind.identifier(),
}))
.into(),
);

View File

@ -284,7 +284,7 @@ impl WalletCore {
.nullifier_public_key;
let account_id = AccountId::from((&npk, identifier));
self.storage
.insert_private_account_data(account_id, &PrivateAccountKind::Account(identifier), Account::default());
.insert_private_account_data(account_id, &PrivateAccountKind::Regular(identifier), Account::default());
(account_id, cci)
}
@ -548,7 +548,7 @@ impl WalletCore {
.map(|(kind, res_acc)| {
let npk = &key_chain.nullifier_public_key;
let account_id = match &kind {
PrivateAccountKind::Account(identifier) => {
PrivateAccountKind::Regular(identifier) => {
nssa::AccountId::from((npk, *identifier))
}
PrivateAccountKind::Pda { program_id, seed, identifier } => {

View File

@ -18,6 +18,17 @@ pub enum PrivacyPreservingAccount {
vpk: ViewingPublicKey,
identifier: Identifier,
},
/// An owned private PDA: wallet holds the nsk/npk; account_id was derived via
/// `AccountId::for_private_pda`. Produces visibility mask 3.
PrivatePdaOwned(AccountId),
/// A foreign private PDA: wallet knows the recipient's npk/vpk but not their nsk.
/// Produces visibility mask 3 with a default (uninitialised) account.
PrivatePdaForeign {
account_id: AccountId,
npk: NullifierPublicKey,
vpk: ViewingPublicKey,
identifier: Identifier,
},
}
impl PrivacyPreservingAccount {
@ -31,11 +42,9 @@ impl PrivacyPreservingAccount {
matches!(
&self,
Self::PrivateOwned(_)
| Self::PrivateForeign {
npk: _,
vpk: _,
identifier: _
}
| Self::PrivateForeign { .. }
| Self::PrivatePdaOwned(_)
| Self::PrivatePdaForeign { .. }
)
}
}
@ -106,6 +115,28 @@ impl AccountManager {
(State::Private(pre), 2)
}
PrivacyPreservingAccount::PrivatePdaOwned(account_id) => {
let pre = private_acc_preparation(wallet, account_id).await?;
(State::Private(pre), 3)
}
PrivacyPreservingAccount::PrivatePdaForeign {
account_id,
npk,
vpk,
identifier,
} => {
let acc = nssa_core::account::Account::default();
let auth_acc = AccountWithMetadata::new(acc, false, account_id);
let pre = AccountPreparedData {
nsk: None,
npk,
identifier,
vpk,
pre_state: auth_acc,
proof: None,
};
(State::Private(pre), 3)
}
};
pre_states.push(state);
@ -235,7 +266,7 @@ async fn private_acc_preparation(
// TODO: Technically we could allow unauthorized owned accounts, but currently we don't have
// support from that in the wallet.
let sender_pre = AccountWithMetadata::new(from_acc.clone(), true, (&from_npk, from_identifier));
let sender_pre = AccountWithMetadata::new(from_acc.clone(), true, account_id);
Ok(AccountPreparedData {
nsk: Some(nsk),