refactor(circuit): use PrivateAddressPlaintext in-guest

This commit is contained in:
Artem Gureev 2026-06-30 10:22:04 +00:00 committed by agureev
parent bfd0a4e9e2
commit 2cee0b3861
4 changed files with 106 additions and 115 deletions

View File

@ -3,7 +3,7 @@ use serde::{Deserialize, Serialize};
use crate::{
Commitment, CommitmentSetDigest, Identifier, MembershipProof, Nullifier, NullifierPublicKey,
NullifierSecretKey,
account::{Account, AccountWithMetadata},
account::{Account, AccountWithMetadata, PrivateAddressPlaintext},
encryption::{EncryptedAccountData, ViewingPublicKey},
program::{BlockValidityWindow, PdaSeed, ProgramId, ProgramOutput, TimestampValidityWindow},
};
@ -28,7 +28,7 @@ pub enum InputAccountIdentity {
Public,
/// Init of an authorized standalone private account: no membership proof. The `pre_state`
/// must be `Account::default()`. The `account_id` is derived as
/// `AccountId::for_regular_private_account(&NullifierPublicKey::from(nsk), vpk, identifier)`
/// `PrivateAddressPlaintext::new(NullifierPublicKey::from(nsk), vpk, identifier).account_id()`
/// and matched against `pre_state.account_id`.
PrivateAuthorizedInit {
vpk: ViewingPublicKey,
@ -55,8 +55,7 @@ pub enum InputAccountIdentity {
},
/// Init of a private PDA, unauthorized. The npk-to-account_id binding is proven upstream
/// via `Claim::Pda(seed)` or a caller's `pda_seeds` match. The identifier diversifies the
/// PDA within the `(program_id, seed, npk)` family: `AccountId::for_private_pda` uses it
/// as the 4th input.
/// PDA within the `(program_id, seed, npk)` family.
PrivatePdaInit {
vpk: ViewingPublicKey,
random_seed: [u8; 32],
@ -64,10 +63,10 @@ pub enum InputAccountIdentity {
identifier: Identifier,
/// When `Some((seed, authority_program_id))`, the circuit binds this position via the
/// external derivation check
/// `AccountId::for_private_pda(authority_program_id, seed, npk, vpk, identifier) ==
/// pre_state.account_id` rather than requiring a `Claim::Pda` or caller
/// `pda_seeds` to establish the binding. The `pre_state` must have `is_authorized
/// == false`.
/// `PrivateAddressPlaintext::new(npk, vpk,
/// identifier).pda_account_id(authority_program_id, seed) == pre_state.account_id`
/// rather than requiring a `Claim::Pda` or caller `pda_seeds` to establish the
/// binding. The `pre_state` must have `is_authorized == false`.
seed: Option<(PdaSeed, ProgramId)>,
},
/// Update of an existing private PDA, with membership proof. `npk` is derived
@ -81,9 +80,10 @@ pub enum InputAccountIdentity {
identifier: Identifier,
/// When `Some((seed, authority_program_id))`, the circuit binds this position via the
/// external derivation check
/// `AccountId::for_private_pda(authority_program_id, seed, npk, vpk, identifier) ==
/// pre_state.account_id` rather than requiring a caller `pda_seeds` to establish
/// the binding. The `pre_state` must have `is_authorized == false`.
/// `PrivateAddressPlaintext::new(npk, vpk,
/// identifier).pda_account_id(authority_program_id, seed) == pre_state.account_id`
/// rather than requiring a caller `pda_seeds` to establish the binding. The
/// `pre_state` must have `is_authorized == false`.
seed: Option<(PdaSeed, ProgramId)>,
},
}
@ -103,22 +103,24 @@ impl InputAccountIdentity {
}
#[must_use]
pub fn npk_vpk_if_private_pda(
&self,
) -> Option<(NullifierPublicKey, ViewingPublicKey, Identifier)> {
pub fn private_pda_address(&self) -> Option<PrivateAddressPlaintext> {
match self {
Self::PrivatePdaInit {
npk,
vpk,
identifier,
..
} => Some((*npk, vpk.clone(), *identifier)),
} => Some(PrivateAddressPlaintext::new(*npk, vpk.clone(), *identifier)),
Self::PrivatePdaUpdate {
nsk,
vpk,
identifier,
..
} => Some((NullifierPublicKey::from(nsk), vpk.clone(), *identifier)),
} => Some(PrivateAddressPlaintext::new(
NullifierPublicKey::from(nsk),
vpk.clone(),
*identifier,
)),
Self::Public
| Self::PrivateAuthorizedInit { .. }
| Self::PrivateAuthorizedUpdate { .. }

View File

@ -4,9 +4,8 @@ use std::{
};
use lee_core::{
Identifier, InputAccountIdentity, NullifierPublicKey,
account::{Account, AccountId, AccountWithMetadata},
encryption::ViewingPublicKey,
InputAccountIdentity,
account::{Account, AccountId, AccountWithMetadata, PrivateAddressPlaintext},
program::{
AccountPostState, BlockValidityWindow, ChainedCall, Claim, DEFAULT_PROGRAM_ID,
MAX_NUMBER_CHAINED_CALLS, PdaSeed, ProgramId, ProgramOutput, TimestampValidityWindow,
@ -22,15 +21,15 @@ pub struct ExecutionState {
block_validity_window: BlockValidityWindow,
timestamp_validity_window: TimestampValidityWindow,
/// Positions (in `pre_states`) of private-PDA accounts whose supplied npk has been bound to
/// their `AccountId` via a proven `AccountId::for_private_pda(program_id, seed, npk, vpk,
/// identifier)` 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 map uses `contains_key`,
/// not `assert!(insert)`. After the main loop, every private-PDA position must appear in this
/// map; otherwise the npk is unbound and the circuit rejects.
/// their `AccountId` via a proven
/// `PrivateAddressPlaintext::new(npk, vpk, identifier).pda_account_id(program_id, seed)`
/// 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 map
/// uses `contains_key`, not `assert!(insert)`. After the main loop, every private-PDA
/// position must appear in this map; otherwise the npk is unbound and the circuit rejects.
/// The stored `(ProgramId, PdaSeed)` is the owner program and seed, used in
/// `compute_circuit_output` to construct `PrivateAccountKind::Pda { program_id, seed,
/// identifier }`.
@ -44,13 +43,12 @@ pub 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 private-PDA `pre_state`'s position in `account_identities` to the (npk, vpk,
/// identifier) supplied for that position. Built once in `derive_from_outputs` by walking
/// `account_identities` and consulting `npk_vpk_if_private_pda`. Used later by the claim and
/// caller-seeds authorization paths to verify
/// `AccountId::for_private_pda(program_id, seed, npk, vpk, identifier) ==
/// pre_state.account_id`.
private_pda_by_position: HashMap<usize, (NullifierPublicKey, ViewingPublicKey, Identifier)>,
/// Map from a private-PDA `pre_state`'s position in `account_identities` to the
/// `PrivateAddressPlaintext` supplied for that position. Built once in `derive_from_outputs`
/// by walking `account_identities` and consulting `private_pda_address`. Used later by the
/// claim and caller-seeds authorization paths to verify
/// `address.pda_account_id(program_id, seed) == pre_state.account_id`.
private_pda_by_position: HashMap<usize, PrivateAddressPlaintext>,
authorized_accounts: HashSet<AccountId>,
}
@ -61,17 +59,14 @@ impl ExecutionState {
program_id: ProgramId,
program_outputs: Vec<ProgramOutput>,
) -> Self {
// Build position → (npk, identifier) map for private-PDA pre_states, indexed by position
// in `account_identities`. The vec is documented as 1:1 with the program's pre_state
// order, so position here matches `pre_state_position` used downstream in
// Build position → `PrivateAddressPlaintext` map for private-PDA pre_states, indexed by
// position in `account_identities`. The vec is documented as 1:1 with the program's
// pre_state order, so position here matches `pre_state_position` used downstream in
// `validate_and_sync_states`.
let mut private_pda_by_position: HashMap<
usize,
(NullifierPublicKey, ViewingPublicKey, Identifier),
> = HashMap::new();
let mut private_pda_by_position: HashMap<usize, PrivateAddressPlaintext> = HashMap::new();
for (pos, account_identity) in account_identities.iter().enumerate() {
if let Some((npk, vpk, identifier)) = account_identity.npk_vpk_if_private_pda() {
private_pda_by_position.insert(pos, (npk, vpk, identifier));
if let Some(address) = account_identity.private_pda_address() {
private_pda_by_position.insert(pos, address);
}
}
@ -312,19 +307,16 @@ impl ExecutionState {
let pre_state_position = self.pre_states.len();
let external_seed = match account_identities.get(pre_state_position) {
Some(InputAccountIdentity::PrivatePdaInit {
npk,
vpk,
identifier,
seed: Some((seed, authority_program_id)),
..
}) => {
let expected = AccountId::for_private_pda(
authority_program_id,
seed,
npk,
vpk,
*identifier,
);
let expected = self
.private_pda_by_position
.get(&pre_state_position)
.expect(
"private PDA pre_state must have an address in the position map",
)
.pda_account_id(authority_program_id, seed);
assert_eq!(
pre_account_id, expected,
"External seed mismatch for PrivatePdaInit at position {pre_state_position}"
@ -332,20 +324,16 @@ impl ExecutionState {
Some((*seed, *authority_program_id))
}
Some(InputAccountIdentity::PrivatePdaUpdate {
nsk,
vpk,
identifier,
seed: Some((seed, authority_program_id)),
..
}) => {
let npk = NullifierPublicKey::from(nsk);
let expected = AccountId::for_private_pda(
authority_program_id,
seed,
&npk,
vpk,
*identifier,
);
let expected = self
.private_pda_by_position
.get(&pre_state_position)
.expect(
"private PDA pre_state must have an address in the position map",
)
.pda_account_id(authority_program_id, seed);
assert_eq!(
pre_account_id, expected,
"External seed mismatch for PrivatePdaUpdate at position {pre_state_position}"
@ -424,19 +412,13 @@ impl ExecutionState {
match claim {
Claim::Authorized => {}
Claim::Pda(seed) => {
let (npk, vpk, identifier) = self
let pda = self
.private_pda_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,
vpk,
*identifier,
);
"private PDA pre_state must have an address in the position map",
)
.pda_account_id(&program_id, &seed);
assert_eq!(
pre_account_id, pda,
"Invalid private PDA claim for account {pre_account_id}"
@ -561,7 +543,7 @@ fn bind_private_pda_position(
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_by_position: &HashMap<usize, (NullifierPublicKey, ViewingPublicKey, Identifier)>,
private_pda_by_position: &HashMap<usize, PrivateAddressPlaintext>,
authorized_accounts: &mut HashSet<AccountId>,
pre_account_id: AccountId,
pre_state_position: usize,
@ -575,10 +557,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, vpk, identifier)) =
private_pda_by_position.get(&pre_state_position)
&& AccountId::for_private_pda(&caller, seed, npk, vpk, *identifier)
== pre_account_id
if let Some(address) = private_pda_by_position.get(&pre_state_position)
&& address.pda_account_id(&caller, seed) == pre_account_id
{
return Some((*seed, true, caller));
}

View File

@ -17,7 +17,7 @@ fn main() {
program_outputs,
);
let output = output::compute_circuit_output(execution_state, &account_identities);
let output = output::compute_circuit_output(execution_state, account_identities);
env::commit(&output);
}

View File

@ -2,7 +2,7 @@ use lee_core::{
Commitment, CommitmentSetDigest, DUMMY_COMMITMENT_HASH, EncryptedAccountData, EncryptionScheme,
EphemeralSecretKey, InputAccountIdentity, MembershipProof, Nullifier, NullifierPublicKey,
NullifierSecretKey, PrivacyPreservingCircuitOutput, PrivateAccountKind, SharedSecretKey,
account::{Account, AccountId, Nonce},
account::{Account, AccountId, Nonce, PrivateAddressPlaintext},
compute_digest_for_path,
encryption::ViewingPublicKey,
};
@ -11,7 +11,7 @@ use crate::execution_state::ExecutionState;
pub fn compute_circuit_output(
execution_state: ExecutionState,
account_identities: &[InputAccountIdentity],
account_identities: Vec<InputAccountIdentity>,
) -> PrivacyPreservingCircuitOutput {
let (block_validity_window, timestamp_validity_window, pda_seed_by_position, states_iter) =
execution_state.into_parts();
@ -33,7 +33,7 @@ pub fn compute_circuit_output(
let mut output_index = 0;
for (pos, (account_identity, (pre_state, post_state))) in
account_identities.iter().zip(states_iter).enumerate()
account_identities.into_iter().zip(states_iter).enumerate()
{
match account_identity {
InputAccountIdentity::Public => {
@ -46,8 +46,9 @@ pub fn compute_circuit_output(
nsk,
identifier,
} => {
let npk = NullifierPublicKey::from(nsk);
let account_id = AccountId::for_regular_private_account(&npk, vpk, *identifier);
let address =
PrivateAddressPlaintext::new(NullifierPublicKey::from(&nsk), vpk, identifier);
let account_id = address.account_id();
assert_eq!(account_id, pre_state.account_id, "AccountId mismatch");
assert!(
@ -71,10 +72,10 @@ pub fn compute_circuit_output(
&mut output_index,
post_state,
&account_id,
&PrivateAccountKind::Regular(*identifier),
&npk,
vpk,
random_seed,
&PrivateAccountKind::Regular(address.identifier),
&address.npk,
&address.vpk,
&random_seed,
new_nullifier,
new_nonce,
);
@ -86,8 +87,9 @@ pub fn compute_circuit_output(
membership_proof,
identifier,
} => {
let npk = NullifierPublicKey::from(nsk);
let account_id = AccountId::for_regular_private_account(&npk, vpk, *identifier);
let address =
PrivateAddressPlaintext::new(NullifierPublicKey::from(&nsk), vpk, identifier);
let account_id = address.account_id();
assert_eq!(account_id, pre_state.account_id, "AccountId mismatch");
assert!(
@ -96,22 +98,25 @@ pub fn compute_circuit_output(
);
let new_nullifier = compute_update_nullifier_and_set_digest(
membership_proof,
&membership_proof,
&pre_state.account,
&account_id,
nsk,
&nsk,
);
let new_nonce = pre_state.account.nonce.private_account_nonce_increment(nsk);
let new_nonce = pre_state
.account
.nonce
.private_account_nonce_increment(&nsk);
emit_private_output(
&mut output,
&mut output_index,
post_state,
&account_id,
&PrivateAccountKind::Regular(*identifier),
&npk,
vpk,
random_seed,
&PrivateAccountKind::Regular(address.identifier),
&address.npk,
&address.vpk,
&random_seed,
new_nullifier,
new_nonce,
);
@ -122,7 +127,8 @@ pub fn compute_circuit_output(
npk,
identifier,
} => {
let account_id = AccountId::for_regular_private_account(npk, vpk, *identifier);
let address = PrivateAddressPlaintext::new(npk, vpk, identifier);
let account_id = address.account_id();
assert_eq!(account_id, pre_state.account_id, "AccountId mismatch");
assert_eq!(
@ -146,10 +152,10 @@ pub fn compute_circuit_output(
&mut output_index,
post_state,
&account_id,
&PrivateAccountKind::Regular(*identifier),
npk,
vpk,
random_seed,
&PrivateAccountKind::Regular(address.identifier),
&address.npk,
&address.vpk,
&random_seed,
new_nullifier,
new_nonce,
);
@ -195,11 +201,11 @@ pub fn compute_circuit_output(
&PrivateAccountKind::Pda {
program_id: *authority_program_id,
seed: *seed,
identifier: *identifier,
identifier,
},
npk,
vpk,
random_seed,
&npk,
&vpk,
&random_seed,
new_nullifier,
new_nonce,
);
@ -223,15 +229,18 @@ pub fn compute_circuit_output(
);
let new_nullifier = compute_update_nullifier_and_set_digest(
membership_proof,
&membership_proof,
&pre_state.account,
&pre_state.account_id,
nsk,
&nsk,
);
let new_nonce = pre_state.account.nonce.private_account_nonce_increment(nsk);
let new_nonce = pre_state
.account
.nonce
.private_account_nonce_increment(&nsk);
let account_id = pre_state.account_id;
let npk = NullifierPublicKey::from(nsk);
let npk = NullifierPublicKey::from(&nsk);
let (authority_program_id, seed) = pda_seed_by_position
.get(&pos)
.expect("PrivatePdaUpdate position must be in pda_seed_by_position");
@ -243,11 +252,11 @@ pub fn compute_circuit_output(
&PrivateAccountKind::Pda {
program_id: *authority_program_id,
seed: *seed,
identifier: *identifier,
identifier,
},
&npk,
vpk,
random_seed,
&vpk,
&random_seed,
new_nullifier,
new_nonce,
);