[token-2022] Add support for record program for Transfer and TransferWithFee instructions (#7127)

* refactor out verify proof logic into a separate function

* clean up transfer proof verification

* re-organize confidential transfer instruction into `Transfer` and `TransferWithFee`

* use `verify_and_extract_context` function directly

* update token client

* update program-2022 tests

* update token-cli

* cargo fmt

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

Co-authored-by: Jon C <[email protected]>

---------

Co-authored-by: Jon C <[email protected]>

Author: samkim-crypto

token/cli/src/command.rs

@@ -41,7 +41,6 @@ use {
                 account_info::{
                     ApplyPendingBalanceAccountInfo, TransferAccountInfo, WithdrawAccountInfo,
                 },
-                instruction::TransferSplitContextStateAccounts,
                 ConfidentialTransferAccount, ConfidentialTransferMint,
             },
             confidential_transfer_fee::ConfidentialTransferFeeConfig,
@@ -1575,6 +1574,7 @@ async fn command_transfer(
             });
 
             let context_state_authority = config.fee_payer()?;
+            let context_state_authority_pubkey = context_state_authority.pubkey();
             let equality_proof_context_state_account = Keypair::new();
             let equality_proof_pubkey = equality_proof_context_state_account.pubkey();
             let ciphertext_validity_proof_context_state_account = Keypair::new();
@@ -1583,91 +1583,90 @@ async fn command_transfer(
             let range_proof_context_state_account = Keypair::new();
             let range_proof_pubkey = range_proof_context_state_account.pubkey();
 
-            let transfer_context_state_accounts = TransferSplitContextStateAccounts {
-                equality_proof: &equality_proof_pubkey,
-                ciphertext_validity_proof: &ciphertext_validity_proof_pubkey,
-                range_proof: &range_proof_pubkey,
-                authority: &context_state_authority.pubkey(),
-                no_op_on_uninitialized_split_context_state: false,
-                close_split_context_state_accounts: None,
-            };
-
             let state = token.get_account_info(&sender).await.unwrap();
             let extension = state
                 .get_extension::<ConfidentialTransferAccount>()
                 .unwrap();
             let transfer_account_info = TransferAccountInfo::new(extension);
 
-            let (
-                equality_proof_data,
-                ciphertext_validity_proof_data,
-                range_proof_data,
-                source_decrypt_handles,
-            ) = transfer_account_info
-                .generate_split_transfer_proof_data(
-                    transfer_balance,
-                    &args.sender_elgamal_keypair,
-                    &args.sender_aes_key,
-                    &recipient_elgamal_pubkey,
-                    auditor_elgamal_pubkey.as_ref(),
-                )
-                .unwrap();
+            let (equality_proof_data, ciphertext_validity_proof_data, range_proof_data) =
+                transfer_account_info
+                    .generate_split_transfer_proof_data(
+                        transfer_balance,
+                        &args.sender_elgamal_keypair,
+                        &args.sender_aes_key,
+                        &recipient_elgamal_pubkey,
+                        auditor_elgamal_pubkey.as_ref(),
+                    )
+                    .unwrap();
 
             // setup proofs
             let _ = try_join!(
                 token.create_range_proof_context_state_for_transfer(
-                    transfer_context_state_accounts,
+                    &range_proof_pubkey,
+                    &context_state_authority_pubkey,
                     &range_proof_data,
                     &range_proof_context_state_account,
                 ),
                 token.create_equality_proof_context_state_for_transfer(
-                    transfer_context_state_accounts,
+                    &equality_proof_pubkey,
+                    &context_state_authority_pubkey,
                     &equality_proof_data,
                     &equality_proof_context_state_account,
                 ),
                 token.create_ciphertext_validity_proof_context_state_for_transfer(
-                    transfer_context_state_accounts,
+                    &ciphertext_validity_proof_pubkey,
+                    &context_state_authority_pubkey,
                     &ciphertext_validity_proof_data,
                     &ciphertext_validity_proof_context_state_account,
                 )
             )?;
 
             // do the transfer
+            let equality_proof_context_proof_account =
+                ProofAccount::ContextAccount(equality_proof_pubkey);
+            let ciphertext_validity_proof_context_proof_account =
+                ProofAccount::ContextAccount(ciphertext_validity_proof_pubkey);
+            let range_proof_context_proof_account =
+                ProofAccount::ContextAccount(range_proof_pubkey);
+
             let transfer_result = token
-                .confidential_transfer_transfer_with_split_proofs(
+                .confidential_transfer_transfer(
                     &sender,
                     &recipient_token_account,
                     &sender_owner,
-                    transfer_context_state_accounts,
+                    Some(&equality_proof_context_proof_account),
+                    Some(&ciphertext_validity_proof_context_proof_account),
+                    Some(&range_proof_context_proof_account),
                     transfer_balance,
                     Some(transfer_account_info),
+                    &args.sender_elgamal_keypair,
                     &args.sender_aes_key,
-                    &source_decrypt_handles,
+                    &recipient_elgamal_pubkey,
+                    auditor_elgamal_pubkey.as_ref(),
                     &bulk_signers,
                 )
                 .await?;
 
             // close context state accounts
-            let context_state_authority_pubkey = context_state_authority.pubkey();
-            let close_context_state_signers = &[context_state_authority];
             let _ = try_join!(
                 token.confidential_transfer_close_context_state(
                     &equality_proof_pubkey,
                     &sender,
                     &context_state_authority_pubkey,
-                    close_context_state_signers,
+                    &context_state_authority,
                 ),
                 token.confidential_transfer_close_context_state(
                     &ciphertext_validity_proof_pubkey,
                     &sender,
                     &context_state_authority_pubkey,
-                    close_context_state_signers,
+                    &context_state_authority,
                 ),
                 token.confidential_transfer_close_context_state(
                     &range_proof_pubkey,
                     &sender,
                     &context_state_authority_pubkey,
-                    close_context_state_signers,
+                    &context_state_authority,
                 ),
             )?;
 
@@ -3369,13 +3368,12 @@ async fn command_deposit_withdraw_confidential_tokens(
 
             // close context state account
             let context_state_authority_pubkey = context_state_authority.pubkey();
-            let close_context_state_signers = &[context_state_authority];
             token
                 .confidential_transfer_close_context_state(
                     &context_state_pubkey,
                     &token_account_address,
                     &context_state_authority_pubkey,
-                    close_context_state_signers,
+                    &context_state_authority,
                 )
                 .await?
         }

token/client/src/proof_generation.rs

@@ -10,13 +10,13 @@ use {
     spl_token_2022::{
         error::TokenError,
         extension::confidential_transfer::{
-            ciphertext_extraction::{transfer_amount_source_ciphertext, SourceDecryptHandles},
+            ciphertext_extraction::transfer_amount_source_ciphertext,
             processor::verify_and_split_deposit_amount,
         },
         solana_zk_token_sdk::{
             encryption::{
                 auth_encryption::{AeCiphertext, AeKey},
-                elgamal::{DecryptHandle, ElGamalCiphertext, ElGamalKeypair, ElGamalPubkey},
+                elgamal::{ElGamalCiphertext, ElGamalKeypair, ElGamalPubkey},
                 grouped_elgamal::GroupedElGamal,
                 pedersen::{Pedersen, PedersenCommitment, PedersenOpening},
             },
@@ -25,7 +25,8 @@ use {
                     try_combine_lo_hi_commitments, try_combine_lo_hi_openings, FeeEncryption,
                     FeeParameters, TransferAmountCiphertext,
                 },
-                BatchedGroupedCiphertext2HandlesValidityProofData, BatchedRangeProofU256Data,
+                BatchedGroupedCiphertext2HandlesValidityProofData,
+                BatchedGroupedCiphertext3HandlesValidityProofData, BatchedRangeProofU256Data,
                 CiphertextCommitmentEqualityProofData, FeeSigmaProofData,
             },
             zk_token_elgamal::ops::subtract_with_lo_hi,
@@ -53,11 +54,10 @@ pub fn transfer_with_fee_split_proof_data(
 ) -> Result<
     (
         CiphertextCommitmentEqualityProofData,
-        BatchedGroupedCiphertext2HandlesValidityProofData,
+        BatchedGroupedCiphertext3HandlesValidityProofData,
         FeeSigmaProofData,
         BatchedGroupedCiphertext2HandlesValidityProofData,
         BatchedRangeProofU256Data,
-        SourceDecryptHandles,
     ),
     TokenError,
 > {
@@ -126,33 +126,31 @@ pub fn transfer_with_fee_split_proof_data(
     )
     .map_err(|_| TokenError::ProofGeneration)?;
 
-    // create source decrypt handle
-    let source_decrypt_handle_lo =
-        DecryptHandle::new(source_elgamal_keypair.pubkey(), &transfer_amount_opening_lo);
-    let source_decrypt_handle_hi =
-        DecryptHandle::new(source_elgamal_keypair.pubkey(), &transfer_amount_opening_hi);
-
-    let source_decrypt_handles = SourceDecryptHandles {
-        lo: source_decrypt_handle_lo.into(),
-        hi: source_decrypt_handle_hi.into(),
-    };
-
-    // encrypt the transfer amount under the destination and auditor ElGamal public
-    // key
+    // encrypt the transfer amount under the source, destination and auditor ElGamal
+    // public key
     let transfer_amount_destination_auditor_ciphertext_lo = GroupedElGamal::encrypt_with(
-        [destination_elgamal_pubkey, auditor_elgamal_pubkey],
+        [
+            source_elgamal_keypair.pubkey(),
+            destination_elgamal_pubkey,
+            auditor_elgamal_pubkey,
+        ],
         transfer_amount_lo,
         &transfer_amount_opening_lo,
     );
     let transfer_amount_destination_auditor_ciphertext_hi = GroupedElGamal::encrypt_with(
-        [destination_elgamal_pubkey, auditor_elgamal_pubkey],
+        [
+            source_elgamal_keypair.pubkey(),
+            destination_elgamal_pubkey,
+            auditor_elgamal_pubkey,
+        ],
         transfer_amount_hi,
         &transfer_amount_opening_hi,
     );
 
     // generate transfer amount ciphertext validity data
     let transfer_amount_ciphertext_validity_proof_data =
-        BatchedGroupedCiphertext2HandlesValidityProofData::new(
+        BatchedGroupedCiphertext3HandlesValidityProofData::new(
+            source_elgamal_keypair.pubkey(),
             destination_elgamal_pubkey,
             auditor_elgamal_pubkey,
             &transfer_amount_destination_auditor_ciphertext_lo,
@@ -332,7 +330,6 @@ pub fn transfer_with_fee_split_proof_data(
         fee_sigma_proof_data,
         fee_ciphertext_validity_proof_data,
         range_proof_data,
-        source_decrypt_handles,
     ))
 }