feat: error on unclaimed modified accounts

integration_tests/tests/pinata.rs

@@ -11,11 +11,13 @@ use tokio::test;
 use wallet::cli::{
     Command, SubcommandReturnValue,
     account::{AccountSubcommand, NewSubcommand},
-    programs::pinata::PinataProgramAgnosticSubcommand,
+    programs::{
+        native_token_transfer::AuthTransferSubcommand, pinata::PinataProgramAgnosticSubcommand,
+    },
 };
 
 #[test]
-async fn claim_pinata_to_public_account() -> Result<()> {
+async fn claim_pinata_to_existing_public_account() -> Result<()> {
     let mut ctx = TestContext::new().await?;
 
     let pinata_prize = 150;
@@ -120,8 +122,26 @@ async fn claim_pinata_to_new_private_account() -> Result<()> {
         anyhow::bail!("Expected RegisterAccount return value");
     };
 
+    let winner_account_id_formatted = format_private_account_id(&winner_account_id.to_string());
+
+    // Initialize account under auth transfer program
+    let command = Command::AuthTransfer(AuthTransferSubcommand::Init {
+        account_id: winner_account_id_formatted.clone(),
+    });
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    let new_commitment = ctx
+        .wallet()
+        .get_private_account_commitment(&winner_account_id)
+        .context("Failed to get private account commitment")?;
+    assert!(verify_commitment_is_in_state(new_commitment, ctx.sequencer_client()).await);
+
+    // Claim pinata to the new private account
     let command = Command::Pinata(PinataProgramAgnosticSubcommand::Claim {
-        to: format_private_account_id(&winner_account_id.to_string()),
+        to: winner_account_id_formatted,
     });
 
     let pinata_balance_pre = ctx

nssa/core/src/program.rs

@@ -107,6 +107,13 @@ impl AccountPostState {
         }
     }
 
+    /// Creates a post state that requests ownership of the account
+    /// if the account's program owner is the default program ID.
+    pub fn new_claimed_if_default(account: Account) -> Self {
+        let claim = account.program_owner == DEFAULT_PROGRAM_ID;
+        Self { account, claim }
+    }
+
     /// Returns `true` if this post state requests that the account
     /// be claimed (owned) by the executing program.
     pub fn requires_claim(&self) -> bool {

nssa/src/program.rs

@@ -226,6 +226,15 @@ mod tests {
             }
         }
 
+        pub fn changer_claimer() -> Self {
+            use test_program_methods::{CHANGER_CLAIMER_ELF, CHANGER_CLAIMER_ID};
+
+            Program {
+                id: CHANGER_CLAIMER_ID,
+                elf: CHANGER_CLAIMER_ELF.to_vec(),
+            }
+        }
+
         pub fn noop() -> Self {
             use test_program_methods::{NOOP_ELF, NOOP_ID};
 

nssa/src/public_transaction/transaction.rs

@@ -200,6 +200,22 @@ impl PublicTransaction {
             chain_calls_counter += 1;
         }
 
+        // Check that all modified uninitialized accounts where claimed
+        for post in state_diff.iter().filter_map(|(account_id, post)| {
+            let pre = state.get_account_by_id(account_id);
+            if pre.program_owner != DEFAULT_PROGRAM_ID {
+                return None;
+            }
+            if pre == *post {
+                return None;
+            }
+            Some(post)
+        }) {
+            if post.program_owner == DEFAULT_PROGRAM_ID {
+                return Err(NssaError::InvalidProgramBehavior);
+            }
+        }
+
         Ok(state_diff)
     }
 }

nssa/src/state.rs

@@ -504,6 +504,7 @@ pub mod tests {
             self.insert_program(Program::chain_caller());
             self.insert_program(Program::amm());
             self.insert_program(Program::claimer());
+            self.insert_program(Program::changer_claimer());
             self
         }
 
@@ -4370,6 +4371,108 @@ pub mod tests {
         assert!(matches!(res, Err(NssaError::CircuitProvingError(_))));
     }
 
+    #[test]
+    fn test_public_changer_claimer_no_data_change_no_claim_succeeds() {
+        let initial_data = [];
+        let mut state =
+            V02State::new_with_genesis_accounts(&initial_data, &[]).with_test_programs();
+        let account_id = AccountId::new([1; 32]);
+        let program_id = Program::changer_claimer().id();
+        // Don't change data (None) and don't claim (false)
+        let instruction: (Option<Vec<u8>>, bool) = (None, false);
+
+        let message =
+            public_transaction::Message::try_new(program_id, vec![account_id], vec![], instruction)
+                .unwrap();
+        let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
+        let tx = PublicTransaction::new(message, witness_set);
+
+        let result = state.transition_from_public_transaction(&tx);
+
+        // Should succeed - no changes made, no claim needed
+        assert!(result.is_ok());
+        // Account should remain default/unclaimed
+        assert_eq!(state.get_account_by_id(&account_id), Account::default());
+    }
+
+    #[test]
+    fn test_public_changer_claimer_data_change_no_claim_fails() {
+        let initial_data = [];
+        let mut state =
+            V02State::new_with_genesis_accounts(&initial_data, &[]).with_test_programs();
+        let account_id = AccountId::new([1; 32]);
+        let program_id = Program::changer_claimer().id();
+        // Change data but don't claim (false) - should fail
+        let new_data = vec![1, 2, 3, 4, 5];
+        let instruction: (Option<Vec<u8>>, bool) = (Some(new_data), false);
+
+        let message =
+            public_transaction::Message::try_new(program_id, vec![account_id], vec![], instruction)
+                .unwrap();
+        let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
+        let tx = PublicTransaction::new(message, witness_set);
+
+        let result = state.transition_from_public_transaction(&tx);
+
+        // Should fail - cannot modify data without claiming the account
+        assert!(matches!(result, Err(NssaError::InvalidProgramBehavior)));
+    }
+
+    #[test]
+    fn test_private_changer_claimer_no_data_change_no_claim_succeeds() {
+        let program = Program::changer_claimer();
+        let sender_keys = test_private_account_keys_1();
+        let private_account =
+            AccountWithMetadata::new(Account::default(), true, &sender_keys.npk());
+        // Don't change data (None) and don't claim (false)
+        let instruction: (Option<Vec<u8>>, bool) = (None, false);
+
+        let result = execute_and_prove(
+            vec![private_account],
+            Program::serialize_instruction(instruction).unwrap(),
+            vec![1],
+            vec![2],
+            vec![(
+                sender_keys.npk(),
+                SharedSecretKey::new(&[3; 32], &sender_keys.ivk()),
+            )],
+            vec![sender_keys.nsk],
+            vec![Some((0, vec![]))],
+            &program.into(),
+        );
+
+        // Should succeed - no changes made, no claim needed
+        assert!(result.is_ok());
+    }
+
+    #[test]
+    fn test_private_changer_claimer_data_change_no_claim_fails() {
+        let program = Program::changer_claimer();
+        let sender_keys = test_private_account_keys_1();
+        let private_account =
+            AccountWithMetadata::new(Account::default(), true, &sender_keys.npk());
+        // Change data but don't claim (false) - should fail
+        let new_data = vec![1, 2, 3, 4, 5];
+        let instruction: (Option<Vec<u8>>, bool) = (Some(new_data), false);
+
+        let result = execute_and_prove(
+            vec![private_account],
+            Program::serialize_instruction(instruction).unwrap(),
+            vec![1],
+            vec![2],
+            vec![(
+                sender_keys.npk(),
+                SharedSecretKey::new(&[3; 32], &sender_keys.ivk()),
+            )],
+            vec![sender_keys.nsk],
+            vec![Some((0, vec![]))],
+            &program.into(),
+        );
+
+        // Should fail - cannot modify data without claiming the account
+        assert!(matches!(result, Err(NssaError::CircuitProvingError(_))));
+    }
+
     #[test]
     fn test_malicious_authorization_changer_should_fail_in_privacy_preserving_circuit() {
         // Arrange

program_methods/guest/src/bin/pinata.rs

@@ -77,7 +77,7 @@ fn main() {
         instruction_words,
         vec![pinata, winner],
         vec![
-            AccountPostState::new(pinata_post),
+            AccountPostState::new_claimed_if_default(pinata_post),
             AccountPostState::new(winner_post),
         ],
     );

program_methods/guest/src/bin/privacy_preserving_circuit.rs

@@ -41,7 +41,7 @@ fn main() {
     env::commit(&output);
 }
 
-/// World state before and after program execution.
+/// State of the involved accounts before and after program execution.
 struct ExecutionState {
     pre_states: Vec<AccountWithMetadata>,
     post_states: HashMap<AccountId, Account>,
@@ -66,9 +66,10 @@ impl ExecutionState {
             pre_states: Vec::new(),
             post_states: HashMap::new(),
         };
-        let mut last_program_id = program_id;
+
         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() {
             assert!(
                 chain_calls_counter <= MAX_NUMBER_CHAINED_CALLS,
@@ -116,7 +117,6 @@ impl ExecutionState {
                 program_output.pre_states,
                 program_output.post_states,
             );
-            last_program_id = chained_call.program_id;
             chain_calls_counter += 1;
         }
 
@@ -125,11 +125,25 @@ impl ExecutionState {
             "Inner call without a chained call found",
         );
 
-        // Claim accounts which were not explicitly claimed during execution
-        for account in execution_state.post_states.values_mut() {
-            if account.program_owner == DEFAULT_PROGRAM_ID {
-                account.program_owner = last_program_id;
-            }
+        // Check that all modified uninitialized accounts were claimed
+        for (account_id, post) in execution_state
+            .pre_states
+            .iter()
+            .filter(|a| a.account.program_owner == DEFAULT_PROGRAM_ID)
+            .map(|a| {
+                let post = execution_state
+                    .post_states
+                    .get(&a.account_id)
+                    .expect("Post state must exist for pre state");
+                (a, post)
+            })
+            .filter(|(pre_default, post)| pre_default.account != **post)
+            .map(|(pre, post)| (pre.account_id, post))
+        {
+            assert_ne!(
+                post.program_owner, DEFAULT_PROGRAM_ID,
+                "Account {account_id:?} was modified but not claimed"
+            );
         }
 
         execution_state

test_program_methods/guest/src/bin/changer_claimer.rs

@@ -0,0 +1,38 @@
+use nssa_core::program::{AccountPostState, ProgramInput, read_nssa_inputs, write_nssa_outputs};
+
+type Instruction = (Option<Vec<u8>>, bool);
+
+/// A program that optionally modifies the account data and optionally claims it.
+fn main() {
+    let (
+        ProgramInput {
+            pre_states,
+            instruction: (data_opt, should_claim),
+        },
+        instruction_words,
+    ) = read_nssa_inputs::<Instruction>();
+
+    let [pre] = match pre_states.try_into() {
+        Ok(array) => array,
+        Err(_) => return,
+    };
+
+    let account_pre = &pre.account;
+    let mut account_post = account_pre.clone();
+
+    // Update data if provided
+    if let Some(data) = data_opt {
+        account_post.data = data
+            .try_into()
+            .expect("provided data should fit into data limit");
+    }
+
+    // Claim or not based on the boolean flag
+    let post_state = if should_claim {
+        AccountPostState::new_claimed(account_post)
+    } else {
+        AccountPostState::new(account_post)
+    };
+
+    write_nssa_outputs(instruction_words, vec![pre], vec![post_state]);
+}