refactor(privacy_preserving_circuit): address PR review comments

Bundle shared handle_* arguments into PrivateOutputHandler struct in
output.rs and fix misplaced docstring on resolve_external_seed in
execution_state.rs.
This commit is contained in:
Marvin Jones 2026-06-26 12:33:58 -04:00
parent 095383df26
commit 93fe131a62
3 changed files with 250 additions and 347 deletions

View File

@ -443,13 +443,6 @@ fn intersect_validity_windows(
(block_validity_window, timestamp_validity_window)
}
/// 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 resolve_external_seed(
account_identities: &[InputAccountIdentity],
pre_state_position: usize,
@ -510,6 +503,13 @@ fn resolve_external_seed(
}
}
/// 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,

View File

@ -8,6 +8,201 @@ use lee_core::{
use crate::execution_state::ExecutionState;
struct PrivateOutputHandler<'ctx> {
output: &'ctx mut PrivacyPreservingCircuitOutput,
output_index: &'ctx mut u32,
pre_state: &'ctx lee_core::account::AccountWithMetadata,
post_state: Account,
epk: &'ctx EphemeralPublicKey,
view_tag: u8,
ssk: &'ctx SharedSecretKey,
identifier: u128,
}
impl PrivateOutputHandler<'_> {
fn authorized_init(self, nsk: &NullifierSecretKey) {
let npk = NullifierPublicKey::from(nsk);
let account_id =
derive_and_verify_account_id(&npk, self.identifier, self.pre_state.account_id);
assert!(
self.pre_state.is_authorized,
"Pre-state not authorized for authenticated private account"
);
assert_eq!(
self.pre_state.account,
Account::default(),
"Found new private account with non default values"
);
let (new_nullifier, new_nonce) = init_nullifier_and_nonce(&account_id);
let kind = PrivateAccountKind::Regular(self.identifier);
self.emit_private_output(&account_id, &kind, new_nullifier, new_nonce);
}
fn authorized_update(self, nsk: &NullifierSecretKey, membership_proof: &MembershipProof) {
let npk = NullifierPublicKey::from(nsk);
let account_id =
derive_and_verify_account_id(&npk, self.identifier, self.pre_state.account_id);
assert!(
self.pre_state.is_authorized,
"Pre-state not authorized for authenticated private account"
);
let new_nullifier = compute_update_nullifier_and_set_digest(
membership_proof,
&self.pre_state.account,
&account_id,
nsk,
);
let new_nonce = self
.pre_state
.account
.nonce
.private_account_nonce_increment(nsk);
let kind = PrivateAccountKind::Regular(self.identifier);
self.emit_private_output(&account_id, &kind, new_nullifier, new_nonce);
}
fn unauthorized(self, npk: &NullifierPublicKey) {
let account_id =
derive_and_verify_account_id(npk, self.identifier, self.pre_state.account_id);
assert_eq!(
self.pre_state.account,
Account::default(),
"Found new private account with non default values",
);
assert!(
!self.pre_state.is_authorized,
"Found new private account marked as authorized."
);
let (new_nullifier, new_nonce) = init_nullifier_and_nonce(&account_id);
let kind = PrivateAccountKind::Regular(self.identifier);
self.emit_private_output(&account_id, &kind, new_nullifier, new_nonce);
}
fn pda_init(
self,
pos: usize,
pda_seed_by_position: &std::collections::HashMap<
usize,
(lee_core::program::ProgramId, lee_core::program::PdaSeed),
>,
) {
// The npk-to-account_id binding is established upstream in
// `validate_and_sync_states` via `Claim::Pda(seed)` or a caller `pda_seeds`
// match. Here we only enforce the init pre-conditions.
assert!(
!self.pre_state.is_authorized,
"PrivatePdaInit requires unauthorized pre_state"
);
assert_eq!(
self.pre_state.account,
Account::default(),
"New private PDA must be default"
);
let (new_nullifier, new_nonce) = init_nullifier_and_nonce(&self.pre_state.account_id);
let account_id = self.pre_state.account_id;
let (authority_program_id, seed) = pda_seed_by_position
.get(&pos)
.expect("PrivatePdaInit position must be in pda_seed_by_position");
let kind = PrivateAccountKind::Pda {
program_id: *authority_program_id,
seed: *seed,
identifier: self.identifier,
};
self.emit_private_output(&account_id, &kind, new_nullifier, new_nonce);
}
fn pda_update(
self,
nsk: &NullifierSecretKey,
membership_proof: &MembershipProof,
external_seed: Option<&(lee_core::program::PdaSeed, lee_core::program::ProgramId)>,
pos: usize,
pda_seed_by_position: &std::collections::HashMap<
usize,
(lee_core::program::ProgramId, lee_core::program::PdaSeed),
>,
) {
// With an external seed the binding comes from the circuit input and the
// pre_state is intentionally unauthorized; without one the binding comes from
// a Claim or caller pda_seeds, so the pre_state must already be authorized.
assert!(
self.pre_state.is_authorized ^ external_seed.is_some(),
"PrivatePdaUpdate requires authorized pre_state or external seed"
);
let new_nullifier = compute_update_nullifier_and_set_digest(
membership_proof,
&self.pre_state.account,
&self.pre_state.account_id,
nsk,
);
let new_nonce = self
.pre_state
.account
.nonce
.private_account_nonce_increment(nsk);
let account_id = self.pre_state.account_id;
let (authority_program_id, seed) = pda_seed_by_position
.get(&pos)
.expect("PrivatePdaUpdate position must be in pda_seed_by_position");
let kind = PrivateAccountKind::Pda {
program_id: *authority_program_id,
seed: *seed,
identifier: self.identifier,
};
self.emit_private_output(&account_id, &kind, new_nullifier, new_nonce);
}
fn emit_private_output(
self,
account_id: &AccountId,
kind: &PrivateAccountKind,
new_nullifier: (Nullifier, CommitmentSetDigest),
new_nonce: Nonce,
) {
self.output.new_nullifiers.push(new_nullifier);
let mut post_with_updated_nonce = self.post_state;
post_with_updated_nonce.nonce = new_nonce;
let commitment_post = Commitment::new(account_id, &post_with_updated_nonce);
let encrypted_account = EncryptionScheme::encrypt(
&post_with_updated_nonce,
kind,
self.ssk,
&commitment_post,
*self.output_index,
);
self.output.new_commitments.push(commitment_post);
self.output
.encrypted_private_post_states
.push(EncryptedAccountData {
ciphertext: encrypted_account,
epk: self.epk.clone(),
view_tag: self.view_tag,
});
*self.output_index = self
.output_index
.checked_add(1)
.unwrap_or_else(|| panic!("Too many private accounts, output index overflow"));
}
}
fn init_nullifier_and_nonce(account_id: &AccountId) -> ((Nullifier, CommitmentSetDigest), Nonce) {
let nullifier = (
Nullifier::for_account_initialization(account_id),
@ -27,253 +222,16 @@ fn derive_and_verify_account_id(
account_id
}
#[expect(
clippy::too_many_arguments,
reason = "extracted match arm with many destructured fields"
)]
fn handle_private_authorized_init(
output: &mut PrivacyPreservingCircuitOutput,
output_index: &mut u32,
pre_state: &lee_core::account::AccountWithMetadata,
post_state: Account,
epk: &EphemeralPublicKey,
view_tag: u8,
ssk: &SharedSecretKey,
nsk: &NullifierSecretKey,
identifier: u128,
) {
let npk = NullifierPublicKey::from(nsk);
let account_id = derive_and_verify_account_id(&npk, identifier, pre_state.account_id);
assert!(
pre_state.is_authorized,
"Pre-state not authorized for authenticated private account"
);
assert_eq!(
pre_state.account,
Account::default(),
"Found new private account with non default values"
);
let (new_nullifier, new_nonce) = init_nullifier_and_nonce(&account_id);
emit_private_output(
output,
output_index,
post_state,
&account_id,
&PrivateAccountKind::Regular(identifier),
ssk,
epk,
view_tag,
new_nullifier,
new_nonce,
);
}
#[expect(
clippy::too_many_arguments,
reason = "extracted match arm with many destructured fields"
)]
fn handle_private_authorized_update(
output: &mut PrivacyPreservingCircuitOutput,
output_index: &mut u32,
pre_state: &lee_core::account::AccountWithMetadata,
post_state: Account,
epk: &EphemeralPublicKey,
view_tag: u8,
ssk: &SharedSecretKey,
nsk: &NullifierSecretKey,
fn compute_update_nullifier_and_set_digest(
membership_proof: &MembershipProof,
identifier: u128,
) {
let npk = NullifierPublicKey::from(nsk);
let account_id = derive_and_verify_account_id(&npk, identifier, pre_state.account_id);
assert!(
pre_state.is_authorized,
"Pre-state not authorized for authenticated private account"
);
let new_nullifier = compute_update_nullifier_and_set_digest(
membership_proof,
&pre_state.account,
&account_id,
nsk,
);
let new_nonce = pre_state.account.nonce.private_account_nonce_increment(nsk);
emit_private_output(
output,
output_index,
post_state,
&account_id,
&PrivateAccountKind::Regular(identifier),
ssk,
epk,
view_tag,
new_nullifier,
new_nonce,
);
}
#[expect(
clippy::too_many_arguments,
reason = "extracted match arm with many destructured fields"
)]
fn handle_private_unauthorized(
output: &mut PrivacyPreservingCircuitOutput,
output_index: &mut u32,
pre_state: &lee_core::account::AccountWithMetadata,
post_state: Account,
epk: &EphemeralPublicKey,
view_tag: u8,
npk: &NullifierPublicKey,
ssk: &SharedSecretKey,
identifier: u128,
) {
let account_id = derive_and_verify_account_id(npk, identifier, pre_state.account_id);
assert_eq!(
pre_state.account,
Account::default(),
"Found new private account with non default values",
);
assert!(
!pre_state.is_authorized,
"Found new private account marked as authorized."
);
let (new_nullifier, new_nonce) = init_nullifier_and_nonce(&account_id);
emit_private_output(
output,
output_index,
post_state,
&account_id,
&PrivateAccountKind::Regular(identifier),
ssk,
epk,
view_tag,
new_nullifier,
new_nonce,
);
}
#[expect(
clippy::too_many_arguments,
reason = "PDA init has many distinct fields from the variant"
)]
fn handle_private_pda_init(
output: &mut PrivacyPreservingCircuitOutput,
output_index: &mut u32,
pre_state: &lee_core::account::AccountWithMetadata,
post_state: Account,
epk: &EphemeralPublicKey,
view_tag: u8,
ssk: &SharedSecretKey,
identifier: u128,
pos: usize,
pda_seed_by_position: &std::collections::HashMap<
usize,
(lee_core::program::ProgramId, lee_core::program::PdaSeed),
>,
) {
// The npk-to-account_id binding is established upstream in
// `validate_and_sync_states` via `Claim::Pda(seed)` or a caller `pda_seeds`
// match. Here we only enforce the init pre-conditions.
assert!(
!pre_state.is_authorized,
"PrivatePdaInit requires unauthorized pre_state"
);
assert_eq!(
pre_state.account,
Account::default(),
"New private PDA must be default"
);
let (new_nullifier, new_nonce) = init_nullifier_and_nonce(&pre_state.account_id);
let account_id = pre_state.account_id;
let (authority_program_id, seed) = pda_seed_by_position
.get(&pos)
.expect("PrivatePdaInit position must be in pda_seed_by_position");
emit_private_output(
output,
output_index,
post_state,
&account_id,
&PrivateAccountKind::Pda {
program_id: *authority_program_id,
seed: *seed,
identifier,
},
ssk,
epk,
view_tag,
new_nullifier,
new_nonce,
);
}
#[expect(
clippy::too_many_arguments,
reason = "PDA update has many distinct fields from the variant"
)]
fn handle_private_pda_update(
output: &mut PrivacyPreservingCircuitOutput,
output_index: &mut u32,
pre_state: &lee_core::account::AccountWithMetadata,
post_state: Account,
epk: &EphemeralPublicKey,
view_tag: u8,
ssk: &SharedSecretKey,
pre_account: &Account,
account_id: &AccountId,
nsk: &NullifierSecretKey,
membership_proof: &MembershipProof,
identifier: u128,
external_seed: Option<&(lee_core::program::PdaSeed, lee_core::program::ProgramId)>,
pos: usize,
pda_seed_by_position: &std::collections::HashMap<
usize,
(lee_core::program::ProgramId, lee_core::program::PdaSeed),
>,
) {
// With an external seed the binding comes from the circuit input and the
// pre_state is intentionally unauthorized; without one the binding comes from
// a Claim or caller pda_seeds, so the pre_state must already be authorized.
assert!(
pre_state.is_authorized ^ external_seed.is_some(),
"PrivatePdaUpdate requires authorized pre_state or external seed"
);
let new_nullifier = compute_update_nullifier_and_set_digest(
membership_proof,
&pre_state.account,
&pre_state.account_id,
nsk,
);
let new_nonce = pre_state.account.nonce.private_account_nonce_increment(nsk);
let account_id = pre_state.account_id;
let (authority_program_id, seed) = pda_seed_by_position
.get(&pos)
.expect("PrivatePdaUpdate position must be in pda_seed_by_position");
emit_private_output(
output,
output_index,
post_state,
&account_id,
&PrivateAccountKind::Pda {
program_id: *authority_program_id,
seed: *seed,
identifier,
},
ssk,
epk,
view_tag,
new_nullifier,
new_nonce,
);
) -> (Nullifier, CommitmentSetDigest) {
let commitment_pre = Commitment::new(account_id, pre_account);
let set_digest = compute_digest_for_path(&commitment_pre, membership_proof);
let nullifier = Nullifier::for_account_update(&commitment_pre, nsk);
(nullifier, set_digest)
}
pub fn compute_circuit_output(
@ -313,17 +271,17 @@ pub fn compute_circuit_output(
ssk,
nsk,
identifier,
} => handle_private_authorized_init(
&mut output,
&mut output_index,
&pre_state,
} => PrivateOutputHandler {
output: &mut output,
output_index: &mut output_index,
pre_state: &pre_state,
post_state,
epk,
*view_tag,
view_tag: *view_tag,
ssk,
nsk,
*identifier,
),
identifier: *identifier,
}
.authorized_init(nsk),
InputAccountIdentity::PrivateAuthorizedUpdate {
epk,
view_tag,
@ -331,35 +289,34 @@ pub fn compute_circuit_output(
nsk,
membership_proof,
identifier,
} => handle_private_authorized_update(
&mut output,
&mut output_index,
&pre_state,
} => PrivateOutputHandler {
output: &mut output,
output_index: &mut output_index,
pre_state: &pre_state,
post_state,
epk,
*view_tag,
view_tag: *view_tag,
ssk,
nsk,
membership_proof,
*identifier,
),
identifier: *identifier,
}
.authorized_update(nsk, membership_proof),
InputAccountIdentity::PrivateUnauthorized {
epk,
view_tag,
npk,
ssk,
identifier,
} => handle_private_unauthorized(
&mut output,
&mut output_index,
&pre_state,
} => PrivateOutputHandler {
output: &mut output,
output_index: &mut output_index,
pre_state: &pre_state,
post_state,
epk,
*view_tag,
npk,
view_tag: *view_tag,
ssk,
*identifier,
),
identifier: *identifier,
}
.unauthorized(npk),
InputAccountIdentity::PrivatePdaInit {
epk,
view_tag,
@ -367,18 +324,17 @@ pub fn compute_circuit_output(
ssk,
identifier,
seed: _,
} => handle_private_pda_init(
&mut output,
&mut output_index,
&pre_state,
} => PrivateOutputHandler {
output: &mut output,
output_index: &mut output_index,
pre_state: &pre_state,
post_state,
epk,
*view_tag,
view_tag: *view_tag,
ssk,
*identifier,
pos,
&pda_seed_by_position,
),
identifier: *identifier,
}
.pda_init(pos, &pda_seed_by_position),
InputAccountIdentity::PrivatePdaUpdate {
epk,
view_tag,
@ -387,17 +343,19 @@ pub fn compute_circuit_output(
membership_proof,
identifier,
seed: external_seed,
} => handle_private_pda_update(
&mut output,
&mut output_index,
&pre_state,
} => PrivateOutputHandler {
output: &mut output,
output_index: &mut output_index,
pre_state: &pre_state,
post_state,
epk,
*view_tag,
view_tag: *view_tag,
ssk,
identifier: *identifier,
}
.pda_update(
nsk,
membership_proof,
*identifier,
external_seed.as_ref(),
pos,
&pda_seed_by_position,
@ -407,58 +365,3 @@ pub fn compute_circuit_output(
output
}
#[expect(
clippy::too_many_arguments,
reason = "Inputs are distinct concerns from the variant arms; bundling would be artificial"
)]
fn emit_private_output(
output: &mut PrivacyPreservingCircuitOutput,
output_index: &mut u32,
post_state: Account,
account_id: &AccountId,
kind: &PrivateAccountKind,
shared_secret: &SharedSecretKey,
epk: &EphemeralPublicKey,
view_tag: u8,
new_nullifier: (Nullifier, CommitmentSetDigest),
new_nonce: 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(account_id, &post_with_updated_nonce);
let encrypted_account = EncryptionScheme::encrypt(
&post_with_updated_nonce,
kind,
shared_secret,
&commitment_post,
*output_index,
);
output.new_commitments.push(commitment_post);
output
.encrypted_private_post_states
.push(EncryptedAccountData {
ciphertext: encrypted_account,
epk: epk.clone(),
view_tag,
});
*output_index = output_index
.checked_add(1)
.unwrap_or_else(|| panic!("Too many private accounts, output index overflow"));
}
fn compute_update_nullifier_and_set_digest(
membership_proof: &MembershipProof,
pre_account: &Account,
account_id: &AccountId,
nsk: &NullifierSecretKey,
) -> (Nullifier, CommitmentSetDigest) {
let commitment_pre = Commitment::new(account_id, pre_account);
let set_digest = compute_digest_for_path(&commitment_pre, membership_proof);
let nullifier = Nullifier::for_account_update(&commitment_pre, nsk);
(nullifier, set_digest)
}