From 6fa05fa847b0eebc2372fe256de303e2c567a853 Mon Sep 17 00:00:00 2001 From: Artem Gureev Date: Tue, 30 Jun 2026 17:57:57 +0000 Subject: [PATCH] feat(lee_core): change private pda id derivation Private PDA ID generation now folds the underlying accound ID generated --- lee/state_machine/core/src/circuit_io.rs | 18 +++---- lee/state_machine/core/src/nullifier.rs | 27 ++++++++-- lee/state_machine/core/src/program.rs | 53 +++++++++++++------ .../execution_state.rs | 22 ++++---- .../bin/privacy_preserving_circuit/output.rs | 33 ++++++------ 5 files changed, 94 insertions(+), 59 deletions(-) diff --git a/lee/state_machine/core/src/circuit_io.rs b/lee/state_machine/core/src/circuit_io.rs index ea3580bc..d1dc5096 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, PrivateAddressPlaintext}, + account::{Account, AccountId, AccountWithMetadata}, 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 - /// `PrivateAddressPlaintext::new(NullifierPublicKey::from(nsk), vpk, identifier).account_id()` + /// `AccountId::for_regular_private_account(NullifierPublicKey::from(nsk), vpk, identifier)` /// and matched against `pre_state.account_id`. PrivateAuthorizedInit { vpk: ViewingPublicKey, @@ -63,8 +63,8 @@ pub enum InputAccountIdentity { identifier: Identifier, /// When `Some((seed, authority_program_id))`, the circuit binds this position via the /// external derivation check - /// `PrivateAddressPlaintext::new(npk, vpk, - /// identifier).pda_account_id(authority_program_id, seed) == pre_state.account_id` + /// `AccountId::for_regular_private_account(npk, vpk, + /// identifier).pda(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)>, @@ -80,8 +80,8 @@ pub enum InputAccountIdentity { identifier: Identifier, /// When `Some((seed, authority_program_id))`, the circuit binds this position via the /// external derivation check - /// `PrivateAddressPlaintext::new(npk, vpk, - /// identifier).pda_account_id(authority_program_id, seed) == pre_state.account_id` + /// `AccountId::for_regular_private_account(npk, vpk, + /// identifier).pda(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,20 +103,20 @@ impl InputAccountIdentity { } #[must_use] - pub fn private_pda_address(&self) -> Option> { + pub fn regular_account_id(&self) -> Option { match self { Self::PrivatePdaInit { npk, vpk, identifier, .. - } => Some(PrivateAddressPlaintext::new(*npk, vpk, *identifier)), + } => Some(AccountId::for_regular_private_account(*npk, vpk, *identifier)), Self::PrivatePdaUpdate { nsk, vpk, identifier, .. - } => Some(PrivateAddressPlaintext::new( + } => Some(AccountId::for_regular_private_account( NullifierPublicKey::from(nsk), vpk, *identifier, diff --git a/lee/state_machine/core/src/nullifier.rs b/lee/state_machine/core/src/nullifier.rs index 6df7bd2a..cd0cfc03 100644 --- a/lee/state_machine/core/src/nullifier.rs +++ b/lee/state_machine/core/src/nullifier.rs @@ -36,6 +36,27 @@ impl PrivateAddressPlaintext<'_> { } } +impl AccountId { + #[must_use] + pub fn for_regular_private_account( + npk: NullifierPublicKey, + vpk: &ViewingPublicKey, + identifier: Identifier, + ) -> Self { + let mut bytes = [0_u8; 32 + 32 + ViewingPublicKey::LEN + 16]; + bytes[0..32].copy_from_slice(PRIVATE_ACCOUNT_ID_PREFIX); + bytes[32..64].copy_from_slice(&npk.0); + bytes[64..64 + ViewingPublicKey::LEN].copy_from_slice(vpk.to_bytes()); + bytes[64 + ViewingPublicKey::LEN..].copy_from_slice(&identifier.to_le_bytes()); + Self::new( + Impl::hash_bytes(&bytes) + .as_bytes() + .try_into() + .expect("Conversion should not fail"), + ) + } +} + impl AsRef<[u8]> for NullifierPublicKey { fn as_ref(&self) -> &[u8] { self.0.as_slice() @@ -162,7 +183,7 @@ mod tests { 220, 68, 135, 10, 171, 182, 80, 54, 74, 228, 244, 236, 7, ]); - let account_id = PrivateAddressPlaintext::new(npk, &vpk, 0).account_id(); + let account_id = AccountId::for_regular_private_account(npk, &vpk, 0); assert_eq!(account_id, expected_account_id); } @@ -180,7 +201,7 @@ mod tests { 189, 170, 32, 181, 255, 231, 19, 92, 235, 59, 153, 185, 172, 206, ]); - let account_id = PrivateAddressPlaintext::new(npk, &vpk, 1).account_id(); + let account_id = AccountId::for_regular_private_account(npk, &vpk, 1); assert_eq!(account_id, expected_account_id); } @@ -199,7 +220,7 @@ mod tests { 159, 112, 84, 100, 133, 244, 16, 34, 221, 35, 128, 131, 98, 159, ]); - let account_id = PrivateAddressPlaintext::new(npk, &vpk, identifier).account_id(); + let account_id = AccountId::for_regular_private_account(npk, &vpk, identifier); assert_eq!(account_id, expected_account_id); } diff --git a/lee/state_machine/core/src/program.rs b/lee/state_machine/core/src/program.rs index fb5d6805..e8cab601 100644 --- a/lee/state_machine/core/src/program.rs +++ b/lee/state_machine/core/src/program.rs @@ -143,6 +143,26 @@ impl AccountId { ) } + #[must_use] + pub fn pda(self, program_id: &ProgramId, seed: &PdaSeed) -> Self { + use risc0_zkvm::sha::{Impl, Sha256 as _}; + const PRIVATE_PDA_PREFIX: &[u8; 32] = b"/LEE/v0.3/AccountId/PrivatePDA/\x00"; + + let mut bytes = [0_u8; 32 + 32 + 32 + 32]; + bytes[0..32].copy_from_slice(PRIVATE_PDA_PREFIX); + let program_id_bytes: &[u8] = + bytemuck::try_cast_slice(program_id).expect("ProgramId should be castable to &[u8]"); + bytes[32..64].copy_from_slice(program_id_bytes); + bytes[64..96].copy_from_slice(&seed.0); + bytes[96..128].copy_from_slice(self.value()); + Self::new( + Impl::hash_bytes(&bytes) + .as_bytes() + .try_into() + .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( @@ -152,14 +172,13 @@ impl AccountId { ) -> Self { match kind { PrivateAccountKind::Regular(identifier) => { - PrivateAddressPlaintext::new(*npk, vpk, *identifier).account_id() + Self::for_regular_private_account(*npk, vpk, *identifier) } PrivateAccountKind::Pda { program_id, seed, identifier, - } => PrivateAddressPlaintext::new(*npk, vpk, *identifier) - .pda_account_id(program_id, seed), + } => Self::for_regular_private_account(*npk, vpk, *identifier).pda(program_id, seed), } } } @@ -964,7 +983,7 @@ mod tests { 156, 13, 55, 32, 139, 91, 222, 209, 83, 172, 148, 123, 179, ]); assert_eq!( - PrivateAddressPlaintext::new(npk, &vpk, identifier).pda_account_id(&program_id, &seed), + AccountId::for_regular_private_account(npk, &vpk, identifier).pda(&program_id, &seed), expected ); } @@ -978,8 +997,8 @@ mod tests { let npk_b = NullifierPublicKey([4; 32]); let vpk = ViewingPublicKey::from_seed(&[1_u8; 32], &[2_u8; 32]); assert_ne!( - PrivateAddressPlaintext::new(npk_a, &vpk, u128::MAX).pda_account_id(&program_id, &seed), - PrivateAddressPlaintext::new(npk_b, &vpk, u128::MAX).pda_account_id(&program_id, &seed), + AccountId::for_regular_private_account(npk_a, &vpk, u128::MAX).pda(&program_id, &seed), + AccountId::for_regular_private_account(npk_b, &vpk, u128::MAX).pda(&program_id, &seed), ); } @@ -992,8 +1011,8 @@ mod tests { let npk = NullifierPublicKey([3; 32]); let vpk = ViewingPublicKey::from_seed(&[1_u8; 32], &[2_u8; 32]); assert_ne!( - PrivateAddressPlaintext::new(npk, &vpk, u128::MAX).pda_account_id(&program_id, &seed_a), - PrivateAddressPlaintext::new(npk, &vpk, u128::MAX).pda_account_id(&program_id, &seed_b), + AccountId::for_regular_private_account(npk, &vpk, u128::MAX).pda(&program_id, &seed_a), + AccountId::for_regular_private_account(npk, &vpk, u128::MAX).pda(&program_id, &seed_b), ); } @@ -1006,8 +1025,8 @@ mod tests { let npk = NullifierPublicKey([3; 32]); let vpk = ViewingPublicKey::from_seed(&[1_u8; 32], &[2_u8; 32]); assert_ne!( - PrivateAddressPlaintext::new(npk, &vpk, u128::MAX).pda_account_id(&program_id_a, &seed), - PrivateAddressPlaintext::new(npk, &vpk, u128::MAX).pda_account_id(&program_id_b, &seed), + AccountId::for_regular_private_account(npk, &vpk, u128::MAX).pda(&program_id_a, &seed), + AccountId::for_regular_private_account(npk, &vpk, u128::MAX).pda(&program_id_b, &seed), ); } @@ -1020,12 +1039,12 @@ mod tests { let npk = NullifierPublicKey([3; 32]); let vpk = ViewingPublicKey::from_seed(&[1_u8; 32], &[2_u8; 32]); assert_ne!( - PrivateAddressPlaintext::new(npk, &vpk, 0).pda_account_id(&program_id, &seed), - PrivateAddressPlaintext::new(npk, &vpk, 1).pda_account_id(&program_id, &seed), + AccountId::for_regular_private_account(npk, &vpk, 0).pda(&program_id, &seed), + AccountId::for_regular_private_account(npk, &vpk, 1).pda(&program_id, &seed), ); assert_ne!( - PrivateAddressPlaintext::new(npk, &vpk, 0).pda_account_id(&program_id, &seed), - PrivateAddressPlaintext::new(npk, &vpk, u128::MAX).pda_account_id(&program_id, &seed), + AccountId::for_regular_private_account(npk, &vpk, 0).pda(&program_id, &seed), + AccountId::for_regular_private_account(npk, &vpk, u128::MAX).pda(&program_id, &seed), ); } @@ -1038,7 +1057,7 @@ mod tests { let npk = NullifierPublicKey([3; 32]); let vpk = ViewingPublicKey::from_seed(&[1_u8; 32], &[2_u8; 32]); let private_id = - PrivateAddressPlaintext::new(npk, &vpk, u128::MAX).pda_account_id(&program_id, &seed); + AccountId::for_regular_private_account(npk, &vpk, u128::MAX).pda(&program_id, &seed); let public_id = AccountId::for_public_pda(&program_id, &seed); assert_ne!(private_id, public_id); } @@ -1080,7 +1099,7 @@ mod tests { assert_eq!( AccountId::for_private_account(&npk, &vpk, &PrivateAccountKind::Regular(identifier)), - PrivateAddressPlaintext::new(npk, &vpk, identifier).account_id(), + AccountId::for_regular_private_account(npk, &vpk, identifier), ); assert_eq!( AccountId::for_private_account( @@ -1092,7 +1111,7 @@ mod tests { identifier } ), - PrivateAddressPlaintext::new(npk, &vpk, identifier).pda_account_id(&program_id, &seed), + AccountId::for_regular_private_account(npk, &vpk, identifier).pda(&program_id, &seed), ); } 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 75c0e9f9..759e5cc5 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 @@ -5,7 +5,7 @@ use std::{ use lee_core::{ InputAccountIdentity, NullifierPublicKey, - account::{Account, AccountId, AccountWithMetadata, PrivateAddressPlaintext}, + account::{Account, AccountId, AccountWithMetadata}, program::{ AccountPostState, BlockValidityWindow, ChainedCall, Claim, DEFAULT_PROGRAM_ID, MAX_NUMBER_CHAINED_CALLS, PdaSeed, ProgramId, ProgramOutput, TimestampValidityWindow, @@ -22,7 +22,7 @@ pub struct ExecutionState { timestamp_validity_window: TimestampValidityWindow, /// Positions (in `pre_states`) of private-PDA accounts whose supplied npk has been bound to /// their `AccountId` via a proven - /// `PrivateAddressPlaintext::new(npk, vpk, identifier).pda_account_id(program_id, seed)` + /// `AccountId::for_regular_private_account(npk, vpk, identifier).pda(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 @@ -295,8 +295,7 @@ impl ExecutionState { seed: Some((seed, authority_program_id)), .. }) => { - let expected = PrivateAddressPlaintext::new(*npk, vpk, *identifier) - .pda_account_id(authority_program_id, seed); + let expected = AccountId::for_regular_private_account(*npk, vpk, *identifier).pda(authority_program_id, seed); assert_eq!( pre_account_id, expected, "External seed mismatch for PrivatePdaInit at position {pre_state_position}" @@ -311,8 +310,7 @@ impl ExecutionState { .. }) => { let npk = NullifierPublicKey::from(nsk); - let expected = PrivateAddressPlaintext::new(npk, vpk, *identifier) - .pda_account_id(authority_program_id, seed); + let expected = AccountId::for_regular_private_account(npk, vpk, *identifier).pda(authority_program_id, seed); assert_eq!( pre_account_id, expected, "External seed mismatch for PrivatePdaUpdate at position {pre_state_position}" @@ -392,9 +390,9 @@ impl ExecutionState { Claim::Authorized => {} Claim::Pda(seed) => { let pda = account_identity - .private_pda_address() + .regular_account_id() .expect("private PDA claim requires a private PDA account identity") - .pda_account_id(&program_id, &seed); + .pda(&program_id, &seed); assert_eq!( pre_account_id, pda, "Invalid private PDA claim for account {pre_account_id}" @@ -529,15 +527,15 @@ fn resolve_authorization_and_record_bindings( ) -> bool { let matched_caller_seed: Option<(PdaSeed, bool, ProgramId)> = caller_program_id.and_then(|caller| { - let pda_address = account_identities + let pda_base_id = account_identities .get(pre_state_position) - .and_then(InputAccountIdentity::private_pda_address); + .and_then(InputAccountIdentity::regular_account_id); 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(address) = &pda_address - && address.pda_account_id(&caller, seed) == pre_account_id + if let Some(base_id) = &pda_base_id + && base_id.pda(&caller, seed) == pre_account_id { return Some((*seed, true, caller)); } 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 0840abae..8ea2cd64 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, PrivateAddressPlaintext}, + account::{Account, AccountId, Nonce}, compute_digest_for_path, encryption::ViewingPublicKey, }; @@ -46,9 +46,8 @@ pub fn compute_circuit_output( nsk, identifier, } => { - let address = - PrivateAddressPlaintext::new(NullifierPublicKey::from(nsk), vpk, *identifier); - let account_id = address.account_id(); + let npk = NullifierPublicKey::from(nsk); + let account_id = AccountId::for_regular_private_account(npk, vpk, *identifier); assert_eq!(account_id, pre_state.account_id, "AccountId mismatch"); assert!( @@ -72,9 +71,9 @@ pub fn compute_circuit_output( &mut output_index, post_state, &account_id, - &PrivateAccountKind::Regular(address.identifier), - &address.npk, - address.vpk, + &PrivateAccountKind::Regular(*identifier), + &npk, + vpk, random_seed, new_nullifier, new_nonce, @@ -87,9 +86,8 @@ pub fn compute_circuit_output( membership_proof, identifier, } => { - let address = - PrivateAddressPlaintext::new(NullifierPublicKey::from(nsk), vpk, *identifier); - let account_id = address.account_id(); + let npk = NullifierPublicKey::from(nsk); + let account_id = AccountId::for_regular_private_account(npk, vpk, *identifier); assert_eq!(account_id, pre_state.account_id, "AccountId mismatch"); assert!( @@ -110,9 +108,9 @@ pub fn compute_circuit_output( &mut output_index, post_state, &account_id, - &PrivateAccountKind::Regular(address.identifier), - &address.npk, - address.vpk, + &PrivateAccountKind::Regular(*identifier), + &npk, + vpk, random_seed, new_nullifier, new_nonce, @@ -124,8 +122,7 @@ pub fn compute_circuit_output( npk, identifier, } => { - let address = PrivateAddressPlaintext::new(*npk, vpk, *identifier); - let account_id = address.account_id(); + let account_id = AccountId::for_regular_private_account(*npk, vpk, *identifier); assert_eq!(account_id, pre_state.account_id, "AccountId mismatch"); assert_eq!( @@ -149,9 +146,9 @@ pub fn compute_circuit_output( &mut output_index, post_state, &account_id, - &PrivateAccountKind::Regular(address.identifier), - &address.npk, - address.vpk, + &PrivateAccountKind::Regular(*identifier), + npk, + vpk, random_seed, new_nullifier, new_nonce,