Merge pull request #414 from logos-blockchain/arjentix/protect-from-pda-griefing-attack

feat: protect from public pda griefing attacks

examples/program_deployment/methods/guest/src/bin/hello_world.rs

@@ -1,6 +1,4 @@
-use nssa_core::program::{
-    AccountPostState, DEFAULT_PROGRAM_ID, ProgramInput, ProgramOutput, read_nssa_inputs,
-};
+use nssa_core::program::{AccountPostState, Claim, ProgramInput, ProgramOutput, read_nssa_inputs};
 
 // Hello-world example program.
 //
@@ -45,13 +43,7 @@ fn main() {
 
     // Wrap the post state account values inside a `AccountPostState` instance.
     // This is used to forward the account claiming request if any
-    let post_state = if post_account.program_owner == DEFAULT_PROGRAM_ID {
-        // This produces a claim request
-        AccountPostState::new_claimed(post_account)
-    } else {
-        // This doesn't produce a claim request
-        AccountPostState::new(post_account)
-    };
+    let post_state = AccountPostState::new_claimed_if_default(post_account, Claim::Authorized);
 
     // The output is a proposed state difference. It will only succeed if the pre states coincide
     // with the previous values of the accounts, and the transition to the post states conforms

examples/program_deployment/methods/guest/src/bin/hello_world_with_authorization.rs

@@ -1,6 +1,4 @@
-use nssa_core::program::{
-    AccountPostState, DEFAULT_PROGRAM_ID, ProgramInput, ProgramOutput, read_nssa_inputs,
-};
+use nssa_core::program::{AccountPostState, Claim, ProgramInput, ProgramOutput, read_nssa_inputs};
 
 // Hello-world with authorization example program.
 //
@@ -52,13 +50,7 @@ fn main() {
 
     // Wrap the post state account values inside a `AccountPostState` instance.
     // This is used to forward the account claiming request if any
-    let post_state = if post_account.program_owner == DEFAULT_PROGRAM_ID {
-        // This produces a claim request
-        AccountPostState::new_claimed(post_account)
-    } else {
-        // This doesn't produce a claim request
-        AccountPostState::new(post_account)
-    };
+    let post_state = AccountPostState::new_claimed_if_default(post_account, Claim::Authorized);
 
     // The output is a proposed state difference. It will only succeed if the pre states coincide
     // with the previous values of the accounts, and the transition to the post states conforms

examples/program_deployment/methods/guest/src/bin/hello_world_with_move_function.rs

@@ -1,8 +1,6 @@
 use nssa_core::{
-    account::{Account, AccountWithMetadata, Data},
-    program::{
-        AccountPostState, DEFAULT_PROGRAM_ID, ProgramInput, ProgramOutput, read_nssa_inputs,
-    },
+    account::{AccountWithMetadata, Data},
+    program::{AccountPostState, Claim, ProgramInput, ProgramOutput, read_nssa_inputs},
 };
 
 // Hello-world with write + move_data example program.
@@ -26,16 +24,6 @@ const MOVE_DATA_FUNCTION_ID: u8 = 1;
 
 type Instruction = (u8, Vec<u8>);
 
-fn build_post_state(post_account: Account) -> AccountPostState {
-    if post_account.program_owner == DEFAULT_PROGRAM_ID {
-        // This produces a claim request
-        AccountPostState::new_claimed(post_account)
-    } else {
-        // This doesn't produce a claim request
-        AccountPostState::new(post_account)
-    }
-}
-
 fn write(pre_state: AccountWithMetadata, greeting: &[u8]) -> AccountPostState {
     // Construct the post state account values
     let post_account = {
@@ -48,7 +36,7 @@ fn write(pre_state: AccountWithMetadata, greeting: &[u8]) -> AccountPostState {
         this
     };
 
-    build_post_state(post_account)
+    AccountPostState::new_claimed_if_default(post_account, Claim::Authorized)
 }
 
 fn move_data(from_pre: AccountWithMetadata, to_pre: AccountWithMetadata) -> Vec<AccountPostState> {
@@ -58,7 +46,7 @@ fn move_data(from_pre: AccountWithMetadata, to_pre: AccountWithMetadata) -> Vec<
     let from_post = {
         let mut this = from_pre.account;
         this.data = Data::default();
-        build_post_state(this)
+        AccountPostState::new_claimed_if_default(this, Claim::Authorized)
     };
 
     let to_post = {
@@ -68,7 +56,7 @@ fn move_data(from_pre: AccountWithMetadata, to_pre: AccountWithMetadata) -> Vec<
         this.data = bytes
             .try_into()
             .expect("Data should fit within the allowed limits");
-        build_post_state(this)
+        AccountPostState::new_claimed_if_default(this, Claim::Authorized)
     };
 
     vec![from_post, to_post]

integration_tests/tests/program_deployment.rs

@@ -11,10 +11,13 @@ use integration_tests::{
     NSSA_PROGRAM_FOR_TEST_DATA_CHANGER, TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext,
 };
 use log::info;
-use nssa::{AccountId, program::Program};
+use nssa::program::Program;
 use sequencer_service_rpc::RpcClient as _;
 use tokio::test;
-use wallet::cli::Command;
+use wallet::cli::{
+    Command, SubcommandReturnValue,
+    account::{AccountSubcommand, NewSubcommand},
+};
 
 #[test]
 async fn deploy_and_execute_program() -> Result<()> {
@@ -40,14 +43,31 @@ async fn deploy_and_execute_program() -> Result<()> {
     // logic)
     let bytecode = std::fs::read(binary_filepath)?;
     let data_changer = Program::new(bytecode)?;
-    let account_id: AccountId = "11".repeat(16).parse()?;
+
+    let SubcommandReturnValue::RegisterAccount { account_id } = wallet::cli::execute_subcommand(
+        ctx.wallet_mut(),
+        Command::Account(AccountSubcommand::New(NewSubcommand::Public {
+            cci: None,
+            label: None,
+        })),
+    )
+    .await?
+    else {
+        panic!("Expected RegisterAccount return value");
+    };
+
+    let nonces = ctx.wallet().get_accounts_nonces(vec![account_id]).await?;
+    let private_key = ctx
+        .wallet()
+        .get_account_public_signing_key(account_id)
+        .unwrap();
     let message = nssa::public_transaction::Message::try_new(
         data_changer.id(),
         vec![account_id],
-        vec![],
+        nonces,
         vec![0],
     )?;
-    let witness_set = nssa::public_transaction::WitnessSet::for_message(&message, &[]);
+    let witness_set = nssa::public_transaction::WitnessSet::for_message(&message, &[private_key]);
     let transaction = nssa::PublicTransaction::new(message, witness_set);
     let _response = ctx
         .sequencer_client()
@@ -64,7 +84,7 @@ async fn deploy_and_execute_program() -> Result<()> {
     assert_eq!(post_state_account.program_owner, data_changer.id());
     assert_eq!(post_state_account.balance, 0);
     assert_eq!(post_state_account.data.as_ref(), &[0]);
-    assert_eq!(post_state_account.nonce.0, 0);
+    assert_eq!(post_state_account.nonce.0, 1);
 
     info!("Successfully deployed and executed program");
 

integration_tests/tests/wallet_ffi.rs

@@ -924,7 +924,7 @@ fn test_wallet_ffi_transfer_deshielded() -> Result<()> {
     let home = tempfile::tempdir()?;
     let wallet_ffi_handle = new_wallet_ffi_with_test_context_config(&ctx, home.path())?;
     let from: FfiBytes32 = (&ctx.ctx().existing_private_accounts()[0]).into();
-    let to = FfiBytes32::from_bytes([37; 32]);
+    let to: FfiBytes32 = (&ctx.ctx().existing_public_accounts()[0]).into();
     let amount: [u8; 16] = 100_u128.to_le_bytes();
 
     let mut transfer_result = FfiTransferResult::default();
@@ -967,7 +967,7 @@ fn test_wallet_ffi_transfer_deshielded() -> Result<()> {
     };
 
     assert_eq!(from_balance, 9900);
-    assert_eq!(to_balance, 100);
+    assert_eq!(to_balance, 10100);
 
     unsafe {
         wallet_ffi_free_transfer_result(&raw mut transfer_result);

nssa/core/src/program.rs

@@ -22,7 +22,7 @@ pub struct ProgramInput<T> {
 /// Each program can derive up to `2^256` unique account IDs by choosing different
 /// seeds. PDAs allow programs to control namespaced account identifiers without
 /// collisions between programs.
-#[derive(Debug, Serialize, Deserialize, Clone, Eq, PartialEq)]
+#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
 pub struct PdaSeed([u8; 32]);
 
 impl PdaSeed {
@@ -91,11 +91,26 @@ impl ChainedCall {
 /// A post state may optionally request that the executing program
 /// becomes the owner of the account (a “claim”). This is used to signal
 /// that the program intends to take ownership of the account.
-#[derive(Debug, Serialize, Deserialize, Clone)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 #[cfg_attr(any(feature = "host", test), derive(PartialEq, Eq))]
 pub struct AccountPostState {
     account: Account,
-    claim: bool,
+    claim: Option<Claim>,
+}
+
+/// A claim request for an account, indicating that the executing program intends to take ownership
+/// of the account.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+pub enum Claim {
+    /// The program requests ownership of the account which was authorized by the signer.
+    ///
+    /// Note that it's possible to successfully execute program outputting [`AccountPostState`] with
+    /// `is_authorized == false` and `claim == Some(Claim::Authorized)`.
+    /// This will give no error if program had authorization in pre state and may be useful
+    /// if program decides to give up authorization for a chained call.
+    Authorized,
+    /// The program requests ownership of the account through a PDA.
+    Pda(PdaSeed),
 }
 
 impl AccountPostState {
@@ -105,7 +120,7 @@ impl AccountPostState {
     pub const fn new(account: Account) -> Self {
         Self {
             account,
-            claim: false,
+            claim: None,
         }
     }
 
@@ -113,25 +128,27 @@ impl AccountPostState {
     /// This indicates that the executing program intends to claim the
     /// account as its own and is allowed to mutate it.
     #[must_use]
-    pub const fn new_claimed(account: Account) -> Self {
+    pub const fn new_claimed(account: Account, claim: Claim) -> Self {
         Self {
             account,
-            claim: true,
+            claim: Some(claim),
         }
     }
 
     /// Creates a post state that requests ownership of the account
     /// if the account's program owner is the default program ID.
     #[must_use]
-    pub fn new_claimed_if_default(account: Account) -> Self {
-        let claim = account.program_owner == DEFAULT_PROGRAM_ID;
-        Self { account, claim }
+    pub fn new_claimed_if_default(account: Account, claim: Claim) -> Self {
+        let is_default_owner = account.program_owner == DEFAULT_PROGRAM_ID;
+        Self {
+            account,
+            claim: is_default_owner.then_some(claim),
+        }
     }
 
-    /// Returns `true` if this post state requests that the account
-    /// be claimed (owned) by the executing program.
+    /// Returns whether this post state requires a claim.
     #[must_use]
-    pub const fn requires_claim(&self) -> bool {
+    pub const fn required_claim(&self) -> Option<Claim> {
         self.claim
     }
 
@@ -142,6 +159,7 @@ impl AccountPostState {
     }
 
     /// Returns the underlying account.
+    #[must_use]
     pub const fn account_mut(&mut self) -> &mut Account {
         &mut self.account
     }
@@ -598,10 +616,10 @@ mod tests {
             nonce: 10_u128.into(),
         };
 
-        let account_post_state = AccountPostState::new_claimed(account.clone());
+        let account_post_state = AccountPostState::new_claimed(account.clone(), Claim::Authorized);
 
         assert_eq!(account, account_post_state.account);
-        assert!(account_post_state.requires_claim());
+        assert_eq!(account_post_state.required_claim(), Some(Claim::Authorized));
     }
 
     #[test]
@@ -616,7 +634,7 @@ mod tests {
         let account_post_state = AccountPostState::new(account.clone());
 
         assert_eq!(account, account_post_state.account);
-        assert!(!account_post_state.requires_claim());
+        assert!(account_post_state.required_claim().is_none());
     }
 
     #[test]

nssa/src/public_transaction/transaction.rs

@@ -4,7 +4,7 @@ use borsh::{BorshDeserialize, BorshSerialize};
 use log::debug;
 use nssa_core::{
     account::{Account, AccountId, AccountWithMetadata},
-    program::{BlockId, ChainedCall, DEFAULT_PROGRAM_ID, validate_execution},
+    program::{BlockId, ChainedCall, Claim, DEFAULT_PROGRAM_ID, validate_execution},
 };
 use sha2::{Digest as _, digest::FixedOutput as _};
 
@@ -157,6 +157,10 @@ impl PublicTransaction {
                 &chained_call.pda_seeds,
             );
 
+            let is_authorized = |account_id: &AccountId| {
+                signer_account_ids.contains(account_id) || authorized_pdas.contains(account_id)
+            };
+
             for pre in &program_output.pre_states {
                 let account_id = pre.account_id;
                 // Check that the program output pre_states coincide with the values in the public
@@ -172,10 +176,8 @@ impl PublicTransaction {
 
                 // Check that authorization flags are consistent with the provided ones or
                 // authorized by program through the PDA mechanism
-                let is_authorized = signer_account_ids.contains(&account_id)
-                    || authorized_pdas.contains(&account_id);
                 ensure!(
-                    pre.is_authorized == is_authorized,
+                    pre.is_authorized == is_authorized(&account_id),
                     NssaError::InvalidProgramBehavior
                 );
             }
@@ -199,17 +201,35 @@ impl PublicTransaction {
                 NssaError::OutOfValidityWindow
             );
 
-            for post in program_output
-                .post_states
-                .iter_mut()
-                .filter(|post| post.requires_claim())
-            {
+            for (i, post) in program_output.post_states.iter_mut().enumerate() {
+                let Some(claim) = post.required_claim() else {
+                    continue;
+                };
                 // The invoked program can only claim accounts with default program id.
-                if post.account().program_owner == DEFAULT_PROGRAM_ID {
-                    post.account_mut().program_owner = chained_call.program_id;
-                } else {
-                    return Err(NssaError::InvalidProgramBehavior);
+                ensure!(
+                    post.account().program_owner == DEFAULT_PROGRAM_ID,
+                    NssaError::InvalidProgramBehavior
+                );
+
+                let account_id = program_output.pre_states[i].account_id;
+
+                match claim {
+                    Claim::Authorized => {
+                        // The program can only claim accounts that were authorized by the signer.
+                        ensure!(
+                            is_authorized(&account_id),
+                            NssaError::InvalidProgramBehavior
+                        );
+                    }
+                    Claim::Pda(seed) => {
+                        // The program can only claim accounts that correspond to the PDAs it is
+                        // authorized to claim.
+                        let pda = AccountId::from((&chained_call.program_id, &seed));
+                        ensure!(account_id == pda, NssaError::InvalidProgramBehavior);
+                    }
                 }
+
+                post.account_mut().program_owner = chained_call.program_id;
             }
 
             // Update the state diff

nssa/src/state.rs

@@ -456,16 +456,19 @@ pub mod tests {
     fn transfer_transaction(
         from: AccountId,
         from_key: &PrivateKey,
-        nonce: u128,
+        from_nonce: u128,
         to: AccountId,
+        to_key: &PrivateKey,
+        to_nonce: u128,
         balance: u128,
     ) -> PublicTransaction {
         let account_ids = vec![from, to];
-        let nonces = vec![Nonce(nonce)];
+        let nonces = vec![Nonce(from_nonce), Nonce(to_nonce)];
         let program_id = Program::authenticated_transfer_program().id();
         let message =
             public_transaction::Message::try_new(program_id, account_ids, nonces, balance).unwrap();
-        let witness_set = public_transaction::WitnessSet::for_message(&message, &[from_key]);
+        let witness_set =
+            public_transaction::WitnessSet::for_message(&message, &[from_key, to_key]);
         PublicTransaction::new(message, witness_set)
     }
 
@@ -567,17 +570,18 @@ pub mod tests {
         let initial_data = [(account_id, 100)];
         let mut state = V03State::new_with_genesis_accounts(&initial_data, &[]);
         let from = account_id;
-        let to = AccountId::new([2; 32]);
+        let to_key = PrivateKey::try_new([2; 32]).unwrap();
+        let to = AccountId::from(&PublicKey::new_from_private_key(&to_key));
         assert_eq!(state.get_account_by_id(to), Account::default());
         let balance_to_move = 5;
 
-        let tx = transfer_transaction(from, &key, 0, to, balance_to_move);
+        let tx = transfer_transaction(from, &key, 0, to, &to_key, 0, balance_to_move);
         state.transition_from_public_transaction(&tx, 1).unwrap();
 
         assert_eq!(state.get_account_by_id(from).balance, 95);
         assert_eq!(state.get_account_by_id(to).balance, 5);
         assert_eq!(state.get_account_by_id(from).nonce, Nonce(1));
-        assert_eq!(state.get_account_by_id(to).nonce, Nonce(0));
+        assert_eq!(state.get_account_by_id(to).nonce, Nonce(1));
     }
 
     #[test]
@@ -588,11 +592,12 @@ pub mod tests {
         let mut state = V03State::new_with_genesis_accounts(&initial_data, &[]);
         let from = account_id;
         let from_key = key;
-        let to = AccountId::new([2; 32]);
+        let to_key = PrivateKey::try_new([2; 32]).unwrap();
+        let to = AccountId::from(&PublicKey::new_from_private_key(&to_key));
         let balance_to_move = 101;
         assert!(state.get_account_by_id(from).balance < balance_to_move);
 
-        let tx = transfer_transaction(from, &from_key, 0, to, balance_to_move);
+        let tx = transfer_transaction(from, &from_key, 0, to, &to_key, 0, balance_to_move);
         let result = state.transition_from_public_transaction(&tx, 1);
 
         assert!(matches!(result, Err(NssaError::ProgramExecutionFailed(_))));
@@ -613,16 +618,17 @@ pub mod tests {
         let from = account_id2;
         let from_key = key2;
         let to = account_id1;
+        let to_key = key1;
         assert_ne!(state.get_account_by_id(to), Account::default());
         let balance_to_move = 8;
 
-        let tx = transfer_transaction(from, &from_key, 0, to, balance_to_move);
+        let tx = transfer_transaction(from, &from_key, 0, to, &to_key, 0, balance_to_move);
         state.transition_from_public_transaction(&tx, 1).unwrap();
 
         assert_eq!(state.get_account_by_id(from).balance, 192);
         assert_eq!(state.get_account_by_id(to).balance, 108);
         assert_eq!(state.get_account_by_id(from).nonce, Nonce(1));
-        assert_eq!(state.get_account_by_id(to).nonce, Nonce(0));
+        assert_eq!(state.get_account_by_id(to).nonce, Nonce(1));
     }
 
     #[test]
@@ -633,21 +639,38 @@ pub mod tests {
         let account_id2 = AccountId::from(&PublicKey::new_from_private_key(&key2));
         let initial_data = [(account_id1, 100)];
         let mut state = V03State::new_with_genesis_accounts(&initial_data, &[]);
-        let account_id3 = AccountId::new([3; 32]);
+        let key3 = PrivateKey::try_new([3; 32]).unwrap();
+        let account_id3 = AccountId::from(&PublicKey::new_from_private_key(&key3));
         let balance_to_move = 5;
 
-        let tx = transfer_transaction(account_id1, &key1, 0, account_id2, balance_to_move);
+        let tx = transfer_transaction(
+            account_id1,
+            &key1,
+            0,
+            account_id2,
+            &key2,
+            0,
+            balance_to_move,
+        );
         state.transition_from_public_transaction(&tx, 1).unwrap();
         let balance_to_move = 3;
-        let tx = transfer_transaction(account_id2, &key2, 0, account_id3, balance_to_move);
+        let tx = transfer_transaction(
+            account_id2,
+            &key2,
+            1,
+            account_id3,
+            &key3,
+            0,
+            balance_to_move,
+        );
         state.transition_from_public_transaction(&tx, 1).unwrap();
 
         assert_eq!(state.get_account_by_id(account_id1).balance, 95);
         assert_eq!(state.get_account_by_id(account_id2).balance, 2);
         assert_eq!(state.get_account_by_id(account_id3).balance, 3);
         assert_eq!(state.get_account_by_id(account_id1).nonce, Nonce(1));
-        assert_eq!(state.get_account_by_id(account_id2).nonce, Nonce(1));
-        assert_eq!(state.get_account_by_id(account_id3).nonce, Nonce(0));
+        assert_eq!(state.get_account_by_id(account_id2).nonce, Nonce(2));
+        assert_eq!(state.get_account_by_id(account_id3).nonce, Nonce(1));
     }
 
     #[test]
@@ -2212,15 +2235,14 @@ pub mod tests {
     #[test]
     fn claiming_mechanism() {
         let program = Program::authenticated_transfer_program();
-        let key = PrivateKey::try_new([1; 32]).unwrap();
-        let account_id = AccountId::from(&PublicKey::new_from_private_key(&key));
+        let from_key = PrivateKey::try_new([1; 32]).unwrap();
+        let from = AccountId::from(&PublicKey::new_from_private_key(&from_key));
         let initial_balance = 100;
-        let initial_data = [(account_id, initial_balance)];
+        let initial_data = [(from, initial_balance)];
         let mut state =
             V03State::new_with_genesis_accounts(&initial_data, &[]).with_test_programs();
-        let from = account_id;
-        let from_key = key;
-        let to = AccountId::new([2; 32]);
+        let to_key = PrivateKey::try_new([2; 32]).unwrap();
+        let to = AccountId::from(&PublicKey::new_from_private_key(&to_key));
         let amount: u128 = 37;
 
         // Check the recipient is an uninitialized account
@@ -2229,17 +2251,19 @@ pub mod tests {
         let expected_recipient_post = Account {
             program_owner: program.id(),
             balance: amount,
+            nonce: Nonce(1),
             ..Account::default()
         };
 
         let message = public_transaction::Message::try_new(
             program.id(),
             vec![from, to],
-            vec![Nonce(0)],
+            vec![Nonce(0), Nonce(0)],
             amount,
         )
         .unwrap();
-        let witness_set = public_transaction::WitnessSet::for_message(&message, &[&from_key]);
+        let witness_set =
+            public_transaction::WitnessSet::for_message(&message, &[&from_key, &to_key]);
         let tx = PublicTransaction::new(message, witness_set);
 
         state.transition_from_public_transaction(&tx, 1).unwrap();
@@ -2249,6 +2273,58 @@ pub mod tests {
         assert_eq!(recipient_post, expected_recipient_post);
     }
 
+    #[test]
+    fn unauthorized_public_account_claiming_fails() {
+        let program = Program::authenticated_transfer_program();
+        let account_key = PrivateKey::try_new([9; 32]).unwrap();
+        let account_id = AccountId::from(&PublicKey::new_from_private_key(&account_key));
+        let mut state = V03State::new_with_genesis_accounts(&[], &[]);
+
+        assert_eq!(state.get_account_by_id(account_id), Account::default());
+
+        let message =
+            public_transaction::Message::try_new(program.id(), vec![account_id], vec![], 0_u128)
+                .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, 1);
+
+        assert!(matches!(result, Err(NssaError::ProgramExecutionFailed(_))));
+        assert_eq!(state.get_account_by_id(account_id), Account::default());
+    }
+
+    #[test]
+    fn authorized_public_account_claiming_succeeds() {
+        let program = Program::authenticated_transfer_program();
+        let account_key = PrivateKey::try_new([10; 32]).unwrap();
+        let account_id = AccountId::from(&PublicKey::new_from_private_key(&account_key));
+        let mut state = V03State::new_with_genesis_accounts(&[], &[]);
+
+        assert_eq!(state.get_account_by_id(account_id), Account::default());
+
+        let message = public_transaction::Message::try_new(
+            program.id(),
+            vec![account_id],
+            vec![Nonce(0)],
+            0_u128,
+        )
+        .unwrap();
+        let witness_set = public_transaction::WitnessSet::for_message(&message, &[&account_key]);
+        let tx = PublicTransaction::new(message, witness_set);
+
+        state.transition_from_public_transaction(&tx, 1).unwrap();
+
+        assert_eq!(
+            state.get_account_by_id(account_id),
+            Account {
+                program_owner: program.id(),
+                nonce: Nonce(1),
+                ..Account::default()
+            }
+        );
+    }
+
     #[test]
     fn public_chained_call() {
         let program = Program::chain_caller();
@@ -2382,15 +2458,14 @@ pub mod tests {
         // program and not the chained_caller program.
         let chain_caller = Program::chain_caller();
         let auth_transfer = Program::authenticated_transfer_program();
-        let key = PrivateKey::try_new([1; 32]).unwrap();
-        let account_id = AccountId::from(&PublicKey::new_from_private_key(&key));
+        let from_key = PrivateKey::try_new([1; 32]).unwrap();
+        let from = AccountId::from(&PublicKey::new_from_private_key(&from_key));
         let initial_balance = 100;
-        let initial_data = [(account_id, initial_balance)];
+        let initial_data = [(from, initial_balance)];
         let mut state =
             V03State::new_with_genesis_accounts(&initial_data, &[]).with_test_programs();
-        let from = account_id;
-        let from_key = key;
-        let to = AccountId::new([2; 32]);
+        let to_key = PrivateKey::try_new([2; 32]).unwrap();
+        let to = AccountId::from(&PublicKey::new_from_private_key(&to_key));
         let amount: u128 = 37;
 
         // Check the recipient is an uninitialized account
@@ -2400,6 +2475,7 @@ pub mod tests {
             // The expected program owner is the authenticated transfer program
             program_owner: auth_transfer.id(),
             balance: amount,
+            nonce: Nonce(1),
             ..Account::default()
         };
 
@@ -2415,11 +2491,12 @@ pub mod tests {
             chain_caller.id(),
             vec![to, from], // The chain_caller program permutes the account order in the chain
             // call
-            vec![Nonce(0)],
+            vec![Nonce(0), Nonce(0)],
             instruction,
         )
         .unwrap();
-        let witness_set = public_transaction::WitnessSet::for_message(&message, &[&from_key]);
+        let witness_set =
+            public_transaction::WitnessSet::for_message(&message, &[&from_key, &to_key]);
         let tx = PublicTransaction::new(message, witness_set);
 
         state.transition_from_public_transaction(&tx, 1).unwrap();
@@ -2430,6 +2507,88 @@ pub mod tests {
         assert_eq!(to_post, expected_to_post);
     }
 
+    #[test]
+    fn unauthorized_public_account_claiming_fails_when_executed_privately() {
+        let program = Program::authenticated_transfer_program();
+        let account_id = AccountId::new([11; 32]);
+        let public_account = AccountWithMetadata::new(Account::default(), false, account_id);
+
+        let result = execute_and_prove(
+            vec![public_account],
+            Program::serialize_instruction(0_u128).unwrap(),
+            vec![0],
+            vec![],
+            vec![],
+            vec![],
+            &program.into(),
+        );
+
+        assert!(matches!(result, Err(NssaError::ProgramProveFailed(_))));
+    }
+
+    #[test]
+    fn authorized_public_account_claiming_succeeds_when_executed_privately() {
+        let program = Program::authenticated_transfer_program();
+        let program_id = program.id();
+        let sender_keys = test_private_account_keys_1();
+        let sender_private_account = Account {
+            program_owner: program_id,
+            balance: 100,
+            ..Account::default()
+        };
+        let sender_commitment = Commitment::new(&sender_keys.npk(), &sender_private_account);
+        let mut state =
+            V03State::new_with_genesis_accounts(&[], std::slice::from_ref(&sender_commitment));
+        let sender_pre = AccountWithMetadata::new(sender_private_account, true, &sender_keys.npk());
+        let recipient_private_key = PrivateKey::try_new([2; 32]).unwrap();
+        let recipient_account_id =
+            AccountId::from(&PublicKey::new_from_private_key(&recipient_private_key));
+        let recipient_pre =
+            AccountWithMetadata::new(Account::default(), true, recipient_account_id);
+        let esk = [5; 32];
+        let shared_secret = SharedSecretKey::new(&esk, &sender_keys.vpk());
+        let epk = EphemeralPublicKey::from_scalar(esk);
+
+        let (output, proof) = execute_and_prove(
+            vec![sender_pre, recipient_pre],
+            Program::serialize_instruction(37_u128).unwrap(),
+            vec![1, 0],
+            vec![(sender_keys.npk(), shared_secret)],
+            vec![sender_keys.nsk],
+            vec![state.get_proof_for_commitment(&sender_commitment)],
+            &program.into(),
+        )
+        .unwrap();
+
+        let message = Message::try_from_circuit_output(
+            vec![recipient_account_id],
+            vec![Nonce(0)],
+            vec![(sender_keys.npk(), sender_keys.vpk(), epk)],
+            output,
+        )
+        .unwrap();
+
+        let witness_set = WitnessSet::for_message(&message, proof, &[&recipient_private_key]);
+        let tx = PrivacyPreservingTransaction::new(message, witness_set);
+
+        state
+            .transition_from_privacy_preserving_transaction(&tx, 1)
+            .unwrap();
+
+        let nullifier = Nullifier::for_account_update(&sender_commitment, &sender_keys.nsk);
+        assert!(state.private_state.1.contains(&nullifier));
+
+        assert_eq!(
+            state.get_account_by_id(recipient_account_id),
+            Account {
+                program_owner: program_id,
+                balance: 37,
+                nonce: Nonce(1),
+                ..Account::default()
+            }
+        );
+    }
+
     #[test_case::test_case(1; "single call")]
     #[test_case::test_case(2; "two calls")]
     fn private_chained_call(number_of_calls: u32) {
@@ -2547,85 +2706,6 @@ pub mod tests {
         );
     }
 
-    #[test]
-    fn pda_mechanism_with_pinata_token_program() {
-        let pinata_token = Program::pinata_token();
-        let token = Program::token();
-
-        let pinata_definition_id = AccountId::new([1; 32]);
-        let pinata_token_definition_id = AccountId::new([2; 32]);
-        // Total supply of pinata token will be in an account under a PDA.
-        let pinata_token_holding_id = AccountId::from((&pinata_token.id(), &PdaSeed::new([0; 32])));
-        let winner_token_holding_id = AccountId::new([3; 32]);
-
-        let expected_winner_account_holding = token_core::TokenHolding::Fungible {
-            definition_id: pinata_token_definition_id,
-            balance: 150,
-        };
-        let expected_winner_token_holding_post = Account {
-            program_owner: token.id(),
-            data: Data::from(&expected_winner_account_holding),
-            ..Account::default()
-        };
-
-        let mut state = V03State::new_with_genesis_accounts(&[], &[]);
-        state.add_pinata_token_program(pinata_definition_id);
-
-        // Execution of the token program to create new token for the pinata token
-        // definition and supply accounts
-        let total_supply: u128 = 10_000_000;
-        let instruction = token_core::Instruction::NewFungibleDefinition {
-            name: String::from("PINATA"),
-            total_supply,
-        };
-        let message = public_transaction::Message::try_new(
-            token.id(),
-            vec![pinata_token_definition_id, pinata_token_holding_id],
-            vec![],
-            instruction,
-        )
-        .unwrap();
-        let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
-        let tx = PublicTransaction::new(message, witness_set);
-        state.transition_from_public_transaction(&tx, 1).unwrap();
-
-        // Execution of winner's token holding account initialization
-        let instruction = token_core::Instruction::InitializeAccount;
-        let message = public_transaction::Message::try_new(
-            token.id(),
-            vec![pinata_token_definition_id, winner_token_holding_id],
-            vec![],
-            instruction,
-        )
-        .unwrap();
-        let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
-        let tx = PublicTransaction::new(message, witness_set);
-        state.transition_from_public_transaction(&tx, 1).unwrap();
-
-        // Submit a solution to the pinata program to claim the prize
-        let solution: u128 = 989_106;
-        let message = public_transaction::Message::try_new(
-            pinata_token.id(),
-            vec![
-                pinata_definition_id,
-                pinata_token_holding_id,
-                winner_token_holding_id,
-            ],
-            vec![],
-            solution,
-        )
-        .unwrap();
-        let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
-        let tx = PublicTransaction::new(message, witness_set);
-        state.transition_from_public_transaction(&tx, 1).unwrap();
-
-        let winner_token_holding_post = state.get_account_by_id(winner_token_holding_id);
-        assert_eq!(
-            winner_token_holding_post,
-            expected_winner_token_holding_post
-        );
-    }
-
     #[test]
     fn claiming_mechanism_cannot_claim_initialied_accounts() {
         let claimer = Program::claimer();
@@ -2769,6 +2849,53 @@ pub mod tests {
         assert!(state.private_state.1.contains(&nullifier));
     }
 
+    #[test]
+    fn private_unauthorized_uninitialized_account_can_still_be_claimed() {
+        let mut state = V03State::new_with_genesis_accounts(&[], &[]).with_test_programs();
+
+        let private_keys = test_private_account_keys_1();
+        // This is intentional: claim authorization was introduced to protect public accounts,
+        // especially PDAs. Private PDAs are not useful in practice because there is no way to
+        // operate them without the corresponding private keys, so unauthorized private claiming
+        // remains allowed.
+        let unauthorized_account =
+            AccountWithMetadata::new(Account::default(), false, &private_keys.npk());
+
+        let program = Program::claimer();
+        let esk = [5; 32];
+        let shared_secret = SharedSecretKey::new(&esk, &private_keys.vpk());
+        let epk = EphemeralPublicKey::from_scalar(esk);
+
+        let (output, proof) = execute_and_prove(
+            vec![unauthorized_account],
+            Program::serialize_instruction(0_u128).unwrap(),
+            vec![2],
+            vec![(private_keys.npk(), shared_secret)],
+            vec![],
+            vec![None],
+            &program.into(),
+        )
+        .unwrap();
+
+        let message = Message::try_from_circuit_output(
+            vec![],
+            vec![],
+            vec![(private_keys.npk(), private_keys.vpk(), epk)],
+            output,
+        )
+        .unwrap();
+
+        let witness_set = WitnessSet::for_message(&message, proof, &[]);
+        let tx = PrivacyPreservingTransaction::new(message, witness_set);
+
+        state
+            .transition_from_privacy_preserving_transaction(&tx, 1)
+            .unwrap();
+
+        let nullifier = Nullifier::for_account_initialization(&private_keys.npk());
+        assert!(state.private_state.1.contains(&nullifier));
+    }
+
     #[test]
     fn private_account_claimed_then_used_without_init_flag_should_fail() {
         let mut state = V03State::new_with_genesis_accounts(&[], &[]).with_test_programs();

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,

programs/amm/src/new_definition.rs

@@ -2,11 +2,11 @@ use std::num::NonZeroU128;
 
 use amm_core::{
     PoolDefinition, compute_liquidity_token_pda, compute_liquidity_token_pda_seed,
-    compute_pool_pda, compute_vault_pda,
+    compute_pool_pda, compute_pool_pda_seed, compute_vault_pda, compute_vault_pda_seed,
 };
 use nssa_core::{
     account::{Account, AccountWithMetadata, Data},
-    program::{AccountPostState, ChainedCall, ProgramId},
+    program::{AccountPostState, ChainedCall, Claim, ProgramId},
 };
 
 #[expect(clippy::too_many_arguments, reason = "TODO: Fix later")]
@@ -108,36 +108,52 @@ pub fn new_definition(
     };
 
     pool_post.data = Data::from(&pool_post_definition);
-    let pool_post = AccountPostState::new_claimed_if_default(pool_post);
+    let pool_pda_seed = compute_pool_pda_seed(definition_token_a_id, definition_token_b_id);
+    let pool_post = AccountPostState::new_claimed_if_default(pool_post, Claim::Pda(pool_pda_seed));
 
     let token_program_id = user_holding_a.account.program_owner;
 
     // Chain call for Token A (user_holding_a -> Vault_A)
+    let vault_a_seed = compute_vault_pda_seed(pool.account_id, definition_token_a_id);
+    let vault_a_authorized = AccountWithMetadata {
+        is_authorized: true,
+        ..vault_a.clone()
+    };
     let call_token_a = ChainedCall::new(
         token_program_id,
-        vec![user_holding_a.clone(), vault_a.clone()],
+        vec![user_holding_a.clone(), vault_a_authorized],
         &token_core::Instruction::Transfer {
             amount_to_transfer: token_a_amount.into(),
         },
-    );
+    )
+    .with_pda_seeds(vec![vault_a_seed]);
+
     // Chain call for Token B (user_holding_b -> Vault_B)
+    let vault_b_seed = compute_vault_pda_seed(pool.account_id, definition_token_b_id);
+    let vault_b_authorized = AccountWithMetadata {
+        is_authorized: true,
+        ..vault_b.clone()
+    };
     let call_token_b = ChainedCall::new(
         token_program_id,
-        vec![user_holding_b.clone(), vault_b.clone()],
+        vec![user_holding_b.clone(), vault_b_authorized],
         &token_core::Instruction::Transfer {
             amount_to_transfer: token_b_amount.into(),
         },
-    );
-
-    let mut pool_lp_auth = pool_definition_lp.clone();
-    pool_lp_auth.is_authorized = true;
+    )
+    .with_pda_seeds(vec![vault_b_seed]);
 
+    let pool_lp_pda_seed = compute_liquidity_token_pda_seed(pool.account_id);
+    let pool_lp_authorized = AccountWithMetadata {
+        is_authorized: true,
+        ..pool_definition_lp.clone()
+    };
     let call_token_lp = ChainedCall::new(
         token_program_id,
-        vec![pool_lp_auth, user_holding_lp.clone()],
+        vec![pool_lp_authorized, user_holding_lp.clone()],
         &instruction,
     )
-    .with_pda_seeds(vec![compute_liquidity_token_pda_seed(pool.account_id)]);
+    .with_pda_seeds(vec![pool_lp_pda_seed]);
 
     let chained_calls = vec![call_token_lp, call_token_b, call_token_a];
 

programs/amm/src/tests.rs

@@ -1,4 +1,4 @@
-use std::num::NonZero;
+use std::{num::NonZero, vec};
 
 use amm_core::{
     PoolDefinition, compute_liquidity_token_pda, compute_liquidity_token_pda_seed,
@@ -1756,7 +1756,7 @@ impl AccountsForExeTests {
                 definition_id: IdForExeTests::token_lp_definition_id(),
                 balance: BalanceForExeTests::lp_supply_init(),
             }),
-            nonce: 0_u128.into(),
+            nonce: 1_u128.into(),
         }
     }
 
@@ -1801,7 +1801,7 @@ impl AccountsForExeTests {
                 definition_id: IdForExeTests::token_lp_definition_id(),
                 balance: 0,
             }),
-            nonce: 0_u128.into(),
+            nonce: 1.into(),
         }
     }
 }
@@ -2799,7 +2799,7 @@ fn simple_amm_new_definition_inactive_initialized_pool_and_uninit_user_lp() {
             IdForExeTests::user_token_b_id(),
             IdForExeTests::user_token_lp_id(),
         ],
-        vec![0_u128.into(), 0_u128.into()],
+        vec![0_u128.into(), 0_u128.into(), 0_u128.into()],
         instruction,
     )
     .unwrap();
@@ -2809,6 +2809,7 @@ fn simple_amm_new_definition_inactive_initialized_pool_and_uninit_user_lp() {
         &[
             &PrivateKeysForTests::user_token_a_key(),
             &PrivateKeysForTests::user_token_b_key(),
+            &PrivateKeysForTests::user_token_lp_key(),
         ],
     );
 
@@ -2955,7 +2956,7 @@ fn simple_amm_new_definition_uninitialized_pool() {
             IdForExeTests::user_token_b_id(),
             IdForExeTests::user_token_lp_id(),
         ],
-        vec![0_u128.into(), 0_u128.into()],
+        vec![0_u128.into(), 0_u128.into(), 0_u128.into()],
         instruction,
     )
     .unwrap();
@@ -2965,6 +2966,7 @@ fn simple_amm_new_definition_uninitialized_pool() {
         &[
             &PrivateKeysForTests::user_token_a_key(),
             &PrivateKeysForTests::user_token_b_key(),
+            &PrivateKeysForTests::user_token_lp_key(),
         ],
     );
 

programs/associated_token_account/src/create.rs

@@ -1,6 +1,6 @@
 use nssa_core::{
     account::{Account, AccountWithMetadata},
-    program::{AccountPostState, ChainedCall, ProgramId},
+    program::{AccountPostState, ChainedCall, Claim, ProgramId},
 };
 
 pub fn create_associated_token_account(
@@ -11,7 +11,7 @@ pub fn create_associated_token_account(
 ) -> (Vec<AccountPostState>, Vec<ChainedCall>) {
     // No authorization check needed: create is idempotent, so anyone can call it safely.
     let token_program_id = token_definition.account.program_owner;
-    ata_core::verify_ata_and_get_seed(
+    let ata_seed = ata_core::verify_ata_and_get_seed(
         &ata_account,
         &owner,
         token_definition.account_id,
@@ -22,7 +22,7 @@ pub fn create_associated_token_account(
     if ata_account.account != Account::default() {
         return (
             vec![
-                AccountPostState::new_claimed_if_default(owner.account.clone()),
+                AccountPostState::new_claimed_if_default(owner.account.clone(), Claim::Authorized),
                 AccountPostState::new(token_definition.account.clone()),
                 AccountPostState::new(ata_account.account.clone()),
             ],
@@ -31,14 +31,20 @@ pub fn create_associated_token_account(
     }
 
     let post_states = vec![
-        AccountPostState::new_claimed_if_default(owner.account.clone()),
+        AccountPostState::new_claimed_if_default(owner.account.clone(), Claim::Authorized),
         AccountPostState::new(token_definition.account.clone()),
         AccountPostState::new(ata_account.account.clone()),
     ];
+    let ata_account_auth = AccountWithMetadata {
+        is_authorized: true,
+        ..ata_account.clone()
+    };
     let chained_call = ChainedCall::new(
         token_program_id,
-        vec![token_definition.clone(), ata_account.clone()],
+        vec![token_definition.clone(), ata_account_auth],
         &token_core::Instruction::InitializeAccount,
-    );
+    )
+    .with_pda_seeds(vec![ata_seed]);
+
     (post_states, vec![chained_call])
 }

programs/token/core/src/lib.rs

@@ -10,23 +10,23 @@ pub enum Instruction {
     /// Transfer tokens from sender to recipient.
     ///
     /// Required accounts:
-    /// - Sender's Token Holding account (authorized),
-    /// - Recipient's Token Holding account.
+    /// - Sender's Token Holding account (initialized, authorized),
+    /// - Recipient's Token Holding account (initialized or authorized and uninitialized).
     Transfer { amount_to_transfer: u128 },
 
     /// Create a new fungible token definition without metadata.
     ///
     /// Required accounts:
-    /// - Token Definition account (uninitialized),
-    /// - Token Holding account (uninitialized).
+    /// - Token Definition account (uninitialized, authorized),
+    /// - Token Holding account (uninitialized, authorized).
     NewFungibleDefinition { name: String, total_supply: u128 },
 
     /// Create a new fungible or non-fungible token definition with metadata.
     ///
     /// Required accounts:
-    /// - Token Definition account (uninitialized),
-    /// - Token Holding account (uninitialized),
-    /// - Token Metadata account (uninitialized).
+    /// - Token Definition account (uninitialized, authorized),
+    /// - Token Holding account (uninitialized, authorized),
+    /// - Token Metadata account (uninitialized, authorized).
     NewDefinitionWithMetadata {
         new_definition: NewTokenDefinition,
         /// Boxed to avoid large enum variant size.
@@ -36,29 +36,29 @@ pub enum Instruction {
     /// Initialize a token holding account for a given token definition.
     ///
     /// Required accounts:
-    /// - Token Definition account (initialized),
-    /// - Token Holding account (uninitialized),
+    /// - Token Definition account (initialized, any authorization),
+    /// - Token Holding account (uninitialized, authorized),
     InitializeAccount,
 
     /// Burn tokens from the holder's account.
     ///
     /// Required accounts:
-    /// - Token Definition account (initialized),
-    /// - Token Holding account (authorized).
+    /// - Token Definition account (initialized, any authorization),
+    /// - Token Holding account (initialized, authorized).
     Burn { amount_to_burn: u128 },
 
     /// Mint new tokens to the holder's account.
     ///
     /// Required accounts:
-    /// - Token Definition account (authorized),
-    /// - Token Holding account (uninitialized or initialized).
+    /// - Token Definition account (initialized, authorized),
+    /// - Token Holding account (uninitialized or authorized and initialized).
     Mint { amount_to_mint: u128 },
 
     /// Print a new NFT from the master copy.
     ///
     /// Required accounts:
     /// - NFT Master Token Holding account (authorized),
-    /// - NFT Printed Copy Token Holding account (uninitialized).
+    /// - NFT Printed Copy Token Holding account (uninitialized, authorized).
     PrintNft,
 }
 

programs/token/src/initialize.rs

@@ -1,6 +1,6 @@
 use nssa_core::{
     account::{Account, AccountWithMetadata, Data},
-    program::AccountPostState,
+    program::{AccountPostState, Claim},
 };
 use token_core::{TokenDefinition, TokenHolding};
 
@@ -30,6 +30,6 @@ pub fn initialize_account(
 
     vec![
         AccountPostState::new(definition_post),
-        AccountPostState::new_claimed(account_to_initialize),
+        AccountPostState::new_claimed(account_to_initialize, Claim::Authorized),
     ]
 }

programs/token/src/mint.rs

@@ -1,6 +1,6 @@
 use nssa_core::{
     account::{Account, AccountWithMetadata, Data},
-    program::AccountPostState,
+    program::{AccountPostState, Claim},
 };
 use token_core::{TokenDefinition, TokenHolding};
 
@@ -67,6 +67,6 @@ pub fn mint(
 
     vec![
         AccountPostState::new(definition_post),
-        AccountPostState::new_claimed_if_default(holding_post),
+        AccountPostState::new_claimed_if_default(holding_post, Claim::Authorized),
     ]
 }

programs/token/src/new_definition.rs

@@ -1,6 +1,6 @@
 use nssa_core::{
     account::{Account, AccountWithMetadata, Data},
-    program::AccountPostState,
+    program::{AccountPostState, Claim},
 };
 use token_core::{
     NewTokenDefinition, NewTokenMetadata, TokenDefinition, TokenHolding, TokenMetadata,
@@ -42,8 +42,8 @@ pub fn new_fungible_definition(
     holding_target_account_post.data = Data::from(&token_holding);
 
     vec![
-        AccountPostState::new_claimed(definition_target_account_post),
-        AccountPostState::new_claimed(holding_target_account_post),
+        AccountPostState::new_claimed(definition_target_account_post, Claim::Authorized),
+        AccountPostState::new_claimed(holding_target_account_post, Claim::Authorized),
     ]
 }
 
@@ -119,8 +119,8 @@ pub fn new_definition_with_metadata(
     metadata_target_account_post.data = Data::from(&token_metadata);
 
     vec![
-        AccountPostState::new_claimed(definition_target_account_post),
-        AccountPostState::new_claimed(holding_target_account_post),
-        AccountPostState::new_claimed(metadata_target_account_post),
+        AccountPostState::new_claimed(definition_target_account_post, Claim::Authorized),
+        AccountPostState::new_claimed(holding_target_account_post, Claim::Authorized),
+        AccountPostState::new_claimed(metadata_target_account_post, Claim::Authorized),
     ]
 }

programs/token/src/print_nft.rs

@@ -1,6 +1,6 @@
 use nssa_core::{
     account::{Account, AccountWithMetadata, Data},
-    program::AccountPostState,
+    program::{AccountPostState, Claim},
 };
 use token_core::TokenHolding;
 
@@ -50,6 +50,6 @@ pub fn print_nft(
 
     vec![
         AccountPostState::new(master_account_post),
-        AccountPostState::new_claimed(printed_account_post),
+        AccountPostState::new_claimed(printed_account_post, Claim::Authorized),
     ]
 }

programs/token/src/tests.rs

@@ -5,7 +5,10 @@
     reason = "We don't care about it in tests"
 )]
 
-use nssa_core::account::{Account, AccountId, AccountWithMetadata, Data};
+use nssa_core::{
+    account::{Account, AccountId, AccountWithMetadata, Data},
+    program::Claim,
+};
 use token_core::{
     MetadataStandard, NewTokenDefinition, NewTokenMetadata, TokenDefinition, TokenHolding,
 };
@@ -851,7 +854,7 @@ fn mint_uninit_holding_success() {
         *holding_post.account(),
         AccountForTests::init_mint().account
     );
-    assert!(holding_post.requires_claim());
+    assert_eq!(holding_post.required_claim(), Some(Claim::Authorized));
 }
 
 #[test]

programs/token/src/transfer.rs

@@ -1,6 +1,6 @@
 use nssa_core::{
     account::{Account, AccountWithMetadata, Data},
-    program::AccountPostState,
+    program::{AccountPostState, Claim},
 };
 use token_core::TokenHolding;
 
@@ -106,6 +106,6 @@ pub fn transfer(
 
     vec![
         AccountPostState::new(sender_post),
-        AccountPostState::new_claimed_if_default(recipient_post),
+        AccountPostState::new_claimed_if_default(recipient_post, Claim::Authorized),
     ]
 }

test_program_methods/guest/src/bin/chain_caller.rs

@@ -38,7 +38,7 @@ fn main() {
             program_id: auth_transfer_id,
             instruction_data: instruction_data.clone(),
             pre_states: vec![running_sender_pre.clone(), running_recipient_pre.clone()], /* <- Account order permutation here */
-            pda_seeds: pda_seed.iter().cloned().collect(),
+            pda_seeds: pda_seed.iter().copied().collect(),
         };
         chained_calls.push(new_chained_call);
 

test_program_methods/guest/src/bin/changer_claimer.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};
 
 type Instruction = (Option<Vec<u8>>, bool);
 
@@ -28,7 +28,7 @@ fn main() {
 
     // Claim or not based on the boolean flag
     let post_state = if should_claim {
-        AccountPostState::new_claimed(account_post)
+        AccountPostState::new_claimed(account_post, Claim::Authorized)
     } else {
         AccountPostState::new(account_post)
     };

test_program_methods/guest/src/bin/claimer.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};
 
 type Instruction = ();
 
@@ -15,7 +15,7 @@ fn main() {
         return;
     };
 
-    let account_post = AccountPostState::new_claimed(pre.account.clone());
+    let account_post = AccountPostState::new_claimed(pre.account.clone(), Claim::Authorized);
 
     ProgramOutput::new(instruction_words, vec![pre], vec![account_post]).write();
 }

test_program_methods/guest/src/bin/data_changer.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};
 
 type Instruction = Vec<u8>;
 
@@ -25,7 +25,10 @@ fn main() {
     ProgramOutput::new(
         instruction_words,
         vec![pre],
-        vec![AccountPostState::new_claimed(account_post)],
+        vec![AccountPostState::new_claimed(
+            account_post,
+            Claim::Authorized,
+        )],
     )
     .write();
 }

wallet/src/program_facades/amm.rs

@@ -58,18 +58,21 @@ impl Amm<'_> {
             user_holding_lp,
         ];
 
-        let nonces = self
+        let mut nonces = self
             .0
             .get_accounts_nonces(vec![user_holding_a, user_holding_b])
             .await
             .map_err(ExecutionFailureKind::SequencerError)?;
 
+        let mut private_keys = Vec::new();
+
         let signing_key_a = self
             .0
             .storage
             .user_data
             .get_pub_account_signing_key(user_holding_a)
             .ok_or(ExecutionFailureKind::KeyNotFoundError)?;
+        private_keys.push(signing_key_a);
 
         let signing_key_b = self
             .0
@@ -77,6 +80,26 @@ impl Amm<'_> {
             .user_data
             .get_pub_account_signing_key(user_holding_b)
             .ok_or(ExecutionFailureKind::KeyNotFoundError)?;
+        private_keys.push(signing_key_b);
+
+        if let Some(signing_key_lp) = self
+            .0
+            .storage
+            .user_data
+            .get_pub_account_signing_key(user_holding_lp)
+        {
+            private_keys.push(signing_key_lp);
+            let lp_nonces = self
+                .0
+                .get_accounts_nonces(vec![user_holding_lp])
+                .await
+                .map_err(ExecutionFailureKind::SequencerError)?;
+            nonces.extend(lp_nonces);
+        } else {
+            println!(
+                "Liquidity pool tokens receiver's account ({user_holding_lp}) private key not found in wallet. Proceeding with only liquidity provider's keys."
+            );
+        }
 
         let message = nssa::public_transaction::Message::try_new(
             program.id(),
@@ -86,10 +109,8 @@ impl Amm<'_> {
         )
         .unwrap();
 
-        let witness_set = nssa::public_transaction::WitnessSet::for_message(
-            &message,
-            &[signing_key_a, signing_key_b],
-        );
+        let witness_set =
+            nssa::public_transaction::WitnessSet::for_message(&message, &private_keys);
 
         let tx = nssa::PublicTransaction::new(message, witness_set);
 

wallet/src/program_facades/native_token_transfer/public.rs

@@ -23,24 +23,40 @@ impl NativeTokenTransfer<'_> {
             .map_err(ExecutionFailureKind::SequencerError)?;
 
         if balance >= balance_to_move {
-            let nonces = self
+            let account_ids = vec![from, to];
+            let program_id = Program::authenticated_transfer_program().id();
+
+            let mut nonces = self
                 .0
                 .get_accounts_nonces(vec![from])
                 .await
                 .map_err(ExecutionFailureKind::SequencerError)?;
 
-            let account_ids = vec![from, to];
-            let program_id = Program::authenticated_transfer_program().id();
-            let message =
-                Message::try_new(program_id, account_ids, nonces, balance_to_move).unwrap();
-
-            let signing_key = self.0.storage.user_data.get_pub_account_signing_key(from);
-
-            let Some(signing_key) = signing_key else {
+            let mut private_keys = Vec::new();
+            let from_signing_key = self.0.storage.user_data.get_pub_account_signing_key(from);
+            let Some(from_signing_key) = from_signing_key else {
                 return Err(ExecutionFailureKind::KeyNotFoundError);
             };
+            private_keys.push(from_signing_key);
 
-            let witness_set = WitnessSet::for_message(&message, &[signing_key]);
+            let to_signing_key = self.0.storage.user_data.get_pub_account_signing_key(to);
+            if let Some(to_signing_key) = to_signing_key {
+                private_keys.push(to_signing_key);
+                let to_nonces = self
+                    .0
+                    .get_accounts_nonces(vec![to])
+                    .await
+                    .map_err(ExecutionFailureKind::SequencerError)?;
+                nonces.extend(to_nonces);
+            } else {
+                println!(
+                    "Receiver's account ({to}) private key not found in wallet. Proceeding with only sender's key."
+                );
+            }
+
+            let message =
+                Message::try_new(program_id, account_ids, nonces, balance_to_move).unwrap();
+            let witness_set = WitnessSet::for_message(&message, &private_keys);
 
             let tx = PublicTransaction::new(message, witness_set);
 

wallet/src/program_facades/token.rs

@@ -19,15 +19,36 @@ impl Token<'_> {
         let account_ids = vec![definition_account_id, supply_account_id];
         let program_id = nssa::program::Program::token().id();
         let instruction = Instruction::NewFungibleDefinition { name, total_supply };
+        let nonces = self
+            .0
+            .get_accounts_nonces(account_ids.clone())
+            .await
+            .map_err(ExecutionFailureKind::SequencerError)?;
         let message = nssa::public_transaction::Message::try_new(
             program_id,
             account_ids,
-            vec![],
+            nonces,
             instruction,
         )
         .unwrap();
 
-        let witness_set = nssa::public_transaction::WitnessSet::for_message(&message, &[]);
+        let def_private_key = self
+            .0
+            .storage
+            .user_data
+            .get_pub_account_signing_key(definition_account_id)
+            .ok_or(ExecutionFailureKind::KeyNotFoundError)?;
+        let supply_private_key = self
+            .0
+            .storage
+            .user_data
+            .get_pub_account_signing_key(supply_account_id)
+            .ok_or(ExecutionFailureKind::KeyNotFoundError)?;
+
+        let witness_set = nssa::public_transaction::WitnessSet::for_message(
+            &message,
+            &[def_private_key, supply_private_key],
+        );
 
         let tx = nssa::PublicTransaction::new(message, witness_set);
 
@@ -138,11 +159,40 @@ impl Token<'_> {
         let instruction = Instruction::Transfer {
             amount_to_transfer: amount,
         };
-        let nonces = self
+        let mut nonces = self
             .0
             .get_accounts_nonces(vec![sender_account_id])
             .await
             .map_err(ExecutionFailureKind::SequencerError)?;
+
+        let mut private_keys = Vec::new();
+        let sender_sk = self
+            .0
+            .storage
+            .user_data
+            .get_pub_account_signing_key(sender_account_id)
+            .ok_or(ExecutionFailureKind::KeyNotFoundError)?;
+        private_keys.push(sender_sk);
+
+        if let Some(recipient_sk) = self
+            .0
+            .storage
+            .user_data
+            .get_pub_account_signing_key(recipient_account_id)
+        {
+            private_keys.push(recipient_sk);
+            let recipient_nonces = self
+                .0
+                .get_accounts_nonces(vec![recipient_account_id])
+                .await
+                .map_err(ExecutionFailureKind::SequencerError)?;
+            nonces.extend(recipient_nonces);
+        } else {
+            println!(
+                "Receiver's account ({recipient_account_id}) private key not found in wallet. Proceeding with only sender's key."
+            );
+        }
+
         let message = nssa::public_transaction::Message::try_new(
             program_id,
             account_ids,
@@ -150,17 +200,8 @@ impl Token<'_> {
             instruction,
         )
         .unwrap();
-
-        let Some(signing_key) = self
-            .0
-            .storage
-            .user_data
-            .get_pub_account_signing_key(sender_account_id)
-        else {
-            return Err(ExecutionFailureKind::KeyNotFoundError);
-        };
         let witness_set =
-            nssa::public_transaction::WitnessSet::for_message(&message, &[signing_key]);
+            nssa::public_transaction::WitnessSet::for_message(&message, &private_keys);
 
         let tx = nssa::PublicTransaction::new(message, witness_set);
 
@@ -477,11 +518,40 @@ impl Token<'_> {
             amount_to_mint: amount,
         };
 
-        let nonces = self
+        let mut nonces = self
             .0
             .get_accounts_nonces(vec![definition_account_id])
             .await
             .map_err(ExecutionFailureKind::SequencerError)?;
+
+        let mut private_keys = Vec::new();
+        let definition_sk = self
+            .0
+            .storage
+            .user_data
+            .get_pub_account_signing_key(definition_account_id)
+            .ok_or(ExecutionFailureKind::KeyNotFoundError)?;
+        private_keys.push(definition_sk);
+
+        if let Some(holder_sk) = self
+            .0
+            .storage
+            .user_data
+            .get_pub_account_signing_key(holder_account_id)
+        {
+            private_keys.push(holder_sk);
+            let recipient_nonces = self
+                .0
+                .get_accounts_nonces(vec![holder_account_id])
+                .await
+                .map_err(ExecutionFailureKind::SequencerError)?;
+            nonces.extend(recipient_nonces);
+        } else {
+            println!(
+                "Holder's account ({holder_account_id}) private key not found in wallet. Proceeding with only definition's key."
+            );
+        }
+
         let message = nssa::public_transaction::Message::try_new(
             Program::token().id(),
             account_ids,
@@ -489,17 +559,8 @@ impl Token<'_> {
             instruction,
         )
         .unwrap();
-
-        let Some(signing_key) = self
-            .0
-            .storage
-            .user_data
-            .get_pub_account_signing_key(definition_account_id)
-        else {
-            return Err(ExecutionFailureKind::KeyNotFoundError);
-        };
         let witness_set =
-            nssa::public_transaction::WitnessSet::for_message(&message, &[signing_key]);
+            nssa::public_transaction::WitnessSet::for_message(&message, &private_keys);
 
         let tx = nssa::PublicTransaction::new(message, witness_set);