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.

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<usize, (NullifierPublicKey, ViewingPublicKey, Identifier)>,
-    authorized_accounts: HashSet<AccountId>,
+    /// The set containing regular accounts authorized at the root of the call-tree, remaining
+    /// authorized throughout all calls.
+    globally_authorized: HashSet<AccountId>,
 }
 
 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<ProgramId>,
+        caller: CallerData,
         caller_pda_seeds: &[PdaSeed],
         output_pre_states: Vec<AccountWithMetadata>,
         output_post_states: Vec<AccountPostState>,
-    ) {
+    ) -> HashSet<AccountId> {
+        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<usize, (ProgramId, PdaSeed)>,
     private_pda_by_position: &HashMap<usize, (NullifierPublicKey, ViewingPublicKey, Identifier)>,
-    authorized_accounts: &mut HashSet<AccountId>,
+    globally_authorized: &HashSet<AccountId>,
+    caller_authorized: &HashSet<AccountId>,
     pre_account_id: AccountId,
     pre_state_position: usize,
     caller_program_id: Option<ProgramId>,
     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)
 }

lee/state_machine/core/src/program/mod.rs

@@ -198,6 +198,12 @@ impl AccountId {
     }
 }
 
+#[derive(Debug)]
+pub struct CallerData {
+    pub program_id: Option<ProgramId>,
+    pub authorized_accounts: HashSet<AccountId>,
+}
+
 #[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq)]
 pub struct ChainedCall {
     /// The program ID of the program to execute.

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(

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<bool>)> = 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::<bool>)),
+        ))
+        .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<bool>)> = 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

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<ProgramId>,
-    authorized_accounts: HashSet<AccountId>,
-}
-
 fn authenticate_public_transaction_signers(
     tx: &PublicTransaction,
     state: &V03State,

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::<Instruction>();
+
+    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();
+}

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<bool>)>,
+);
+
+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::<Instruction>();
+
+    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();
+}