Merge branch 'main' into schouhy/generalize-npk-to-multiple-accounts

This commit is contained in:
Sergio Chouhy
2026-04-24 01:04:55 -03:00
60 changed files with 948 additions and 104 deletions
@@ -11,7 +11,7 @@ use nssa_core::{
compute_digest_for_path,
program::{
AccountPostState, BlockValidityWindow, ChainedCall, Claim, DEFAULT_PROGRAM_ID,
MAX_NUMBER_CHAINED_CALLS, ProgramId, ProgramOutput, TimestampValidityWindow,
MAX_NUMBER_CHAINED_CALLS, PdaSeed, ProgramId, ProgramOutput, TimestampValidityWindow,
validate_execution,
},
};
@@ -23,15 +23,62 @@ struct ExecutionState {
post_states: HashMap<AccountId, Account>,
block_validity_window: BlockValidityWindow,
timestamp_validity_window: TimestampValidityWindow,
/// Positions (in `pre_states`) of mask-3 accounts whose supplied npk has been bound to
/// their `AccountId` via a proven `AccountId::for_private_pda(program_id, seed, npk)`
/// check.
/// Two proof paths populate this set: a `Claim::Pda(seed)` in a program's `post_state` on
/// that `pre_state`, or a caller's `ChainedCall.pda_seeds` entry matching that `pre_state`
/// under the private derivation. Binding is an idempotent property, not an event: the same
/// position can legitimately be bound through both paths in the same tx (e.g. a program
/// claims a private PDA and then delegates it to a callee), and the set uses `contains`,
/// not `assert!(insert)`. After the main loop, every mask-3 position must appear in this
/// set; otherwise the npk is unbound and the circuit rejects.
private_pda_bound_positions: HashSet<usize>,
/// Across the whole transaction, each `(program_id, seed)` pair may resolve to at most one
/// `AccountId`. A seed under a program can derive a family of accounts, one public PDA and
/// one private PDA per distinct npk. Without this check, a single `pda_seeds: [S]` entry in
/// a chained call could authorize multiple family members at once (different npks under the
/// same seed) and let a callee mix balances across them. Every claim and every
/// caller-authorization resolution is recorded here, either as a new `(program, seed)` →
/// `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>,
}
impl ExecutionState {
/// Validate program outputs and derive the overall execution state.
pub fn derive_from_outputs(
visibility_mask: &[u8],
private_account_keys: &[(NullifierPublicKey, SharedSecretKey)],
program_id: ProgramId,
program_outputs: Vec<ProgramOutput>,
) -> Self {
// Build position → npk map for mask-3 pre_states. `private_account_keys` is consumed in
// 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 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(|| {
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);
}
}
}
}
let block_valid_from = program_outputs
.iter()
.filter_map(|output| output.block_validity_window.start())
@@ -66,6 +113,9 @@ impl ExecutionState {
post_states: HashMap::new(),
block_validity_window,
timestamp_validity_window,
private_pda_bound_positions: HashSet::new(),
pda_family_binding: HashMap::new(),
private_pda_npk_by_position,
};
let Some(first_output) = program_outputs.first() else {
@@ -141,14 +191,11 @@ impl ExecutionState {
chained_calls.push_front((next_call.clone(), Some(chained_call.program_id)));
}
let authorized_pdas = nssa_core::program::compute_authorized_pdas(
caller_program_id,
&chained_call.pda_seeds,
);
execution_state.validate_and_sync_states(
visibility_mask,
chained_call.program_id,
&authorized_pdas,
caller_program_id,
&chained_call.pda_seeds,
program_output.pre_states,
program_output.post_states,
);
@@ -162,6 +209,19 @@ impl ExecutionState {
"Inner call without a chained call found",
);
// Every mask-3 pre_state must have had its npk bound to its account_id, either via a
// `Claim::Pda(seed)` in some program's post_state or via a caller's `pda_seeds` matching
// the private derivation. An unbound mask-3 pre_state has no cryptographic link between
// the supplied npk and the account_id, and must be rejected.
for (pos, &mask) in visibility_mask.iter().enumerate() {
if mask == 3 {
assert!(
execution_state.private_pda_bound_positions.contains(&pos),
"private PDA pre_state at position {pos} has no proven (seed, npk) binding via Claim::Pda or caller pda_seeds"
);
}
}
// Check that all modified uninitialized accounts were claimed
for (account_id, post) in execution_state
.pre_states
@@ -191,7 +251,8 @@ impl ExecutionState {
&mut self,
visibility_mask: &[u8],
program_id: ProgramId,
authorized_pdas: &HashSet<AccountId>,
caller_program_id: Option<ProgramId>,
caller_pda_seeds: &[PdaSeed],
pre_states: Vec<AccountWithMetadata>,
post_states: Vec<AccountPostState>,
) {
@@ -218,19 +279,28 @@ impl ExecutionState {
"Inconsistent pre state for account {pre_account_id}",
);
let previous_is_authorized = self
let (previous_is_authorized, pre_state_position) = self
.pre_states
.iter()
.find(|acc| acc.account_id == pre_account_id)
.enumerate()
.find(|(_, acc)| acc.account_id == pre_account_id)
.map_or_else(
|| panic!(
"Pre state must exist in execution state for account {pre_account_id}",
),
|acc| acc.is_authorized
|(pos, acc)| (acc.is_authorized, pos)
);
let is_authorized =
previous_is_authorized || authorized_pdas.contains(&pre_account_id);
let is_authorized = resolve_authorization_and_record_bindings(
&mut self.pda_family_binding,
&mut self.private_pda_bound_positions,
&self.private_pda_npk_by_position,
pre_account_id,
pre_state_position,
caller_program_id,
caller_pda_seeds,
previous_is_authorized,
);
assert_eq!(
pre_is_authorized, is_authorized,
@@ -257,9 +327,9 @@ impl ExecutionState {
.position(|acc| acc.account_id == pre_account_id)
.expect("Pre state must exist at this point");
let is_public_account = visibility_mask[pre_state_position] == 0;
if is_public_account {
match claim {
let mask = visibility_mask[pre_state_position];
match mask {
0 => match claim {
Claim::Authorized => {
// Note: no need to check authorized pdas because we have already
// checked consistency of authorization above.
@@ -269,18 +339,52 @@ impl ExecutionState {
);
}
Claim::Pda(seed) => {
let pda = AccountId::from((&program_id, &seed));
let pda = AccountId::for_public_pda(&program_id, &seed);
assert_eq!(
pre_account_id, pda,
"Invalid PDA claim for account {pre_account_id} which does not match derived PDA {pda}"
);
assert_family_binding(
&mut self.pda_family_binding,
program_id,
seed,
pre_account_id,
);
}
},
3 => {
match claim {
Claim::Authorized => {
assert!(
pre_is_authorized,
"Cannot claim unauthorized private PDA {pre_account_id}"
);
}
Claim::Pda(seed) => {
let npk = 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);
assert_eq!(
pre_account_id, pda,
"Invalid private PDA claim for account {pre_account_id}"
);
self.private_pda_bound_positions.insert(pre_state_position);
assert_family_binding(
&mut self.pda_family_binding,
program_id,
seed,
pre_account_id,
);
}
}
}
} else {
// We don't care about the exact claim mechanism for private accounts.
// This is because the main reason to have it is to protect against PDA griefing
// attacks in public execution, while private PDA doesn't make much sense
// anyway.
_ => {
// Mask 1/2: standard private accounts don't enforce the claim semantics.
// Unauthorized private claiming is intentionally allowed since operating
// these accounts requires the npk/nsk keypair anyway.
}
}
post.account_mut().program_owner = program_id;
@@ -304,6 +408,82 @@ impl ExecutionState {
}
}
/// Record or re-verify the `(program_id, seed) → account_id` family binding for the
/// transaction. Any claim or caller-seed authorization that resolves a `pre_state` under
/// `(program_id, seed)` must agree with every prior resolution of the same pair; otherwise a
/// single `pda_seeds: [seed]` entry could authorize multiple private-PDA family members at
/// once (different npks under the same seed) and let a callee mix balances across them. Free
/// function so callers can pass `&mut self.pda_family_binding` without holding a borrow on
/// the surrounding struct's other fields.
fn assert_family_binding(
bindings: &mut HashMap<(ProgramId, PdaSeed), AccountId>,
program_id: ProgramId,
seed: PdaSeed,
account_id: AccountId,
) {
match bindings.entry((program_id, seed)) {
Entry::Vacant(e) => {
e.insert(account_id);
}
Entry::Occupied(e) => {
assert_eq!(
*e.get(),
account_id,
"Two different accounts resolved under the same (program, seed) in one transaction: existing {}, new {account_id}",
e.get()
);
}
}
}
/// Resolve the authorization state of a `pre_state` seen again in a chained call and record
/// any resulting bindings. Returns `true` if the `pre_state` is authorized through either a
/// previously-seen authorization or a matching caller seed (under the public or private
/// derivation). When a caller seed matches, also records the `(caller, seed) → account_id`
/// family binding and, for the private form, marks the position in
/// `private_pda_bound_positions`. Only reachable when `caller_program_id.is_some()`,
/// top-level flows have no caller-emitted seeds, so binding at top level must come from the
/// claim path. Free function so callers can pass individual `&mut self.*` field borrows
/// without holding a borrow on the surrounding struct's other fields.
#[expect(
clippy::too_many_arguments,
reason = "breaking out a context struct does not buy us anything here"
)]
fn resolve_authorization_and_record_bindings(
pda_family_binding: &mut HashMap<(ProgramId, PdaSeed), AccountId>,
private_pda_bound_positions: &mut HashSet<usize>,
private_pda_npk_by_position: &HashMap<usize, NullifierPublicKey>,
pre_account_id: AccountId,
pre_state_position: usize,
caller_program_id: Option<ProgramId>,
caller_pda_seeds: &[PdaSeed],
previous_is_authorized: bool,
) -> bool {
let matched_caller_seed: Option<(PdaSeed, bool, ProgramId)> =
caller_program_id.and_then(|caller| {
caller_pda_seeds.iter().find_map(|seed| {
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) == pre_account_id
{
return Some((*seed, true, caller));
}
None
})
});
if let Some((seed, is_private_form, caller)) = matched_caller_seed {
assert_family_binding(pda_family_binding, caller, seed, pre_account_id);
if is_private_form {
private_pda_bound_positions.insert(pre_state_position);
}
}
previous_is_authorized || matched_caller_seed.is_some()
}
fn compute_circuit_output(
execution_state: ExecutionState,
visibility_mask: &[u8],
@@ -438,6 +618,88 @@ fn compute_circuit_output(
.checked_add(1)
.unwrap_or_else(|| panic!("Too many private accounts, output index overflow"));
}
3 => {
// Private PDA account. The supplied npk has already been bound to
// `pre_state.account_id` upstream in `validate_and_sync_states`, either via a
// `Claim::Pda(seed)` match or via a caller `pda_seeds` match, both of which
// assert `AccountId::for_private_pda(owner, seed, npk) == account_id`. The
// post-loop assertion in `derive_from_outputs` (see the
// `private_pda_bound_positions` check) guarantees that every mask-3
// position has been through at least one such binding, so this
// branch can safely use the wallet npk without re-verifying.
let Some((npk, shared_secret)) = private_keys_iter.next() else {
panic!("Missing private account key");
};
let (new_nullifier, new_nonce) = if pre_state.is_authorized {
// Existing private PDA with authentication (like mask 1)
let Some(nsk) = private_nsks_iter.next() else {
panic!("Missing private account nullifier secret key");
};
assert_eq!(
npk,
&NullifierPublicKey::from(nsk),
"Nullifier public key mismatch"
);
let Some(membership_proof_opt) = private_membership_proofs_iter.next() else {
panic!("Missing membership proof");
};
let new_nullifier = compute_nullifier_and_set_digest(
membership_proof_opt.as_ref(),
&pre_state.account,
npk,
nsk,
);
let new_nonce = pre_state.account.nonce.private_account_nonce_increment(nsk);
(new_nullifier, new_nonce)
} else {
// New private PDA (like mask 2). The default + unauthorized requirement
// here rules out use cases like a fully-private multisig, which would need
// a non-default, non-authorized private PDA input account.
// TODO(private-pdas-pr-2/3): relax this once the wallet can supply a
// `(seed, owner)` side input so the npk-to-account_id binding can be
// re-verified for an existing private PDA without a `Claim::Pda` or caller
// `pda_seeds` match.
assert_eq!(
pre_state.account,
Account::default(),
"New private PDA must be default"
);
let Some(membership_proof_opt) = private_membership_proofs_iter.next() else {
panic!("Missing membership proof");
};
assert!(
membership_proof_opt.is_none(),
"Membership proof must be None for new accounts"
);
let nullifier = Nullifier::for_account_initialization(npk);
let new_nonce = Nonce::private_account_nonce_init(npk);
((nullifier, DUMMY_COMMITMENT_HASH), new_nonce)
};
output.new_nullifiers.push(new_nullifier);
let mut post_with_updated_nonce = post_state;
post_with_updated_nonce.nonce = new_nonce;
let commitment_post = Commitment::new(npk, &post_with_updated_nonce);
let encrypted_account = EncryptionScheme::encrypt(
&post_with_updated_nonce,
shared_secret,
&commitment_post,
output_index,
);
output.new_commitments.push(commitment_post);
output.ciphertexts.push(encrypted_account);
output_index = output_index
.checked_add(1)
.unwrap_or_else(|| panic!("Too many private accounts, output index overflow"));
}
_ => panic!("Invalid visibility mask value"),
}
}
@@ -500,8 +762,12 @@ fn main() {
program_id,
} = env::read();
let execution_state =
ExecutionState::derive_from_outputs(&visibility_mask, program_id, program_outputs);
let execution_state = ExecutionState::derive_from_outputs(
&visibility_mask,
&private_account_keys,
program_id,
program_outputs,
);
let output = compute_circuit_output(
execution_state,