From 2cee0b3861a76eb6852617c14611d29ad0066782 Mon Sep 17 00:00:00 2001 From: Artem Gureev Date: Tue, 30 Jun 2026 10:22:04 +0000 Subject: [PATCH] refactor(circuit): use PrivateAddressPlaintext in-guest --- lee/state_machine/core/src/circuit_io.rs | 34 +++--- .../execution_state.rs | 108 +++++++----------- .../bin/privacy_preserving_circuit/main.rs | 2 +- .../bin/privacy_preserving_circuit/output.rs | 77 +++++++------ 4 files changed, 106 insertions(+), 115 deletions(-) diff --git a/lee/state_machine/core/src/circuit_io.rs b/lee/state_machine/core/src/circuit_io.rs index 9bead310..1eed0053 100644 --- a/lee/state_machine/core/src/circuit_io.rs +++ b/lee/state_machine/core/src/circuit_io.rs @@ -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 { 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 { .. } diff --git a/program_methods/guest/src/bin/privacy_preserving_circuit/execution_state.rs b/program_methods/guest/src/bin/privacy_preserving_circuit/execution_state.rs index 09ad30af..1cc3dc37 100644 --- a/program_methods/guest/src/bin/privacy_preserving_circuit/execution_state.rs +++ b/program_methods/guest/src/bin/privacy_preserving_circuit/execution_state.rs @@ -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, + /// 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, authorized_accounts: HashSet, } @@ -61,17 +59,14 @@ impl ExecutionState { program_id: ProgramId, program_outputs: Vec, ) -> 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 = 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, - private_pda_by_position: &HashMap, + private_pda_by_position: &HashMap, authorized_accounts: &mut HashSet, 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)); } diff --git a/program_methods/guest/src/bin/privacy_preserving_circuit/main.rs b/program_methods/guest/src/bin/privacy_preserving_circuit/main.rs index a342665d..c09072e9 100644 --- a/program_methods/guest/src/bin/privacy_preserving_circuit/main.rs +++ b/program_methods/guest/src/bin/privacy_preserving_circuit/main.rs @@ -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); } diff --git a/program_methods/guest/src/bin/privacy_preserving_circuit/output.rs b/program_methods/guest/src/bin/privacy_preserving_circuit/output.rs index c799a1db..c5fe0570 100644 --- a/program_methods/guest/src/bin/privacy_preserving_circuit/output.rs +++ b/program_methods/guest/src/bin/privacy_preserving_circuit/output.rs @@ -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, ) -> 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, );