From 79c43dc5a083ffb53c7dbdc46f2015591ab35348 Mon Sep 17 00:00:00 2001 From: agureev Date: Mon, 10 Aug 2026 18:27:42 +0400 Subject: [PATCH] fix!(lee): scope private PDA authorization to the callee subtree BREAKING! Before: authorized private PDAs remain authorized for the rest of the calls after. After: the authorized preivate PDAs remain authorized for the rest of the callee subtree. --- .../src/execution_state.rs | 110 +++++---- lee/state_machine/core/src/program/mod.rs | 6 + lee/state_machine/src/lib.rs | 16 ++ lee/state_machine/src/state/tests/circuit.rs | 226 ++++++++++++++++++ .../src/validated_state_diff/mod.rs | 8 +- .../guest/src/bin/non_delegating_forwarder.rs | 45 ++++ .../guest/src/bin/selective_pda_delegator.rs | 82 +++++++ 7 files changed, 445 insertions(+), 48 deletions(-) create mode 100644 lee/state_machine/test_methods/guest/src/bin/non_delegating_forwarder.rs create mode 100644 lee/state_machine/test_methods/guest/src/bin/selective_pda_delegator.rs diff --git a/lee/privacy_preserving_circuit/src/execution_state.rs b/lee/privacy_preserving_circuit/src/execution_state.rs index 334cad7c2..5ef7a8eb7 100644 --- a/lee/privacy_preserving_circuit/src/execution_state.rs +++ b/lee/privacy_preserving_circuit/src/execution_state.rs @@ -8,7 +8,7 @@ use lee_core::{ account::{Account, AccountId, AccountWithMetadata}, encryption::ViewingPublicKey, program::{ - AccountPostState, BlockValidityWindow, ChainedCall, Claim, DEFAULT_PROGRAM_ID, + AccountPostState, BlockValidityWindow, CallerData, ChainedCall, Claim, DEFAULT_PROGRAM_ID, MAX_NUMBER_CHAINED_CALLS, PdaSeed, ProgramId, ProgramOutput, TimestampValidityWindow, validate_execution, }, @@ -51,7 +51,9 @@ pub struct ExecutionState { /// `AccountId::for_private_pda(program_id, seed, npk, vpk, identifier) == /// pre_state.account_id`. private_pda_by_position: HashMap, - authorized_accounts: HashSet, + /// The set containing regular accounts authorized at the root of the call-tree, remaining + /// authorized throughout all calls. + globally_authorized: HashSet, } impl ExecutionState { @@ -112,7 +114,7 @@ impl ExecutionState { private_pda_bound_positions: HashMap::new(), pda_family_binding: HashMap::new(), private_pda_by_position, - authorized_accounts: HashSet::new(), + globally_authorized: HashSet::new(), }; let Some(first_output) = program_outputs.first() else { @@ -125,12 +127,17 @@ impl ExecutionState { pre_states: first_output.pre_states.clone(), pda_seeds: Vec::new(), }; - let mut chained_calls = VecDeque::from_iter([(initial_call, None)]); + let initial_caller_data = CallerData { + program_id: None, + authorized_accounts: HashSet::new(), + }; + let mut chained_calls = + VecDeque::<(ChainedCall, CallerData)>::from_iter([(initial_call, initial_caller_data)]); let mut program_outputs_iter = program_outputs.into_iter(); let mut chain_calls_counter = 0; - while let Some((chained_call, caller_program_id)) = chained_calls.pop_front() { + while let Some((chained_call, caller_data)) = chained_calls.pop_front() { assert!( chain_calls_counter <= MAX_NUMBER_CHAINED_CALLS, "Max chained calls depth is exceeded" @@ -166,7 +173,7 @@ impl ExecutionState { // by spoofing caller_program_id (e.g. passing caller_program_id = self_program_id // to bypass access control checks). assert_eq!( - program_output.caller_program_id, caller_program_id, + program_output.caller_program_id, caller_data.program_id, "Program output caller_program_id does not match actual caller" ); @@ -184,18 +191,25 @@ impl ExecutionState { ); } - for next_call in program_output.chained_calls.iter().rev() { - chained_calls.push_front((next_call.clone(), Some(chained_call.program_id))); - } - - execution_state.validate_and_sync_states( + let authorized_accounts = execution_state.validate_and_sync_states( account_identities, chained_call.program_id, - caller_program_id, + caller_data, &chained_call.pda_seeds, program_output.pre_states, program_output.post_states, ); + + for next_call in program_output.chained_calls.into_iter().rev() { + // Push the call with newly-authorized account set. + chained_calls.push_front(( + next_call, + CallerData { + program_id: Some(chained_call.program_id), + authorized_accounts: authorized_accounts.clone(), + }, + )); + } chain_calls_counter = chain_calls_counter.checked_add(1).expect( "Chain calls counter should not overflow as it checked before incrementing", ); @@ -246,15 +260,19 @@ impl ExecutionState { } /// Validate program pre and post states and populate the execution state. + /// + /// Return the set of authorized accounts as the result of the processed + /// call. fn validate_and_sync_states( &mut self, account_identities: &[InputAccountIdentity], program_id: ProgramId, - caller_program_id: Option, + caller: CallerData, caller_pda_seeds: &[PdaSeed], output_pre_states: Vec, output_post_states: Vec, - ) { + ) -> HashSet { + let mut authorized_output_accounts = Vec::new(); for (pre, mut post) in output_pre_states.into_iter().zip(output_post_states) { let pre_account_id = pre.account_id; let pre_is_authorized = pre.is_authorized; @@ -278,28 +296,26 @@ impl ExecutionState { "Inconsistent pre state for account {pre_account_id}", ); - let (previous_is_authorized, pre_state_position) = self + let pre_state_position = self .pre_states .iter() - .enumerate() - .find(|(_, acc)| acc.account_id == pre_account_id) - .map_or_else( - || panic!( + .position(|acc| acc.account_id == pre_account_id) + .unwrap_or_else(|| { + panic!( "Pre state must exist in execution state for account {pre_account_id}", - ), - |(pos, acc)| (acc.is_authorized, pos) - ); + ) + }); let is_authorized = resolve_authorization_and_record_bindings( &mut self.pda_family_binding, &mut self.private_pda_bound_positions, &self.private_pda_by_position, - &mut self.authorized_accounts, + &self.globally_authorized, + &caller.authorized_accounts, pre_account_id, pre_state_position, - caller_program_id, + caller.program_id, caller_pda_seeds, - previous_is_authorized, ); assert_eq!( @@ -357,10 +373,24 @@ impl ExecutionState { pre_account_id, ); } + + // If an account is regular and authorized, make it globally-authorized. + if pre_is_authorized + && !self + .private_pda_by_position + .contains_key(&pre_state_position) + { + self.globally_authorized.insert(pre_account_id); + } self.pre_states.push(pre); } } + // If an account it authorized, push it to the autorized set. + if pre_is_authorized { + authorized_output_accounts.push(pre_account_id); + } + if let Some(claim) = post.required_claim() { // The invoked program can only claim accounts with default program id. assert_eq!( @@ -444,6 +474,12 @@ impl ExecutionState { post_states_entry.insert_entry(post.into_account()); } + + caller + .authorized_accounts + .into_iter() + .chain(authorized_output_accounts) + .collect() } /// Consume self and yield the validity windows, the per-position PDA seed/program map @@ -528,11 +564,9 @@ fn bind_private_pda_position( } /// Resolve the authorization state of a `pre_state` seen again in a chained call and record -/// any resulting bindings. Returns `true` if the `pre_state` is authorized through either a -/// previously-seen authorization or a matching caller seed (under the public or private -/// derivation). When a caller seed matches, also records the `(caller, seed) → account_id` -/// family binding and, for the private form, marks the position in -/// `private_pda_bound_positions`. Only reachable when `caller_program_id.is_some()`, +/// any resulting bindings. When a caller seed matches, also records the +/// `(caller, seed) → account_id` family binding and, for the private form, marks the position +/// in `private_pda_bound_positions`. Only reachable when `caller_program_id.is_some()`, /// top-level flows have no caller-emitted seeds, so binding at top level must come from the /// claim path. Free function so callers can pass individual `&mut self.*` field borrows /// without holding a borrow on the surrounding struct's other fields. @@ -544,12 +578,12 @@ 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, - authorized_accounts: &mut HashSet, + globally_authorized: &HashSet, + caller_authorized: &HashSet, pre_account_id: AccountId, pre_state_position: usize, caller_program_id: Option, caller_pda_seeds: &[PdaSeed], - previous_is_authorized: bool, ) -> bool { let matched_caller_seed: Option<(PdaSeed, bool, ProgramId)> = caller_program_id.and_then(|caller| { @@ -580,13 +614,7 @@ fn resolve_authorization_and_record_bindings( } } - if authorized_accounts.contains(&pre_account_id) { - return true; - } - - let authorized = previous_is_authorized || matched_caller_seed.is_some(); - if authorized { - authorized_accounts.insert(pre_account_id); - } - authorized + matched_caller_seed.is_some() + || globally_authorized.contains(&pre_account_id) + || caller_authorized.contains(&pre_account_id) } diff --git a/lee/state_machine/core/src/program/mod.rs b/lee/state_machine/core/src/program/mod.rs index 770bcf2db..d8d0acd40 100644 --- a/lee/state_machine/core/src/program/mod.rs +++ b/lee/state_machine/core/src/program/mod.rs @@ -198,6 +198,12 @@ impl AccountId { } } +#[derive(Debug)] +pub struct CallerData { + pub program_id: Option, + pub authorized_accounts: HashSet, +} + #[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq)] pub struct ChainedCall { /// The program ID of the program to execute. diff --git a/lee/state_machine/src/lib.rs b/lee/state_machine/src/lib.rs index 9886127b1..4c69c12d5 100644 --- a/lee/state_machine/src/lib.rs +++ b/lee/state_machine/src/lib.rs @@ -135,6 +135,22 @@ mod test_methods { ) } + #[must_use] + pub const fn selective_pda_delegator() -> Program { + Program::new_unchecked( + test_methods::SELECTIVE_PDA_DELEGATOR_ID, + Cow::Borrowed(test_methods::SELECTIVE_PDA_DELEGATOR_ELF), + ) + } + + #[must_use] + pub const fn non_delegating_forwarder() -> Program { + Program::new_unchecked( + test_methods::NON_DELEGATING_FORWARDER_ID, + Cow::Borrowed(test_methods::NON_DELEGATING_FORWARDER_ELF), + ) + } + #[must_use] pub const fn pda_claimer() -> Program { Program::new_unchecked( diff --git a/lee/state_machine/src/state/tests/circuit.rs b/lee/state_machine/src/state/tests/circuit.rs index 758a7c85f..ccbc705ed 100644 --- a/lee/state_machine/src/state/tests/circuit.rs +++ b/lee/state_machine/src/state/tests/circuit.rs @@ -589,6 +589,232 @@ fn caller_pda_seeds_with_wrong_seed_rejects_private_pda_for_callee() { assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); } +fn sibling_declaring_delegated_pda(pda_is_authorized: bool) -> Result<(), LeeError> { + let delegator = crate::test_methods::selective_pda_delegator(); + let callee = crate::test_methods::auth_asserting_noop(); + let sibling = crate::test_methods::noop(); + let keys = test_private_account_keys_1(); + let npk = keys.npk(); + let seed = PdaSeed::new([77; 32]); + + let account_id = AccountId::for_private_pda(&delegator.id(), &seed, &npk, &keys.vpk(), 0); + let pre_state = AccountWithMetadata::new(Account::default(), false, account_id); + + let callee_id = callee.id(); + let sibling_id = sibling.id(); + let program_with_deps = ProgramWithDependencies::new( + delegator, + [(callee_id, callee), (sibling_id, sibling)].into(), + ); + + execute_and_prove( + vec![pre_state], + Program::serialize_instruction(( + seed, + seed, + callee_id, + Program::serialize_instruction(()).unwrap(), + Some((sibling_id, Some(pda_is_authorized))), + )) + .unwrap(), + vec![InputAccountIdentity::Private(PrivateWitness { + vpk: keys.vpk(), + random_seed: [0; 32], + identifier: 0, + kind: WitnessKind::Pda { binding: None }, + nullifier: NullifierWitness::Init { + npk, + commitment_root: DUMMY_COMMITMENT_HASH, + }, + })], + &program_with_deps, + ) + .map(|_| ()) +} + +#[test] +fn delegated_pda_is_not_authorized_in_sibling_call() { + let result = sibling_declaring_delegated_pda(true); + + assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); +} + +#[test] +fn sibling_call_may_declare_delegated_pda_unauthorized() { + sibling_declaring_delegated_pda(false) + .expect("a sibling declaring the delegated PDA unauthorized must be accepted"); +} + +#[test] +fn delegated_pda_stays_authorized_in_delegated_subtree() { + let delegator = crate::test_methods::selective_pda_delegator(); + let forwarder = crate::test_methods::non_delegating_forwarder(); + let callee = crate::test_methods::auth_asserting_noop(); + let keys = test_private_account_keys_1(); + let npk = keys.npk(); + let seed = PdaSeed::new([77; 32]); + + let account_id = AccountId::for_private_pda(&delegator.id(), &seed, &npk, &keys.vpk(), 0); + let pre_state = AccountWithMetadata::new(Account::default(), false, account_id); + + let forwarder_id = forwarder.id(); + let callee_id = callee.id(); + let program_with_deps = ProgramWithDependencies::new( + delegator, + [(forwarder_id, forwarder), (callee_id, callee)].into(), + ); + let no_sibling: Option<(ProgramId, Option)> = None; + + execute_and_prove( + vec![pre_state], + Program::serialize_instruction(( + seed, + seed, + forwarder_id, + Program::serialize_instruction(( + callee_id, + Program::serialize_instruction(()).unwrap(), + true, + )) + .unwrap(), + no_sibling, + )) + .unwrap(), + vec![InputAccountIdentity::Private(PrivateWitness { + vpk: keys.vpk(), + random_seed: [0; 32], + identifier: 0, + kind: WitnessKind::Pda { binding: None }, + nullifier: NullifierWitness::Init { + npk, + commitment_root: DUMMY_COMMITMENT_HASH, + }, + })], + &program_with_deps, + ) + .expect("a callee that forwards without re-delegating must keep the PDA authorized"); +} + +#[test] +fn holder_authorization_survives_across_sibling_calls() { + let delegator = crate::test_methods::selective_pda_delegator(); + let callee = crate::test_methods::auth_asserting_noop(); + let sibling = crate::test_methods::noop(); + let pda_keys = test_private_account_keys_1(); + let holder_keys = test_private_account_keys_2(); + let npk = pda_keys.npk(); + let holder_npk = holder_keys.npk(); + let seed = PdaSeed::new([77; 32]); + + let account_id = AccountId::for_private_pda(&delegator.id(), &seed, &npk, &pda_keys.vpk(), 0); + let holder_id = AccountId::for_regular_private_account(&holder_npk, &holder_keys.vpk(), 0); + let pre_state = AccountWithMetadata::new(Account::default(), false, account_id); + let holder_pre_state = AccountWithMetadata::new(Account::default(), true, holder_id); + + let callee_id = callee.id(); + let sibling_id = sibling.id(); + let program_with_deps = ProgramWithDependencies::new( + delegator, + [(callee_id, callee), (sibling_id, sibling)].into(), + ); + + execute_and_prove( + vec![pre_state, holder_pre_state], + Program::serialize_instruction(( + seed, + seed, + callee_id, + Program::serialize_instruction(()).unwrap(), + Some((sibling_id, None::)), + )) + .unwrap(), + vec![ + InputAccountIdentity::Private(PrivateWitness { + vpk: pda_keys.vpk(), + random_seed: [0; 32], + identifier: 0, + kind: WitnessKind::Pda { binding: None }, + nullifier: NullifierWitness::Init { + npk, + commitment_root: DUMMY_COMMITMENT_HASH, + }, + }), + InputAccountIdentity::Private(PrivateWitness { + vpk: holder_keys.vpk(), + random_seed: [0; 32], + identifier: 0, + kind: WitnessKind::Regular { + ask: Some(holder_keys.ask), + }, + nullifier: NullifierWitness::Init { + npk: holder_npk, + commitment_root: DUMMY_COMMITMENT_HASH, + }, + }), + ], + &program_with_deps, + ) + .expect("an account authorized by its own credential stays authorized in a sibling call"); +} + +#[test] +fn inherited_scope_passes_through_intermediate_calls() { + let delegator = crate::test_methods::selective_pda_delegator(); + let forwarder = crate::test_methods::non_delegating_forwarder(); + let callee = crate::test_methods::auth_asserting_noop(); + let keys = test_private_account_keys_1(); + let npk = keys.npk(); + let seed = PdaSeed::new([77; 32]); + + let account_id = AccountId::for_private_pda(&delegator.id(), &seed, &npk, &keys.vpk(), 0); + let pre_state = AccountWithMetadata::new(Account::default(), false, account_id); + + let forwarder_id = forwarder.id(); + let callee_id = callee.id(); + let program_with_deps = ProgramWithDependencies::new( + delegator, + [(forwarder_id, forwarder), (callee_id, callee)].into(), + ); + let no_sibling: Option<(ProgramId, Option)> = None; + let forward_through_undeclaring_call = Program::serialize_instruction(( + forwarder_id, + Program::serialize_instruction(( + callee_id, + Program::serialize_instruction(()).unwrap(), + false, + )) + .unwrap(), + true, + )) + .unwrap(); + + execute_and_prove( + vec![pre_state], + Program::serialize_instruction(( + seed, + seed, + forwarder_id, + forward_through_undeclaring_call, + no_sibling, + )) + .unwrap(), + vec![InputAccountIdentity::Private(PrivateWitness { + vpk: keys.vpk(), + random_seed: [0; 32], + identifier: 0, + kind: WitnessKind::Pda { binding: None }, + nullifier: NullifierWitness::Init { + npk, + commitment_root: DUMMY_COMMITMENT_HASH, + }, + })], + &program_with_deps, + ) + .expect( + "an account authorized in an ancestor's output stays authorized below a call that never mentions it", + ); +} + /// Exploit-scenario pin. A single `(program_id, seed)` pair can derive a family of /// `AccountId`s, one public PDA and one private PDA per distinct npk. Without the tx-wide /// family-binding check, a program could claim `PDA_alice` (`alice_npk`) and diff --git a/lee/state_machine/src/validated_state_diff/mod.rs b/lee/state_machine/src/validated_state_diff/mod.rs index 738cc5153..3598b16ba 100644 --- a/lee/state_machine/src/validated_state_diff/mod.rs +++ b/lee/state_machine/src/validated_state_diff/mod.rs @@ -7,7 +7,7 @@ use lee_core::{ BlockId, Commitment, Nullifier, PrivacyPreservingCircuitOutput, PublicAction, Timestamp, account::{Account, AccountId, AccountWithMetadata}, program::{ - ChainedCall, Claim, DEFAULT_PROGRAM_ID, ProgramId, compute_public_authorized_pdas, + CallerData, ChainedCall, Claim, DEFAULT_PROGRAM_ID, compute_public_authorized_pdas, validate_execution, }, }; @@ -470,12 +470,6 @@ impl ValidatedStateDiff { } } -#[derive(Debug)] -struct CallerData { - program_id: Option, - authorized_accounts: HashSet, -} - fn authenticate_public_transaction_signers( tx: &PublicTransaction, state: &V03State, diff --git a/lee/state_machine/test_methods/guest/src/bin/non_delegating_forwarder.rs b/lee/state_machine/test_methods/guest/src/bin/non_delegating_forwarder.rs new file mode 100644 index 000000000..81b30f666 --- /dev/null +++ b/lee/state_machine/test_methods/guest/src/bin/non_delegating_forwarder.rs @@ -0,0 +1,45 @@ +use lee_core::program::{ + AccountPostState, ChainedCall, InstructionData, ProgramId, ProgramInput, ProgramOutput, + read_lee_inputs, +}; + +type Instruction = (ProgramId, InstructionData, bool); + +fn main() { + let ( + ProgramInput { + self_program_id, + caller_program_id, + pre_states, + instruction: (callee_program_id, callee_instruction, declare_pre_states), + }, + instruction_words, + ) = read_lee_inputs::(); + + let (output_pre_states, output_post_states) = if declare_pre_states { + let post_states = pre_states + .iter() + .map(|account| AccountPostState::new(account.account.clone())) + .collect(); + (pre_states.clone(), post_states) + } else { + (Vec::new(), Vec::new()) + }; + + // Make exactly one chained call based on the input instruction with no + // pda seeds, ensuring the target PDAs are never authorized. + ProgramOutput::new( + self_program_id, + caller_program_id, + instruction_words, + output_pre_states, + output_post_states, + ) + .with_chained_calls(vec![ChainedCall { + program_id: callee_program_id, + instruction_data: callee_instruction, + pre_states, + pda_seeds: vec![], + }]) + .write(); +} diff --git a/lee/state_machine/test_methods/guest/src/bin/selective_pda_delegator.rs b/lee/state_machine/test_methods/guest/src/bin/selective_pda_delegator.rs new file mode 100644 index 000000000..54e5020f0 --- /dev/null +++ b/lee/state_machine/test_methods/guest/src/bin/selective_pda_delegator.rs @@ -0,0 +1,82 @@ +use lee_core::program::{ + AccountPostState, ChainedCall, Claim, InstructionData, PdaSeed, ProgramId, ProgramInput, + ProgramOutput, read_lee_inputs, +}; +use risc0_zkvm::serde::to_vec; + +type Instruction = ( + PdaSeed, + PdaSeed, + ProgramId, + InstructionData, + Option<(ProgramId, Option)>, +); + +fn main() { + let ( + ProgramInput { + self_program_id, + caller_program_id, + pre_states, + instruction: + (claim_seed, delegated_seed, callee_program_id, callee_instruction, sibling), + }, + instruction_words, + ) = read_lee_inputs::(); + + let Some((pda, rest)) = pre_states.split_first() else { + return; + }; + + let pda_for_callee = |is_authorized| { + let mut for_callee = pda.clone(); + for_callee.is_authorized = is_authorized; + for_callee.account.program_owner = self_program_id; + for_callee + }; + + // Send a call to the specified program with the same pre-states + // but authorized first PDA supplied. + // Push all the delegated seeds. + let mut chained_calls = vec![ChainedCall { + program_id: callee_program_id, + instruction_data: callee_instruction, + pre_states: std::iter::once(pda_for_callee(true)) + .chain(rest.iter().cloned()) + .collect(), + pda_seeds: vec![delegated_seed], + }]; + + // If sibling is present in instruction, send out a call + // with no seeds so that PDAs stay unauthorized in parallel + // branches. + if let Some((sibling_program_id, sibling_pda)) = sibling { + chained_calls.push(ChainedCall { + program_id: sibling_program_id, + instruction_data: to_vec(&()).unwrap(), + pre_states: sibling_pda.map_or_else( + || rest.to_vec(), + |is_authorized| { + std::iter::once(pda_for_callee(is_authorized)) + .chain(rest.iter().cloned()) + .collect() + }, + ), + pda_seeds: vec![], + }); + } + + ProgramOutput::new( + self_program_id, + caller_program_id, + instruction_words, + vec![pda.clone()], + // Claim first PDA supplied + vec![AccountPostState::new_claimed( + pda.account.clone(), + Claim::Pda(claim_seed), + )], + ) + .with_chained_calls(chained_calls) + .write(); +}