update token-cli

samkim-crypto · samkim-crypto · commit a50287177b89 · 2024-08-15T13:36:40.000+09:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/token/cli/Cargo.toml b/token/cli/Cargo.toml
@@ -35,6 +35,7 @@ spl-token-2022 = { version = "4.0.0", path = "../program-2022", features = [
   "no-entrypoint",
 ] }
 spl-token-client = { version = "0.11.0", path = "../client" }
+spl-token-confidential-transfer-proof-generation = { version = "0.1.0", path = "../confidential-transfer/proof-generation" }
 spl-token-metadata-interface = { version = "0.4.0", path = "../../token-metadata/interface" }
 spl-token-group-interface = { version = "0.3.0", path = "../../token-group/interface" }
 spl-associated-token-account = { version = "4.0.0", path = "../../associated-token-account/program", features = [
diff --git a/token/cli/src/command.rs b/token/cli/src/command.rs
@@ -57,19 +57,19 @@ use {
             transfer_hook::TransferHook,
             BaseStateWithExtensions, ExtensionType, StateWithExtensionsOwned,
         },
-        solana_zk_token_sdk::{
-            encryption::{
-                auth_encryption::AeKey,
-                elgamal::{self, ElGamalKeypair},
-            },
-            zk_token_elgamal::pod::ElGamalPubkey,
+        solana_zk_sdk::encryption::{
+            auth_encryption::AeKey,
+            elgamal::{self, ElGamalKeypair},
+            pod::elgamal::PodElGamalPubkey,
         },
         state::{Account, AccountState, Mint},
     },
     spl_token_client::{
         client::{ProgramRpcClientSendTransaction, RpcClientResponse},
-        proof_generation::ProofAccount,
-        token::{ComputeUnitLimit, ExtensionInitializationParams, Token},
+        token::{ComputeUnitLimit, ExtensionInitializationParams, ProofAccount, Token},
+    },
+    spl_token_confidential_transfer_proof_generation::{
+        transfer::TransferProofData, withdraw::WithdrawProofData,
     },
     spl_token_group_interface::state::TokenGroup,
     spl_token_metadata_interface::state::{Field, TokenMetadata},
@@ -1479,7 +1479,7 @@ async fn command_transfer(
             let auditor_elgamal_pubkey = if let Ok(confidential_transfer_mint) =
                 mint_state.get_extension::<ConfidentialTransferMint>()
             {
-                let expected_auditor_elgamal_pubkey = Option::<ElGamalPubkey>::from(
+                let expected_auditor_elgamal_pubkey = Option::<PodElGamalPubkey>::from(
                     confidential_transfer_mint.auditor_elgamal_pubkey,
                 );
 
@@ -1589,35 +1589,41 @@ async fn command_transfer(
                 .unwrap();
             let transfer_account_info = TransferAccountInfo::new(extension);
 
-            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();
+            let TransferProofData {
+                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(
+                token.confidential_transfer_create_context_state_account(
                     &range_proof_pubkey,
                     &context_state_authority_pubkey,
                     &range_proof_data,
+                    true,
                     &range_proof_context_state_account,
                 ),
-                token.create_equality_proof_context_state_for_transfer(
+                token.confidential_transfer_create_context_state_account(
                     &equality_proof_pubkey,
                     &context_state_authority_pubkey,
                     &equality_proof_data,
+                    false,
                     &equality_proof_context_state_account,
                 ),
-                token.create_ciphertext_validity_proof_context_state_for_transfer(
+                token.confidential_transfer_create_context_state_account(
                     &ciphertext_validity_proof_pubkey,
                     &context_state_authority_pubkey,
                     &ciphertext_validity_proof_data,
+                    false,
                     &ciphertext_validity_proof_context_state_account,
                 )
             )?;
@@ -2944,7 +2950,7 @@ async fn command_update_confidential_transfer_settings(
             let new_auditor_pubkey = if let Some(auditor_pubkey) = auditor_pubkey {
                 auditor_pubkey.into()
             } else {
-                Option::<ElGamalPubkey>::from(confidential_transfer_mint.auditor_elgamal_pubkey)
+                Option::<PodElGamalPubkey>::from(confidential_transfer_mint.auditor_elgamal_pubkey)
             };
 
             (new_auto_approve, new_auditor_pubkey)
@@ -3335,28 +3341,47 @@ async fn command_deposit_withdraw_confidential_tokens(
             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 equality_proof_context_state_keypair = Keypair::new();
+            let equality_proof_context_state_pubkey = equality_proof_context_state_keypair.pubkey();
+            let range_proof_context_state_keypair = Keypair::new();
+            let range_proof_context_state_pubkey = range_proof_context_state_keypair.pubkey();
 
-            let withdraw_proof_data =
-                withdraw_account_info.generate_proof_data(amount, elgamal_keypair, aes_key)?;
+            let WithdrawProofData {
+                equality_proof_data,
+                range_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,
+            // set up context state accounts
+            let context_state_authority_pubkey = context_state_authority.pubkey();
+
+            let _ = try_join!(
+                token.confidential_transfer_create_context_state_account(
+                    &equality_proof_context_state_pubkey,
+                    &context_state_authority_pubkey,
+                    &equality_proof_data,
+                    false,
+                    &equality_proof_context_state_keypair,
+                ),
+                token.confidential_transfer_create_context_state_account(
+                    &range_proof_context_state_pubkey,
+                    &context_state_authority_pubkey,
+                    &range_proof_data,
+                    true,
+                    &range_proof_context_state_keypair,
                 )
-                .await?;
+            )?;
 
             // do the withdrawal
-            token
+            let withdraw_result = token
                 .confidential_transfer_withdraw(
                     &token_account_address,
                     &owner,
-                    Some(&ProofAccount::ContextAccount(context_state_pubkey)),
+                    Some(&ProofAccount::ContextAccount(
+                        equality_proof_context_state_pubkey,
+                    )),
+                    Some(&ProofAccount::ContextAccount(
+                        range_proof_context_state_pubkey,
+                    )),
                     amount,
                     decimals,
                     Some(withdraw_account_info),
@@ -3367,15 +3392,22 @@ async fn command_deposit_withdraw_confidential_tokens(
                 .await?;
 
             // close context state account
-            let context_state_authority_pubkey = context_state_authority.pubkey();
-            token
-                .confidential_transfer_close_context_state(
-                    &context_state_pubkey,
+            let _ = try_join!(
+                token.confidential_transfer_close_context_state(
+                    &equality_proof_context_state_pubkey,
+                    &token_account_address,
+                    &context_state_authority_pubkey,
+                    &context_state_authority,
+                ),
+                token.confidential_transfer_close_context_state(
+                    &range_proof_context_state_pubkey,
                     &token_account_address,
                     &context_state_authority_pubkey,
                     &context_state_authority,
                 )
-                .await?
+            )?;
+
+            withdraw_result
         }
     };
 
@@ -3447,8 +3479,8 @@ async fn command_apply_pending_balance(
 struct ConfidentialTransferArgs {
     sender_elgamal_keypair: ElGamalKeypair,
     sender_aes_key: AeKey,
-    recipient_elgamal_pubkey: Option<ElGamalPubkey>,
-    auditor_elgamal_pubkey: Option<ElGamalPubkey>,
+    recipient_elgamal_pubkey: Option<PodElGamalPubkey>,
+    auditor_elgamal_pubkey: Option<PodElGamalPubkey>,
 }
 
 pub async fn process_command<'a>(
diff --git a/token/cli/src/encryption_keypair.rs b/token/cli/src/encryption_keypair.rs
@@ -5,9 +5,9 @@
 use {
     base64::{prelude::BASE64_STANDARD, Engine},
     clap::ArgMatches,
-    spl_token_2022::solana_zk_token_sdk::{
-        encryption::elgamal::{ElGamalKeypair, ElGamalPubkey},
-        zk_token_elgamal::pod::ElGamalPubkey as PodElGamalPubkey,
+    spl_token_2022::solana_zk_sdk::encryption::{
+        elgamal::{ElGamalKeypair, ElGamalPubkey},
+        pod::elgamal::PodElGamalPubkey,
     },
 };
 
diff --git a/token/cli/tests/command.rs b/token/cli/tests/command.rs
@@ -32,7 +32,7 @@ use {
             BaseStateWithExtensions, StateWithExtensionsOwned,
         },
         instruction::create_native_mint,
-        solana_zk_token_sdk::zk_token_elgamal::pod::ElGamalPubkey,
+        solana_zk_sdk::encryption::pod::elgamal::PodElGamalPubkey,
         state::{Account, AccountState, Mint, Multisig},
     },
     spl_token_cli::{
@@ -2609,7 +2609,7 @@ async fn transfer_fee_basis_point(test_validator: &TestValidator, payer: &Keypai
 }
 
 async fn confidential_transfer(test_validator: &TestValidator, payer: &Keypair) {
-    use spl_token_2022::solana_zk_token_sdk::encryption::elgamal::ElGamalKeypair;
+    use spl_token_2022::solana_zk_sdk::encryption::elgamal::ElGamalKeypair;
 
     let config = test_config_with_default_signer(test_validator, payer, &spl_token_2022::id());
 
@@ -2650,13 +2650,13 @@ async fn confidential_transfer(test_validator: &TestValidator, payer: &Keypair)
         auto_approve,
     );
     assert_eq!(
-        Option::<ElGamalPubkey>::from(extension.auditor_elgamal_pubkey),
+        Option::<PodElGamalPubkey>::from(extension.auditor_elgamal_pubkey),
         None,
     );
 
     // update confidential transfer mint settings
     let auditor_keypair = ElGamalKeypair::new_rand();
-    let auditor_pubkey: ElGamalPubkey = (*auditor_keypair.pubkey()).into();
+    let auditor_pubkey: PodElGamalPubkey = (*auditor_keypair.pubkey()).into();
     let new_auto_approve = true;
 
     process_test_command(
@@ -2686,7 +2686,7 @@ async fn confidential_transfer(test_validator: &TestValidator, payer: &Keypair)
         new_auto_approve,
     );
     assert_eq!(
-        Option::<ElGamalPubkey>::from(extension.auditor_elgamal_pubkey),
+        Option::<PodElGamalPubkey>::from(extension.auditor_elgamal_pubkey),
         Some(auditor_pubkey),
     );
 
@@ -2946,7 +2946,7 @@ async fn confidential_transfer_with_fee(test_validator: &TestValidator, payer: &
         auto_approve,
     );
     assert_eq!(
-        Option::<ElGamalPubkey>::from(extension.auditor_elgamal_pubkey),
+        Option::<PodElGamalPubkey>::from(extension.auditor_elgamal_pubkey),
         None,
     );
 
diff --git a/token/client/src/token.rs b/token/client/src/token.rs
@@ -2358,6 +2358,7 @@ where
         context_state_account: &Pubkey,
         context_state_authority: &Pubkey,
         proof_data: &ZK,
+        split_instruction: bool,
         signer: &S,
     ) -> TokenResult<T::Output> {
         let instruction_type = zk_proof_type_to_instruction(ZK::PROOF_TYPE)?;
@@ -2373,20 +2374,54 @@ where
             context_state_authority,
         };
 
-        self.process_ixs(
-            &[
-                system_instruction::create_account(
+        // Some proof instructions are right at the transaction size limit, but in the
+        // future it might be able to support the transfer too
+        if split_instruction {
+            self.process_ixs(
+                &[system_instruction::create_account(
                     &self.payer.pubkey(),
                     context_state_account,
                     rent,
                     space as u64,
                     &zk_elgamal_proof_program::id(),
-                ),
-                instruction_type.encode_verify_proof(Some(context_state_info), proof_data),
-            ],
-            &[signer],
-        )
-        .await
+                )],
+                &[signer],
+            )
+            .await?;
+
+            let blockhash = self
+                .client
+                .get_latest_blockhash()
+                .await
+                .map_err(TokenError::Client)?;
+
+            let transaction = Transaction::new_signed_with_payer(
+                &[instruction_type.encode_verify_proof(Some(context_state_info), proof_data)],
+                Some(&self.payer.pubkey()),
+                &[self.payer.as_ref()],
+                blockhash,
+            );
+
+            self.client
+                .send_transaction(&transaction)
+                .await
+                .map_err(TokenError::Client)
+        } else {
+            self.process_ixs(
+                &[
+                    system_instruction::create_account(
+                        &self.payer.pubkey(),
+                        context_state_account,
+                        rent,
+                        space as u64,
+                        &zk_elgamal_proof_program::id(),
+                    ),
+                    instruction_type.encode_verify_proof(Some(context_state_info), proof_data),
+                ],
+                &[signer],
+            )
+            .await
+        }
     }
 
     /// Close a ZK Token proof program context state
diff --git a/token/program-2022-test/tests/confidential_transfer.rs b/token/program-2022-test/tests/confidential_transfer.rs
