[token-cli] sequentially process confidential transfer withdraw in to…

…ken-cli (#5674)

* add proof context function for withdraw in token-client

* use proof context for withdraw command in token-cli

token/cli/src/main.rs

@@ -3409,18 +3409,48 @@ async fn command_deposit_withdraw_confidential_tokens(
                 state_with_extension.get_extension::<ConfidentialTransferAccount>()?;
             let withdraw_account_info = WithdrawAccountInfo::new(extension_state);
 
+            let context_state_authority = config.fee_payer()?;
+            let context_state_keypair = Keypair::new();
+            let context_state_pubkey = context_state_keypair.pubkey();
+
+            let withdraw_proof_data =
+                withdraw_account_info.generate_proof_data(amount, elgamal_keypair, aes_key)?;
+
+            // setup proof
+            token
+                .create_withdraw_proof_context_state(
+                    &context_state_pubkey,
+                    &context_state_authority.pubkey(),
+                    &withdraw_proof_data,
+                    &context_state_keypair,
+                )
+                .await?;
+
+            // do the withdrawal
             token
                 .confidential_transfer_withdraw(
                     &token_account_address,
                     &owner,
-                    None,
+                    Some(&context_state_pubkey),
                     amount,
                     decimals,
                     Some(withdraw_account_info),
                     elgamal_keypair,
                     aes_key,
                     &bulk_signers,
                 )
+                .await?;
+
+            // 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,
+                )
                 .await?
         }
     };
@@ -5315,7 +5345,6 @@ fn app<'a, 'b>(
                         .validator(is_valid_pubkey)
                         .value_name("TOKEN_ACCOUNT_ADDRESS")
                         .takes_value(true)
-                        .conflicts_with("token")
                         .help("The address of the token account to configure confidential transfers for \
                             [default: owner's associated token account]")
                 )
@@ -9175,23 +9204,35 @@ mod tests {
         .unwrap();
 
         // withdraw confidential tokens
-        //
-        // NOTE: the test fails due to transaction size limit :(
-
-        // let withdraw_amount = 100.0;
-        //
-        // process_test_command(
-        //     &config,
-        //     &payer,
-        //     &[
-        //         "spl-token",
-        //         CommandName::WithdrawConfidentialTokens.into(),
-        //         &token_pubkey.to_string(),
-        //         &withdraw_amount.to_string(),
-        //     ],
-        // )
-        // .await
-        // .unwrap();
+        process_test_command(
+            &config,
+            &payer,
+            &[
+                "spl-token",
+                CommandName::ApplyPendingBalance.into(),
+                "--address",
+                &destination_account.to_string(),
+            ],
+        )
+        .await
+        .unwrap(); // apply pending balance first
+
+        let withdraw_amount = 100.0;
+
+        process_test_command(
+            &config,
+            &payer,
+            &[
+                "spl-token",
+                CommandName::WithdrawConfidentialTokens.into(),
+                &token_pubkey.to_string(),
+                &withdraw_amount.to_string(),
+                "--address",
+                &destination_account.to_string(),
+            ],
+        )
+        .await
+        .unwrap();
 
         // disable confidential transfers for mint
         process_test_command(

token/client/src/token.rs

@@ -2135,6 +2135,49 @@ where
         .await
     }
 
+    /// Create withdraw proof context state account for a confidential transfer withdraw
+    /// instruction.
+    pub async fn create_withdraw_proof_context_state<S: Signer>(
+        &self,
+        context_state_account: &Pubkey,
+        context_state_authority: &Pubkey,
+        withdraw_proof_data: &WithdrawData,
+        withdraw_proof_signer: &S,
+    ) -> TokenResult<T::Output> {
+        // create withdraw proof context state
+        let instruction_type = ProofInstruction::VerifyWithdraw;
+        let space = size_of::<ProofContextState<WithdrawProofContext>>();
+        let rent = self
+            .client
+            .get_minimum_balance_for_rent_exemption(space)
+            .await
+            .map_err(TokenError::Client)?;
+
+        let withdraw_proof_context_state_info = ContextStateInfo {
+            context_state_account,
+            context_state_authority,
+        };
+
+        self.process_ixs(
+            &[system_instruction::create_account(
+                &self.payer.pubkey(),
+                context_state_account,
+                rent,
+                space as u64,
+                &zk_token_proof_program::id(),
+            )],
+            &[withdraw_proof_signer],
+        )
+        .await?;
+
+        self.process_ixs(
+            &[instruction_type
+                .encode_verify_proof(Some(withdraw_proof_context_state_info), withdraw_proof_data)],
+            &[] as &[&dyn Signer; 0],
+        )
+        .await
+    }
+
     /// Withdraw SPL Tokens from the available balance of a confidential token account using custom
     /// keys
     #[allow(clippy::too_many_arguments)]