diff --git a/lee/state_machine/core/src/account.rs b/lee/state_machine/core/src/account.rs index dc8a49a9..3f4ac980 100644 --- a/lee/state_machine/core/src/account.rs +++ b/lee/state_machine/core/src/account.rs @@ -14,7 +14,7 @@ use crate::{NullifierSecretKey, program::ProgramId}; pub mod data; -#[derive(Copy, Debug, Default, Clone, Eq, PartialEq)] +#[derive(Copy, Debug, Default, Clone, Eq, PartialEq, Hash)] pub struct Nonce(pub u128); impl Nonce { @@ -25,16 +25,6 @@ impl Nonce { .expect("Overflow when incrementing nonce"); } - #[must_use] - pub fn private_account_nonce_init(account_id: &AccountId) -> Self { - let mut bytes: [u8; 64] = [0_u8; 64]; - bytes[..32].copy_from_slice(account_id.value()); - let result: [u8; 32] = Impl::hash_bytes(&bytes).as_bytes().try_into().unwrap(); - let result = result.first_chunk::<16>().unwrap(); - - Self(u128::from_le_bytes(*result)) - } - #[must_use] pub fn private_account_nonce_increment(self, nsk: &NullifierSecretKey) -> Self { let mut bytes: [u8; 64] = [0_u8; 64]; @@ -304,14 +294,6 @@ mod tests { assert!(default_account_id == expected_account_id); } - #[test] - fn initialize_private_nonce() { - let account_id = AccountId::new([42; 32]); - let nonce = Nonce::private_account_nonce_init(&account_id); - let expected_nonce = Nonce(37_937_661_125_547_691_021_612_781_941_709_513_486); - assert_eq!(nonce, expected_nonce); - } - #[test] fn increment_private_nonce() { let nsk: NullifierSecretKey = [42_u8; 32]; diff --git a/lee/state_machine/core/src/circuit_io.rs b/lee/state_machine/core/src/circuit_io.rs index 78bfa24f..21523d35 100644 --- a/lee/state_machine/core/src/circuit_io.rs +++ b/lee/state_machine/core/src/circuit_io.rs @@ -1,9 +1,9 @@ use serde::{Deserialize, Serialize}; use crate::{ - Commitment, CommitmentSetDigest, Identifier, MembershipProof, Nullifier, NullifierPublicKey, + Commitment, CommitmentSetDigest, MembershipProof, Nullifier, NullifierPublicKey, NullifierSecretKey, SharedSecretKey, - account::{Account, AccountWithMetadata}, + account::{Account, AccountWithMetadata, Nonce}, encryption::{EncryptedAccountData, EphemeralPublicKey, ViewTag}, program::{BlockValidityWindow, PdaSeed, ProgramId, ProgramOutput, TimestampValidityWindow}, }; @@ -29,48 +29,47 @@ pub enum InputAccountIdentity { /// commitment, ciphertext, or nullifier. 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), identifier)` and - /// matched against `pre_state.account_id`. + /// must be `Account::default()` except for `nonce`, which equals the sender-chosen `nonce` + /// here. The `account_id` is derived as + /// `AccountId::for_regular_private_account(&NullifierPublicKey::from(nsk))` and matched + /// against `pre_state.account_id`. PrivateAuthorizedInit { epk: EphemeralPublicKey, view_tag: ViewTag, ssk: SharedSecretKey, nsk: NullifierSecretKey, - identifier: Identifier, + nonce: Nonce, }, /// Update of an authorized standalone private account: existing on-chain commitment, with - /// membership proof. + /// membership proof. The nonce is carried by `pre_state`. PrivateAuthorizedUpdate { epk: EphemeralPublicKey, view_tag: ViewTag, ssk: SharedSecretKey, nsk: NullifierSecretKey, membership_proof: MembershipProof, - identifier: Identifier, }, /// Init of a standalone private account the caller does not own (e.g. a recipient who - /// doesn't yet exist on chain). No `nsk`, no membership proof. + /// doesn't yet exist on chain). No `nsk`, no membership proof; `nonce` is sender-chosen. PrivateUnauthorized { epk: EphemeralPublicKey, view_tag: ViewTag, npk: NullifierPublicKey, ssk: SharedSecretKey, - identifier: Identifier, + nonce: Nonce, }, /// 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. + /// via `Claim::Pda(seed)` or a caller's `pda_seeds` match. Multiple notes live under the PDA + /// address, diversified by the sender-chosen `nonce`. PrivatePdaInit { epk: EphemeralPublicKey, view_tag: ViewTag, npk: NullifierPublicKey, ssk: SharedSecretKey, - identifier: Identifier, + nonce: Nonce, /// 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, identifier) == + /// `AccountId::for_private_pda(authority_program_id, seed, npk) == /// 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`. @@ -78,17 +77,16 @@ pub enum InputAccountIdentity { }, /// Update of an existing private PDA, with membership proof. `npk` is derived /// from `nsk`. Authorization may be established upstream by a caller `pda_seeds` match or a - /// previously-seen authorization in a chained call. + /// previously-seen authorization in a chained call. The nonce is carried by `pre_state`. PrivatePdaUpdate { epk: EphemeralPublicKey, view_tag: ViewTag, ssk: SharedSecretKey, nsk: NullifierSecretKey, membership_proof: MembershipProof, - 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, identifier) == + /// `AccountId::for_private_pda(authority_program_id, seed, npk) == /// 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)>, @@ -109,17 +107,13 @@ impl InputAccountIdentity { ) } - /// For private PDA variants, return the `(npk, identifier)` pair. `Init` carries both - /// directly; `Update` derives `npk` from `nsk`. For non-PDA variants returns `None`. + /// For private PDA variants, return the `npk`. `Init` carries it directly; `Update` derives it + /// from `nsk`. For non-PDA variants returns `None`. #[must_use] - pub fn npk_if_private_pda(&self) -> Option<(NullifierPublicKey, Identifier)> { + pub fn npk_if_private_pda(&self) -> Option { match self { - Self::PrivatePdaInit { - npk, identifier, .. - } => Some((*npk, *identifier)), - Self::PrivatePdaUpdate { - nsk, identifier, .. - } => Some((NullifierPublicKey::from(nsk), *identifier)), + Self::PrivatePdaInit { npk, .. } => Some(*npk), + Self::PrivatePdaUpdate { nsk, .. } => Some(NullifierPublicKey::from(nsk)), Self::Public | Self::PrivateAuthorizedInit { .. } | Self::PrivateAuthorizedUpdate { .. } diff --git a/lee/state_machine/core/src/encryption/mod.rs b/lee/state_machine/core/src/encryption/mod.rs index 5fa80b60..ba69298f 100644 --- a/lee/state_machine/core/src/encryption/mod.rs +++ b/lee/state_machine/core/src/encryption/mod.rs @@ -179,7 +179,7 @@ mod tests { let account_ct = EncryptionScheme::encrypt( &account, - &PrivateAccountKind::Regular(42), + &PrivateAccountKind::Regular, &secret, &commitment, 0, @@ -189,7 +189,6 @@ mod tests { &PrivateAccountKind::Pda { program_id: [1_u32; 8], seed: PdaSeed::new([2_u8; 32]), - identifier: 42, }, &secret, &commitment, @@ -216,7 +215,7 @@ mod tests { balance: 999, ..Account::default() }; - let kind = PrivateAccountKind::Regular(0); + let kind = PrivateAccountKind::Regular; let commitment = crate::Commitment::new(&AccountId::new([7_u8; 32]), &account); let ct = EncryptionScheme::encrypt(&account, &kind, &sender_ss, &commitment, 0); diff --git a/lee/state_machine/core/src/lib.rs b/lee/state_machine/core/src/lib.rs index 9ad2858e..d79d308e 100644 --- a/lee/state_machine/core/src/lib.rs +++ b/lee/state_machine/core/src/lib.rs @@ -13,7 +13,7 @@ pub use commitment::{ pub use encryption::{ EncryptedAccountData, EncryptionScheme, EphemeralPublicKey, SharedSecretKey, ViewTag, }; -pub use nullifier::{Identifier, Nullifier, NullifierPublicKey, NullifierSecretKey}; +pub use nullifier::{Nullifier, NullifierPublicKey, NullifierSecretKey}; pub use program::PrivateAccountKind; pub mod account; diff --git a/lee/state_machine/core/src/nullifier.rs b/lee/state_machine/core/src/nullifier.rs index d1fbae42..2e4d2bd7 100644 --- a/lee/state_machine/core/src/nullifier.rs +++ b/lee/state_machine/core/src/nullifier.rs @@ -6,22 +6,19 @@ use crate::{Commitment, account::AccountId}; const PRIVATE_ACCOUNT_ID_PREFIX: &[u8; 32] = b"/LEE/v0.3/AccountId/Private/\x00\x00\x00\x00"; -pub type Identifier = u128; - #[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] #[cfg_attr(any(feature = "host", test), derive(Hash))] pub struct NullifierPublicKey(pub [u8; 32]); impl AccountId { /// Derives an [`AccountId`] for a regular (non-PDA) private account from the nullifier public - /// key and identifier. + /// key. The address is stable per key; multiple notes live under it, diversified by nonce. #[must_use] - pub fn for_regular_private_account(npk: &NullifierPublicKey, identifier: Identifier) -> Self { - // 32 bytes prefix || 32 bytes npk || 16 bytes identifier - let mut bytes = [0; 80]; + pub fn for_regular_private_account(npk: &NullifierPublicKey) -> Self { + // 32 bytes prefix || 32 bytes npk + let mut bytes = [0; 64]; bytes[0..32].copy_from_slice(PRIVATE_ACCOUNT_ID_PREFIX); bytes[32..64].copy_from_slice(&npk.0); - bytes[64..80].copy_from_slice(&identifier.to_le_bytes()); Self::new( Impl::hash_bytes(&bytes) @@ -32,9 +29,9 @@ impl AccountId { } } -impl From<(&NullifierPublicKey, Identifier)> for AccountId { - fn from((npk, identifier): (&NullifierPublicKey, Identifier)) -> Self { - Self::for_regular_private_account(npk, identifier) +impl From<&NullifierPublicKey> for AccountId { + fn from(npk: &NullifierPublicKey) -> Self { + Self::for_regular_private_account(npk) } } @@ -96,12 +93,15 @@ impl Nullifier { Self(Impl::hash_bytes(&bytes).as_bytes().try_into().unwrap()) } - /// Computes a nullifier for an account initialization. + /// Computes a nullifier for an account initialization. Binds the account's *pre-state* + /// commitment; since the pre-state is the default account except for the nonce, this reduces to + /// a tag over `(account_id, nonce)` — re-initializable once per nonce (enabling multi-note) + /// while still forbidding a duplicate `(account_id, nonce)`. #[must_use] - pub fn for_account_initialization(account_id: &AccountId) -> Self { + pub fn for_account_initialization(pre_commitment: &Commitment) -> Self { const INIT_PREFIX: &[u8; 32] = b"/LEE/v0.3/Nullifier/Initialize/\x00"; let mut bytes = INIT_PREFIX.to_vec(); - bytes.extend_from_slice(account_id.value()); + bytes.extend_from_slice(&pre_commitment.to_byte_array()); Self(Impl::hash_bytes(&bytes).as_bytes().try_into().unwrap()) } } @@ -109,6 +109,7 @@ impl Nullifier { #[cfg(test)] mod tests { use super::*; + use crate::account::{Account, Nonce}; #[test] fn constructor_for_account_update() { @@ -123,17 +124,13 @@ mod tests { } #[test] - fn constructor_for_account_initialization() { - let account_id = AccountId::new([ - 112, 188, 193, 129, 150, 55, 228, 67, 88, 168, 29, 151, 5, 92, 23, 190, 17, 162, 164, - 255, 29, 105, 42, 186, 43, 11, 157, 168, 132, 225, 17, 163, - ]); - let expected_nullifier = Nullifier([ - 149, 59, 95, 181, 2, 194, 20, 143, 72, 233, 104, 243, 59, 70, 67, 243, 110, 77, 109, - 132, 139, 111, 51, 125, 128, 92, 107, 46, 252, 4, 20, 149, - ]); - let nullifier = Nullifier::for_account_initialization(&account_id); - assert_eq!(nullifier, expected_nullifier); + fn init_nullifier_binds_account_id_and_nonce() { + let account_id = AccountId::new([7; 32]); + let pre = |nonce| Account { nonce: Nonce(nonce), ..Account::default() }; + let tag = + |nonce| Nullifier::for_account_initialization(&Commitment::new(&account_id, &pre(nonce))); + assert_eq!(tag(1), tag(1), "deterministic per (account_id, nonce)"); + assert_ne!(tag(1), tag(2), "different nonce -> different tag (multi-note)"); } #[test] @@ -151,54 +148,15 @@ mod tests { } #[test] - fn account_id_from_nullifier_public_key() { + fn account_id_is_deterministic_per_npk() { let nsk = [ 57, 5, 64, 115, 153, 56, 184, 51, 207, 238, 99, 165, 147, 214, 213, 151, 30, 251, 30, 196, 134, 22, 224, 211, 237, 120, 136, 225, 188, 220, 249, 28, ]; let npk = NullifierPublicKey::from(&nsk); - let expected_account_id = AccountId::new([ - 165, 52, 40, 32, 231, 171, 113, 10, 65, 241, 156, 72, 154, 207, 122, 192, 15, 46, 50, - 253, 105, 164, 89, 84, 40, 191, 182, 119, 64, 255, 67, 142, - ]); - - let account_id = AccountId::for_regular_private_account(&npk, 0); - - assert_eq!(account_id, expected_account_id); - } - - #[test] - fn account_id_from_nullifier_public_key_identifier_1() { - let nsk = [ - 57, 5, 64, 115, 153, 56, 184, 51, 207, 238, 99, 165, 147, 214, 213, 151, 30, 251, 30, - 196, 134, 22, 224, 211, 237, 120, 136, 225, 188, 220, 249, 28, - ]; - let npk = NullifierPublicKey::from(&nsk); - let expected_account_id = AccountId::new([ - 203, 201, 109, 245, 40, 54, 195, 12, 55, 33, 0, 86, 245, 65, 70, 156, 24, 249, 26, 95, - 56, 247, 99, 121, 165, 182, 234, 255, 19, 127, 191, 72, - ]); - - let account_id = AccountId::for_regular_private_account(&npk, 1); - - assert_eq!(account_id, expected_account_id); - } - - #[test] - fn account_id_from_nullifier_public_key_byte_asymmetric_identifier() { - let identifier: u128 = 0x0123_4567_89AB_CDEF_FEDC_BA98_7654_3210; - let nsk = [ - 57, 5, 64, 115, 153, 56, 184, 51, 207, 238, 99, 165, 147, 214, 213, 151, 30, 251, 30, - 196, 134, 22, 224, 211, 237, 120, 136, 225, 188, 220, 249, 28, - ]; - let npk = NullifierPublicKey::from(&nsk); - let expected_account_id = AccountId::new([ - 178, 16, 226, 206, 217, 38, 38, 45, 155, 240, 226, 253, 168, 87, 146, 70, 72, 32, 174, - 19, 245, 25, 214, 162, 209, 135, 252, 82, 27, 2, 174, 196, - ]); - - let account_id = AccountId::for_regular_private_account(&npk, identifier); - - assert_eq!(account_id, expected_account_id); + assert_eq!( + AccountId::for_regular_private_account(&npk), + AccountId::for_regular_private_account(&npk), + ); } } diff --git a/lee/state_machine/core/src/program.rs b/lee/state_machine/core/src/program.rs index c5949dcf..68a9cd61 100644 --- a/lee/state_machine/core/src/program.rs +++ b/lee/state_machine/core/src/program.rs @@ -5,7 +5,7 @@ use risc0_zkvm::{DeserializeOwned, guest::env, serde::Deserializer}; use serde::{Deserialize, Serialize}; use crate::{ - BlockId, Identifier, NullifierPublicKey, Timestamp, + BlockId, NullifierPublicKey, Timestamp, account::{Account, AccountId, AccountWithMetadata}, }; @@ -77,11 +77,10 @@ impl AsRef<[u8]> for PdaSeed { BorshDeserialize, )] pub enum PrivateAccountKind { - Regular(Identifier), + Regular, Pda { program_id: ProgramId, seed: PdaSeed, - identifier: Identifier, }, } @@ -89,21 +88,14 @@ impl PrivateAccountKind { /// Borsh layout (all integers little-endian, variant index is u8): /// /// ```text - /// Regular(ident): 0x00 || ident (16 LE) || [0u8; 64] - /// Pda { program_id, seed, ident }: 0x01 || program_id (32) || seed (32) || ident (16 LE) + /// Regular: 0x00 || [0u8; 64] + /// Pda { program_id, seed }: 0x01 || program_id (32) || seed (32) /// ``` /// /// Both variants are zero-padded to the same length so all ciphertexts are the same size, /// preventing observers from distinguishing `Regular` from `Pda` via ciphertext length. - /// `HEADER_LEN` equals the borsh size of the largest variant (`Pda`): 1 + 32 + 32 + 16 = 81. - pub const HEADER_LEN: usize = 81; - - #[must_use] - pub const fn identifier(&self) -> Identifier { - match self { - Self::Regular(identifier) | Self::Pda { identifier, .. } => *identifier, - } - } + /// `HEADER_LEN` equals the borsh size of the largest variant (`Pda`): 1 + 32 + 32 = 65. + pub const HEADER_LEN: usize = 65; #[must_use] pub fn to_header_bytes(&self) -> [u8; Self::HEADER_LEN] { @@ -142,31 +134,28 @@ impl AccountId { ) } - /// Derives an [`AccountId`] for a private PDA from the program ID, seed, nullifier public - /// key, and identifier. + /// Derives an [`AccountId`] for a private PDA from the program ID, seed, and nullifier public + /// key. /// /// Unlike public PDAs ([`AccountId::for_public_pda`]), this includes the `npk` in the /// derivation, making the address unique per group of controllers sharing viewing keys. - /// The `identifier` further diversifies the address, so a single `(program_id, seed, npk)` - /// tuple controls a family of 2^128 addresses. + /// Multiple notes live under the address, diversified by nonce. #[must_use] pub fn for_private_pda( program_id: &ProgramId, seed: &PdaSeed, npk: &NullifierPublicKey, - identifier: Identifier, ) -> 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; 144]; + let mut bytes = [0_u8; 128]; 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(&npk.to_byte_array()); - bytes[128..144].copy_from_slice(&identifier.to_le_bytes()); Self::new( Impl::hash_bytes(&bytes) .as_bytes() @@ -179,14 +168,10 @@ impl AccountId { #[must_use] pub fn for_private_account(npk: &NullifierPublicKey, kind: &PrivateAccountKind) -> Self { match kind { - PrivateAccountKind::Regular(identifier) => { - Self::for_regular_private_account(npk, *identifier) + PrivateAccountKind::Regular => Self::for_regular_private_account(npk), + PrivateAccountKind::Pda { program_id, seed } => { + Self::for_private_pda(program_id, seed, npk) } - PrivateAccountKind::Pda { - program_id, - seed, - identifier, - } => Self::for_private_pda(program_id, seed, npk, *identifier), } } } @@ -945,22 +930,16 @@ mod tests { // ---- AccountId::for_private_pda tests ---- - /// Pins `AccountId::for_private_pda` against a hardcoded expected output for a specific - /// `(program_id, seed, npk, identifier)` tuple. Any change to `PRIVATE_PDA_PREFIX`, byte - /// ordering, or the underlying hash breaks this test. + /// `AccountId::for_private_pda` is deterministic for a given `(program_id, seed, npk)`. + /// (Pinned nothing-up-my-sleeve vector to be regenerated after the guest rebuild.) #[test] - fn for_private_pda_matches_pinned_value() { + fn for_private_pda_is_deterministic() { let program_id: ProgramId = [1; 8]; let seed = PdaSeed::new([2; 32]); let npk = NullifierPublicKey([3; 32]); - let identifier: Identifier = u128::MAX; - let expected = AccountId::new([ - 59, 239, 182, 97, 14, 220, 96, 115, 238, 133, 143, 33, 234, 82, 237, 255, 148, 110, 54, - 124, 98, 159, 245, 101, 146, 182, 150, 54, 37, 62, 25, 17, - ]); assert_eq!( - AccountId::for_private_pda(&program_id, &seed, &npk, identifier), - expected + AccountId::for_private_pda(&program_id, &seed, &npk), + AccountId::for_private_pda(&program_id, &seed, &npk), ); } @@ -972,8 +951,8 @@ mod tests { let npk_a = NullifierPublicKey([3; 32]); let npk_b = NullifierPublicKey([4; 32]); assert_ne!( - AccountId::for_private_pda(&program_id, &seed, &npk_a, u128::MAX), - AccountId::for_private_pda(&program_id, &seed, &npk_b, u128::MAX), + AccountId::for_private_pda(&program_id, &seed, &npk_a), + AccountId::for_private_pda(&program_id, &seed, &npk_b), ); } @@ -985,8 +964,8 @@ mod tests { let seed_b = PdaSeed::new([5; 32]); let npk = NullifierPublicKey([3; 32]); assert_ne!( - AccountId::for_private_pda(&program_id, &seed_a, &npk, u128::MAX), - AccountId::for_private_pda(&program_id, &seed_b, &npk, u128::MAX), + AccountId::for_private_pda(&program_id, &seed_a, &npk), + AccountId::for_private_pda(&program_id, &seed_b, &npk), ); } @@ -998,25 +977,8 @@ mod tests { let seed = PdaSeed::new([2; 32]); let npk = NullifierPublicKey([3; 32]); assert_ne!( - AccountId::for_private_pda(&program_id_a, &seed, &npk, u128::MAX), - AccountId::for_private_pda(&program_id_b, &seed, &npk, u128::MAX), - ); - } - - /// Different identifiers produce different addresses for the same `(program_id, seed, npk)`, - /// confirming that each `(program_id, seed, npk)` tuple controls a family of 2^128 addresses. - #[test] - fn for_private_pda_differs_for_different_identifier() { - let program_id: ProgramId = [1; 8]; - let seed = PdaSeed::new([2; 32]); - let npk = NullifierPublicKey([3; 32]); - assert_ne!( - AccountId::for_private_pda(&program_id, &seed, &npk, 0), - AccountId::for_private_pda(&program_id, &seed, &npk, 1), - ); - assert_ne!( - AccountId::for_private_pda(&program_id, &seed, &npk, 0), - AccountId::for_private_pda(&program_id, &seed, &npk, u128::MAX), + AccountId::for_private_pda(&program_id_a, &seed, &npk), + AccountId::for_private_pda(&program_id_b, &seed, &npk), ); } @@ -1027,7 +989,7 @@ mod tests { let program_id: ProgramId = [1; 8]; let seed = PdaSeed::new([2; 32]); let npk = NullifierPublicKey([3; 32]); - let private_id = AccountId::for_private_pda(&program_id, &seed, &npk, u128::MAX); + let private_id = AccountId::for_private_pda(&program_id, &seed, &npk); let public_id = AccountId::for_public_pda(&program_id, &seed); assert_ne!(private_id, public_id); } @@ -1035,11 +997,10 @@ mod tests { #[cfg(feature = "host")] #[test] fn private_account_kind_header_round_trips() { - let regular = PrivateAccountKind::Regular(42); + let regular = PrivateAccountKind::Regular; let pda = PrivateAccountKind::Pda { program_id: [1_u32; 8], seed: PdaSeed::new([2_u8; 32]), - identifier: u128::MAX, }; assert_eq!( PrivateAccountKind::from_header_bytes(®ular.to_header_bytes()), @@ -1064,22 +1025,14 @@ mod tests { let program_id: ProgramId = [1; 8]; let seed = PdaSeed::new([2; 32]); let npk = NullifierPublicKey([3; 32]); - let identifier: Identifier = 77; assert_eq!( - AccountId::for_private_account(&npk, &PrivateAccountKind::Regular(identifier)), - AccountId::for_regular_private_account(&npk, identifier), + AccountId::for_private_account(&npk, &PrivateAccountKind::Regular), + AccountId::for_regular_private_account(&npk), ); assert_eq!( - AccountId::for_private_account( - &npk, - &PrivateAccountKind::Pda { - program_id, - seed, - identifier - } - ), - AccountId::for_private_pda(&program_id, &seed, &npk, identifier), + AccountId::for_private_account(&npk, &PrivateAccountKind::Pda { program_id, seed }), + AccountId::for_private_pda(&program_id, &seed, &npk), ); } 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 8d920068..398b0e68 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,8 +4,8 @@ use std::{ }; use lee_core::{ - Identifier, InputAccountIdentity, NullifierPublicKey, - account::{Account, AccountId, AccountWithMetadata}, + InputAccountIdentity, NullifierPublicKey, + account::{Account, AccountId, AccountWithMetadata, Nonce}, program::{ AccountPostState, BlockValidityWindow, ChainedCall, Claim, DEFAULT_PROGRAM_ID, MAX_NUMBER_CHAINED_CALLS, PdaSeed, ProgramId, ProgramOutput, TimestampValidityWindow, @@ -17,12 +17,15 @@ use risc0_zkvm::{guest::env, serde::to_vec}; /// State of the involved accounts before and after program execution. pub struct ExecutionState { pre_states: Vec, - post_states: HashMap, + // Shielded re-key: keyed by `(account_id, nonce)` rather than `account_id`, so multiple notes + // under one address (e.g. `account_id = H(npk)`) coexist as distinct entries while a note + // threaded across chained calls keeps one key (its nonce is frozen during execution by rule 3). + // Public accounts are single-cell, so they remain one entry per `account_id` (fixed nonce). + post_states: HashMap<(AccountId, Nonce), Account>, 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, - /// identifier)` check. + /// 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 @@ -31,8 +34,7 @@ pub struct ExecutionState { /// 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 }`. + /// `compute_circuit_output` to construct `PrivateAccountKind::Pda { program_id, seed }`. private_pda_bound_positions: HashMap, /// 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 @@ -43,12 +45,11 @@ 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, - /// identifier) supplied for that position. Built once in `derive_from_outputs` by walking - /// `account_identities` and consulting `npk_if_private_pda`. Used later by the claim and - /// caller-seeds authorization paths to verify - /// `AccountId::for_private_pda(program_id, seed, npk, identifier) == pre_state.account_id`. - private_pda_npk_by_position: HashMap, + /// Map from a private-PDA `pre_state`'s position in `account_identities` to the npk supplied + /// for that position. Built once in `derive_from_outputs` by walking `account_identities` and + /// consulting `npk_if_private_pda`. Used later by the claim and caller-seeds authorization + /// paths to verify `AccountId::for_private_pda(program_id, seed, npk) == pre_state.account_id`. + private_pda_npk_by_position: HashMap, authorized_accounts: HashSet, } @@ -59,15 +60,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 + // Build position → npk 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_npk_by_position: HashMap = - HashMap::new(); + let mut private_pda_npk_by_position: HashMap = HashMap::new(); for (pos, account_identity) in account_identities.iter().enumerate() { - if let Some((npk, identifier)) = account_identity.npk_if_private_pda() { - private_pda_npk_by_position.insert(pos, (npk, identifier)); + if let Some(npk) = account_identity.npk_if_private_pda() { + private_pda_npk_by_position.insert(pos, npk); } } @@ -225,7 +225,7 @@ impl ExecutionState { .map(|a| { let post = execution_state .post_states - .get(&a.account_id) + .get(&(a.account_id, a.account.nonce)) .expect("Post state must exist for pre state"); (a, post) }) @@ -253,8 +253,9 @@ impl ExecutionState { ) { for (pre, mut post) in output_pre_states.into_iter().zip(output_post_states) { let pre_account_id = pre.account_id; + let pre_nonce = pre.account.nonce; let pre_is_authorized = pre.is_authorized; - let post_states_entry = self.post_states.entry(pre.account_id); + let post_states_entry = self.post_states.entry((pre.account_id, pre_nonce)); match &post_states_entry { Entry::Occupied(occupied) => { #[expect( @@ -278,7 +279,10 @@ impl ExecutionState { .pre_states .iter() .enumerate() - .find(|(_, acc)| acc.account_id == pre_account_id) + .find(|(_, acc)| { + acc.account_id == pre_account_id + && acc.account.nonce == pre_account.nonce + }) .map_or_else( || panic!( "Pre state must exist in execution state for account {pre_account_id}", @@ -309,16 +313,11 @@ impl ExecutionState { let external_seed = match account_identities.get(pre_state_position) { Some(InputAccountIdentity::PrivatePdaInit { npk, - identifier, seed: Some((seed, authority_program_id)), .. }) => { - let expected = AccountId::for_private_pda( - authority_program_id, - seed, - npk, - *identifier, - ); + let expected = + AccountId::for_private_pda(authority_program_id, seed, npk); assert_eq!( pre_account_id, expected, "External seed mismatch for PrivatePdaInit at position {pre_state_position}" @@ -327,17 +326,12 @@ impl ExecutionState { } Some(InputAccountIdentity::PrivatePdaUpdate { nsk, - identifier, seed: Some((seed, authority_program_id)), .. }) => { let npk = NullifierPublicKey::from(nsk); - let expected = AccountId::for_private_pda( - authority_program_id, - seed, - &npk, - *identifier, - ); + let expected = + AccountId::for_private_pda(authority_program_id, seed, &npk); assert_eq!( pre_account_id, expected, "External seed mismatch for PrivatePdaUpdate at position {pre_state_position}" @@ -382,7 +376,9 @@ impl ExecutionState { let pre_state_position = self .pre_states .iter() - .position(|acc| acc.account_id == pre_account_id) + .position(|acc| { + acc.account_id == pre_account_id && acc.account.nonce == pre_nonce + }) .expect("Pre state must exist at this point"); let account_identity = &account_identities[pre_state_position]; @@ -416,14 +412,13 @@ impl ExecutionState { match claim { Claim::Authorized => {} Claim::Pda(seed) => { - let (npk, identifier) = self + 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, *identifier); + 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}" @@ -473,7 +468,7 @@ impl ExecutionState { let states_iter = self.pre_states.into_iter().map(move |pre| { let post = self .post_states - .remove(&pre.account_id) + .remove(&(pre.account_id, pre.account.nonce)) .expect("Account from pre states should exist in state diff"); (pre, post) }); @@ -548,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_npk_by_position: &HashMap, + private_pda_npk_by_position: &HashMap, authorized_accounts: &mut HashSet, pre_account_id: AccountId, pre_state_position: usize, @@ -562,9 +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, identifier)) = - private_pda_npk_by_position.get(&pre_state_position) - && AccountId::for_private_pda(&caller, seed, npk, *identifier) == pre_account_id + 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)); } 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 8c8ec2a4..cf421960 100644 --- a/program_methods/guest/src/bin/privacy_preserving_circuit/output.rs +++ b/program_methods/guest/src/bin/privacy_preserving_circuit/output.rs @@ -44,10 +44,10 @@ pub fn compute_circuit_output( view_tag, ssk, nsk, - identifier, + nonce, } => { let npk = NullifierPublicKey::from(nsk); - let account_id = AccountId::for_regular_private_account(&npk, *identifier); + let account_id = AccountId::for_regular_private_account(&npk); assert_eq!(account_id, pre_state.account_id, "AccountId mismatch"); assert!( @@ -56,27 +56,27 @@ pub fn compute_circuit_output( ); assert_eq!( pre_state.account, - Account::default(), - "Found new private account with non default values" + Account { nonce: *nonce, ..Account::default() }, + "New private account must be default except for the sender-chosen nonce" ); + let pre_commitment = Commitment::new(&account_id, &pre_state.account); let new_nullifier = ( - Nullifier::for_account_initialization(&account_id), + Nullifier::for_account_initialization(&pre_commitment), DUMMY_COMMITMENT_HASH, ); - let new_nonce = Nonce::private_account_nonce_init(&account_id); emit_private_output( &mut output, &mut output_index, post_state, &account_id, - &PrivateAccountKind::Regular(*identifier), + &PrivateAccountKind::Regular, ssk, epk, *view_tag, new_nullifier, - new_nonce, + *nonce, ); } InputAccountIdentity::PrivateAuthorizedUpdate { @@ -85,10 +85,9 @@ pub fn compute_circuit_output( ssk, nsk, membership_proof, - identifier, } => { let npk = NullifierPublicKey::from(nsk); - let account_id = AccountId::for_regular_private_account(&npk, *identifier); + let account_id = AccountId::for_regular_private_account(&npk); assert_eq!(account_id, pre_state.account_id, "AccountId mismatch"); assert!( @@ -109,7 +108,7 @@ pub fn compute_circuit_output( &mut output_index, post_state, &account_id, - &PrivateAccountKind::Regular(*identifier), + &PrivateAccountKind::Regular, ssk, epk, *view_tag, @@ -122,38 +121,38 @@ pub fn compute_circuit_output( view_tag, npk, ssk, - identifier, + nonce, } => { - let account_id = AccountId::for_regular_private_account(npk, *identifier); + let account_id = AccountId::for_regular_private_account(npk); assert_eq!(account_id, pre_state.account_id, "AccountId mismatch"); assert_eq!( pre_state.account, - Account::default(), - "Found new private account with non default values", + Account { nonce: *nonce, ..Account::default() }, + "New private account must be default except for the sender-chosen nonce", ); assert!( !pre_state.is_authorized, "Found new private account marked as authorized." ); + let pre_commitment = Commitment::new(&account_id, &pre_state.account); let new_nullifier = ( - Nullifier::for_account_initialization(&account_id), + Nullifier::for_account_initialization(&pre_commitment), DUMMY_COMMITMENT_HASH, ); - let new_nonce = Nonce::private_account_nonce_init(&account_id); emit_private_output( &mut output, &mut output_index, post_state, &account_id, - &PrivateAccountKind::Regular(*identifier), + &PrivateAccountKind::Regular, ssk, epk, *view_tag, new_nullifier, - new_nonce, + *nonce, ); } InputAccountIdentity::PrivatePdaInit { @@ -161,7 +160,7 @@ pub fn compute_circuit_output( view_tag, npk: _, ssk, - identifier, + nonce, seed: _, } => { // The npk-to-account_id binding is established upstream in @@ -176,17 +175,17 @@ pub fn compute_circuit_output( ); assert_eq!( pre_state.account, - Account::default(), - "New private PDA must be default" + Account { nonce: *nonce, ..Account::default() }, + "New private PDA must be default except for the sender-chosen nonce" ); - let new_nullifier = ( - Nullifier::for_account_initialization(&pre_state.account_id), - DUMMY_COMMITMENT_HASH, - ); - let new_nonce = Nonce::private_account_nonce_init(&pre_state.account_id); - let account_id = pre_state.account_id; + let pre_commitment = Commitment::new(&account_id, &pre_state.account); + let new_nullifier = ( + Nullifier::for_account_initialization(&pre_commitment), + DUMMY_COMMITMENT_HASH, + ); + let (authority_program_id, seed) = pda_seed_by_position .get(&pos) .expect("PrivatePdaInit position must be in pda_seed_by_position"); @@ -198,13 +197,12 @@ pub fn compute_circuit_output( &PrivateAccountKind::Pda { program_id: *authority_program_id, seed: *seed, - identifier: *identifier, }, ssk, epk, *view_tag, new_nullifier, - new_nonce, + *nonce, ); } InputAccountIdentity::PrivatePdaUpdate { @@ -213,7 +211,6 @@ pub fn compute_circuit_output( ssk, nsk, membership_proof, - identifier, seed: external_seed, } => { // With an external seed the binding comes from the circuit input and the @@ -246,7 +243,6 @@ pub fn compute_circuit_output( &PrivateAccountKind::Pda { program_id: *authority_program_id, seed: *seed, - identifier: *identifier, }, ssk, epk,