feat: protect from public pda griefing attacks

program_methods/guest/src/bin/authenticated_transfer.rs

@@ -1,13 +1,13 @@
 use nssa_core::{
     account::{Account, AccountWithMetadata},
     program::{
-        AccountPostState, DEFAULT_PROGRAM_ID, ProgramInput, ProgramOutput, read_nssa_inputs,
+        AccountPostState, Claim, DEFAULT_PROGRAM_ID, ProgramInput, ProgramOutput, read_nssa_inputs,
     },
 };
 
 /// Initializes a default account under the ownership of this program.
 fn initialize_account(pre_state: AccountWithMetadata) -> AccountPostState {
-    let account_to_claim = AccountPostState::new_claimed(pre_state.account);
+    let account_to_claim = AccountPostState::new_claimed(pre_state.account, Claim::Authorized);
     let is_authorized = pre_state.is_authorized;
 
     // Continue only if the account to claim has default values
@@ -52,7 +52,7 @@ fn transfer(
 
         // Claim recipient account if it has default program owner
         if recipient_post_account.program_owner == DEFAULT_PROGRAM_ID {
-            AccountPostState::new_claimed(recipient_post_account)
+            AccountPostState::new_claimed(recipient_post_account, Claim::Authorized)
         } else {
             AccountPostState::new(recipient_post_account)
         }

program_methods/guest/src/bin/pinata.rs

@@ -1,4 +1,4 @@
-use nssa_core::program::{AccountPostState, ProgramInput, ProgramOutput, read_nssa_inputs};
+use nssa_core::program::{AccountPostState, Claim, ProgramInput, ProgramOutput, read_nssa_inputs};
 use risc0_zkvm::sha::{Impl, Sha256 as _};
 
 const PRIZE: u128 = 150;
@@ -82,7 +82,7 @@ fn main() {
         instruction_words,
         vec![pinata, winner],
         vec![
-            AccountPostState::new_claimed_if_default(pinata_post),
+            AccountPostState::new_claimed_if_default(pinata_post, Claim::Authorized),
             AccountPostState::new(winner_post),
         ],
     )

program_methods/guest/src/bin/privacy_preserving_circuit.rs

@@ -10,8 +10,8 @@ use nssa_core::{
     account::{Account, AccountId, AccountWithMetadata, Nonce},
     compute_digest_for_path,
     program::{
-        AccountPostState, ChainedCall, DEFAULT_PROGRAM_ID, MAX_NUMBER_CHAINED_CALLS, ProgramId,
-        ProgramOutput, ValidityWindow, validate_execution,
+        AccountPostState, ChainedCall, Claim, DEFAULT_PROGRAM_ID, MAX_NUMBER_CHAINED_CALLS,
+        ProgramId, ProgramOutput, ValidityWindow, validate_execution,
     },
 };
 use risc0_zkvm::{guest::env, serde::to_vec};
@@ -25,7 +25,11 @@ struct ExecutionState {
 
 impl ExecutionState {
     /// Validate program outputs and derive the overall execution state.
-    pub fn derive_from_outputs(program_id: ProgramId, program_outputs: Vec<ProgramOutput>) -> Self {
+    pub fn derive_from_outputs(
+        visibility_mask: &[u8],
+        program_id: ProgramId,
+        program_outputs: Vec<ProgramOutput>,
+    ) -> Self {
         let valid_from_id = program_outputs
             .iter()
             .filter_map(|output| output.validity_window.start())
@@ -102,6 +106,7 @@ impl ExecutionState {
                 &chained_call.pda_seeds,
             );
             execution_state.validate_and_sync_states(
+                visibility_mask,
                 chained_call.program_id,
                 &authorized_pdas,
                 program_output.pre_states,
@@ -134,7 +139,7 @@ impl ExecutionState {
         {
             assert_ne!(
                 post.program_owner, DEFAULT_PROGRAM_ID,
-                "Account {account_id:?} was modified but not claimed"
+                "Account {account_id} was modified but not claimed"
             );
         }
 
@@ -144,6 +149,7 @@ impl ExecutionState {
     /// Validate program pre and post states and populate the execution state.
     fn validate_and_sync_states(
         &mut self,
+        visibility_mask: &[u8],
         program_id: ProgramId,
         authorized_pdas: &HashSet<AccountId>,
         pre_states: Vec<AccountWithMetadata>,
@@ -151,14 +157,25 @@ impl ExecutionState {
     ) {
         for (pre, mut post) in pre_states.into_iter().zip(post_states) {
             let pre_account_id = pre.account_id;
+            let pre_is_authorized = pre.is_authorized;
             let post_states_entry = self.post_states.entry(pre.account_id);
             match &post_states_entry {
                 Entry::Occupied(occupied) => {
+                    #[expect(
+                        clippy::shadow_unrelated,
+                        reason = "Shadowing is intentional to use all fields"
+                    )]
+                    let AccountWithMetadata {
+                        account: pre_account,
+                        account_id: pre_account_id,
+                        is_authorized: pre_is_authorized,
+                    } = pre;
+
                     // Ensure that new pre state is the same as known post state
                     assert_eq!(
                         occupied.get(),
-                        &pre.account,
-                        "Inconsistent pre state for account {pre_account_id:?}",
+                        &pre_account,
+                        "Inconsistent pre state for account {pre_account_id}",
                     );
 
                     let previous_is_authorized = self
@@ -167,7 +184,7 @@ impl ExecutionState {
                         .find(|acc| acc.account_id == pre_account_id)
                         .map_or_else(
                             || panic!(
-                                "Pre state must exist in execution state for account {pre_account_id:?}",
+                                "Pre state must exist in execution state for account {pre_account_id}",
                             ),
                             |acc| acc.is_authorized
                         );
@@ -176,22 +193,57 @@ impl ExecutionState {
                         previous_is_authorized || authorized_pdas.contains(&pre_account_id);
 
                     assert_eq!(
-                        pre.is_authorized, is_authorized,
-                        "Inconsistent authorization for account {pre_account_id:?}",
+                        pre_is_authorized, is_authorized,
+                        "Inconsistent authorization for account {pre_account_id}",
                     );
                 }
                 Entry::Vacant(_) => {
+                    // Pre state for the initial call
                     self.pre_states.push(pre);
                 }
             }
 
-            if post.requires_claim() {
+            if let Some(claim) = post.required_claim() {
                 // The invoked program can only claim accounts with default program id.
-                if post.account().program_owner == DEFAULT_PROGRAM_ID {
-                    post.account_mut().program_owner = program_id;
+                assert_eq!(
+                    post.account().program_owner,
+                    DEFAULT_PROGRAM_ID,
+                    "Cannot claim an initialized account {pre_account_id}"
+                );
+
+                let pre_state_position = self
+                    .pre_states
+                    .iter()
+                    .position(|acc| acc.account_id == pre_account_id)
+                    .expect("Pre state must exist at this point");
+
+                let is_public_account = visibility_mask[pre_state_position] == 0;
+                if is_public_account {
+                    match claim {
+                        Claim::Authorized => {
+                            // Note: no need to check authorized pdas because we have already
+                            // checked consistency of authorization above.
+                            assert!(
+                                pre_is_authorized,
+                                "Cannot claim unauthorized account {pre_account_id}"
+                            );
+                        }
+                        Claim::Pda(seed) => {
+                            let pda = AccountId::from((&program_id, &seed));
+                            assert_eq!(
+                                pre_account_id, pda,
+                                "Invalid PDA claim for account {pre_account_id} which does not match derived PDA {pda}"
+                            );
+                        }
+                    }
                 } else {
-                    panic!("Cannot claim an initialized account {pre_account_id:?}");
+                    // We don't care about the exact claim mechanism for private accounts.
+                    // This is because the main reason to have it is to protect against PDA griefing
+                    // attacks in public execution, while private PDA doesn't make much sense
+                    // anyway.
                 }
+
+                post.account_mut().program_owner = program_id;
             }
 
             post_states_entry.insert_entry(post.into_account());
@@ -408,7 +460,8 @@ fn main() {
         program_id,
     } = env::read();
 
-    let execution_state = ExecutionState::derive_from_outputs(program_id, program_outputs);
+    let execution_state =
+        ExecutionState::derive_from_outputs(&visibility_mask, program_id, program_outputs);
 
     let output = compute_circuit_output(
         execution_state,