token-2022: [OS-SPL-ADV-00] Enable and fix account ordering on confidential transfer with split proofs (#5931)

* enable auto cloase of context accounts

* add auto-close on test

* fix account order on transfer with proof contexts

* add test to make sure that the context accounts are cleared

* add test to verify lamport destination

Author: samkim-crypto

token/program-2022-test/tests/confidential_transfer.rs

@@ -20,7 +20,8 @@ use {
                 self,
                 account_info::TransferAccountInfo,
                 instruction::{
-                    TransferSplitContextStateAccounts, TransferWithFeeSplitContextStateAccounts,
+                    CloseSplitContextStateAccounts, TransferSplitContextStateAccounts,
+                    TransferWithFeeSplitContextStateAccounts,
                 },
                 ConfidentialTransferAccount, MAXIMUM_DEPOSIT_TRANSFER_AMOUNT,
             },
@@ -2668,21 +2669,33 @@ async fn confidential_transfer_transfer_with_split_proof_contexts_in_parallel()
     let ciphertext_validity_proof_context_state_account = Keypair::new();
     let range_proof_context_state_account = Keypair::new();
 
+    let lamport_destination = Pubkey::new_unique();
+
+    let close_split_context_state_accounts = CloseSplitContextStateAccounts {
+        lamport_destination: &lamport_destination,
+        zk_token_proof_program: &zk_token_proof_program::id(),
+    };
+
     let transfer_context_state_accounts = TransferSplitContextStateAccounts {
         equality_proof: &equality_proof_context_state_account.pubkey(),
         ciphertext_validity_proof: &ciphertext_validity_proof_context_state_account.pubkey(),
         range_proof: &range_proof_context_state_account.pubkey(),
         authority: &context_state_authority.pubkey(),
         no_op_on_uninitialized_split_context_state: true,
-        close_split_context_state_accounts: None,
+        close_split_context_state_accounts: Some(close_split_context_state_accounts),
     };
 
     let equality_and_ciphertext_proof_signers = vec![
         &alice,
         &equality_proof_context_state_account,
         &ciphertext_validity_proof_context_state_account,
+        &context_state_authority,
+    ];
+    let range_proof_signers = vec![
+        &alice,
+        &range_proof_context_state_account,
+        &context_state_authority,
     ];
-    let range_proof_signers = vec![&alice, &range_proof_context_state_account];
     token
         .confidential_transfer_transfer_with_split_proofs_in_parallel(
             &alice_meta.token_account,
@@ -2724,6 +2737,27 @@ async fn confidential_transfer_transfer_with_split_proof_contexts_in_parallel()
             },
         )
         .await;
+
+    let error = token
+        .get_account(equality_proof_context_state_account.pubkey())
+        .await
+        .unwrap_err();
+    assert_eq!(error, TokenClientError::AccountNotFound);
+
+    let error = token
+        .get_account(ciphertext_validity_proof_context_state_account.pubkey())
+        .await
+        .unwrap_err();
+    assert_eq!(error, TokenClientError::AccountNotFound);
+
+    let error = token
+        .get_account(range_proof_context_state_account.pubkey())
+        .await
+        .unwrap_err();
+    assert_eq!(error, TokenClientError::AccountNotFound);
+
+    let lamport_destination = token.get_account(lamport_destination).await.unwrap();
+    assert!(lamport_destination.lamports > 0);
 }
 
 #[tokio::test]
@@ -3043,6 +3077,13 @@ async fn confidential_transfer_transfer_with_fee_and_split_proof_context_in_para
     let fee_ciphertext_validity_proof_context_state_account = Keypair::new();
     let range_proof_context_state_account = Keypair::new();
 
+    let lamport_destination = Pubkey::new_unique();
+
+    let close_split_context_state_accounts = CloseSplitContextStateAccounts {
+        lamport_destination: &lamport_destination,
+        zk_token_proof_program: &zk_token_proof_program::id(),
+    };
+
     let transfer_context_state_accounts = TransferWithFeeSplitContextStateAccounts {
         equality_proof: &equality_proof_context_state_account.pubkey(),
         transfer_amount_ciphertext_validity_proof:
@@ -3053,20 +3094,26 @@ async fn confidential_transfer_transfer_with_fee_and_split_proof_context_in_para
         range_proof: &range_proof_context_state_account.pubkey(),
         authority: &context_state_authority.pubkey(),
         no_op_on_uninitialized_split_context_state: true,
-        close_split_context_state_accounts: None,
+        close_split_context_state_accounts: Some(close_split_context_state_accounts),
     };
 
     let equality_and_ciphertext_proof_signers = vec![
         &alice,
         &equality_proof_context_state_account,
         &transfer_amount_ciphertext_validity_proof_context_state_account,
+        &context_state_authority,
     ];
     let fee_sigma_proof_signers = vec![
         &alice,
         &fee_sigma_proof_context_state_account,
         &fee_ciphertext_validity_proof_context_state_account,
+        &context_state_authority,
+    ];
+    let range_proof_signers = vec![
+        &alice,
+        &range_proof_context_state_account,
+        &context_state_authority,
     ];
-    let range_proof_signers = vec![&alice, &range_proof_context_state_account];
     token
         .confidential_transfer_transfer_with_fee_and_split_proofs_in_parallel(
             &alice_meta.token_account,
@@ -3113,4 +3160,37 @@ async fn confidential_transfer_transfer_with_fee_and_split_proof_context_in_para
             },
         )
         .await;
+
+    let error = token
+        .get_account(equality_proof_context_state_account.pubkey())
+        .await
+        .unwrap_err();
+    assert_eq!(error, TokenClientError::AccountNotFound);
+
+    let error = token
+        .get_account(transfer_amount_ciphertext_validity_proof_context_state_account.pubkey())
+        .await
+        .unwrap_err();
+    assert_eq!(error, TokenClientError::AccountNotFound);
+
+    let error = token
+        .get_account(fee_sigma_proof_context_state_account.pubkey())
+        .await
+        .unwrap_err();
+    assert_eq!(error, TokenClientError::AccountNotFound);
+
+    let error = token
+        .get_account(fee_ciphertext_validity_proof_context_state_account.pubkey())
+        .await
+        .unwrap_err();
+    assert_eq!(error, TokenClientError::AccountNotFound);
+
+    let error = token
+        .get_account(range_proof_context_state_account.pubkey())
+        .await
+        .unwrap_err();
+    assert_eq!(error, TokenClientError::AccountNotFound);
+
+    let lamport_destination = token.get_account(lamport_destination).await.unwrap();
+    assert!(lamport_destination.lamports > 0);
 }

token/program-2022/src/extension/confidential_transfer/instruction.rs

@@ -399,14 +399,14 @@ pub enum ConfidentialTransferInstruction {
     ///   5. `[]` Context state account for
     ///      `VerifyBatchedGroupedCiphertext2HandlesValidityProof`.
     ///   6. `[]` Context state account for `VerifyBatchedRangeProofU128`.
-    ///   7. `[signer]` The source account owner.
     ///   If `close_split_context_state_on_execution` is set, all context state
-    ///     accounts must be `writable` and the following additional sequence
-    ///     of accounts are needed:
-    ///   8. `[]` The destination account for lamports from the context state
+    ///     accounts must be `writable` and the following sequence
+    ///     of accounts that are marked with asterisk are needed:
+    ///   7*. `[]` The destination account for lamports from the context state
     ///      accounts.
-    ///   9. `[signer]` The context state account owner.
-    ///   10. `[]` The zk token proof program.
+    ///   8*. `[signer]` The context state account owner.
+    ///   9*. `[]` The zk token proof program.
+    ///   10. `[signer]` The source account owner.
     ///
     ///   * Transfer with fee
     ///   1. `[writable]` The source SPL Token account.
@@ -420,14 +420,14 @@ pub enum ConfidentialTransferInstruction {
     ///   7. `[]` Context state account for
     ///      `VerifyBatchedGroupedCiphertext2HandlesValidityProof`.
     ///   8. `[]` Context state account for `VerifyBatchedRangeProofU256`.
-    ///   9. `[signer]` The source account owner.
     ///   If `close_split_context_state_on_execution` is set, all context state
-    ///     accounts must be   `writable` and the following additional sequence
-    ///     of accounts are needed:
-    ///   10. `[]` The destination account for lamports from the context state
+    ///     accounts must be  `writable` and the following sequence
+    ///     of accounts that are marked with asterisk are needed:
+    ///   9*. `[]` The destination account for lamports from the context state
     ///       accounts.
-    ///   11. `[signer]` The context state account owner.
-    ///   12. `[]` The zk token proof program.
+    ///   10*. `[signer]` The context state account owner.
+    ///   11*. `[]` The zk token proof program.
+    ///   12. `[signer]` The source account owner.
     ///
     /// Data expected by this instruction:
     ///   `TransferWithSplitProofsInstructionData`
@@ -1389,7 +1389,6 @@ pub fn transfer_with_split_proofs(
                 false,
             ));
             accounts.push(AccountMeta::new(*context_accounts.range_proof, false));
-            accounts.push(AccountMeta::new_readonly(*source_account_authority, true));
             accounts.push(AccountMeta::new(
                 *close_split_context_state_on_execution_accounts.lamport_destination,
                 false,
@@ -1399,6 +1398,7 @@ pub fn transfer_with_split_proofs(
                 *close_split_context_state_on_execution_accounts.zk_token_proof_program,
                 false,
             ));
+            accounts.push(AccountMeta::new_readonly(*source_account_authority, true));
             true
         } else {
             // If `close_split_context_state_accounts` is not set, then context state
@@ -1473,7 +1473,6 @@ pub fn transfer_with_fee_and_split_proofs(
                 false,
             ));
             accounts.push(AccountMeta::new(*context_accounts.range_proof, false));
-            accounts.push(AccountMeta::new_readonly(*source_account_authority, true));
             accounts.push(AccountMeta::new(
                 *close_split_context_state_on_execution_accounts.lamport_destination,
                 false,
@@ -1483,6 +1482,7 @@ pub fn transfer_with_fee_and_split_proofs(
                 *close_split_context_state_on_execution_accounts.zk_token_proof_program,
                 false,
             ));
+            accounts.push(AccountMeta::new_readonly(*source_account_authority, true));
             true
         } else {
             // If `close_split_context_state_accounts` is not set, then context state
@@ -1508,7 +1508,6 @@ pub fn transfer_with_fee_and_split_proofs(
                 false,
             ));
             accounts.push(AccountMeta::new_readonly(*source_account_authority, true));
-
             false
         };
 

token/program-2022/src/extension/confidential_transfer/processor.rs

@@ -1164,13 +1164,6 @@ pub(crate) fn process_instruction(
             {
                 let data =
                     decode_instruction_data::<TransferWithSplitProofsInstructionData>(input)?;
-
-                // Remove this error on the next Solana version upgrade.
-                if data.close_split_context_state_on_execution.into() {
-                    msg!("Auto-close of context state account is not yet supported");
-                    return Err(ProgramError::InvalidInstructionData);
-                }
-
                 process_transfer(
                     program_id,
                     accounts,

token/program-2022/src/extension/confidential_transfer/verify_proof.rs

@@ -188,6 +188,7 @@ pub fn verify_transfer_proof(
         if close_split_context_state_on_execution {
             let lamport_destination_account_info = next_account_info(account_info_iter)?;
             let context_state_account_authority_info = next_account_info(account_info_iter)?;
+            let _zk_token_proof_program = next_account_info(account_info_iter)?;
 
             msg!("Closing equality proof context state account");
             invoke(
@@ -363,6 +364,7 @@ pub fn verify_transfer_with_fee_proof(
         if close_split_context_state_on_execution {
             let lamport_destination_account_info = next_account_info(account_info_iter)?;
             let context_state_account_authority_info = next_account_info(account_info_iter)?;
+            let _zk_token_proof_program = next_account_info(account_info_iter)?;
 
             msg!("Closing equality proof context state account");
             invoke(