[confidential-extension] Use OptionalNonZeroPubkey and OptionalNonZeroEncryptionPubkey for confidential extension (#3943)

* use OptionalNonZeroPubkey for the confidential mint authority pubkey

* add OptionalNonZeroEncryptionPubkey

* update tests

* Update token/program-2022/src/pod.rs

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

* remove unnecessary COption convertion for OptionalNonZeroEncryptionPubkey

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

Author: samkim-crypto

token/client/src/token.rs

@@ -20,11 +20,12 @@ use {
     },
     spl_token_2022::{
         extension::{
-            confidential_transfer, confidential_transfer::EncryptionPubkey, cpi_guard,
-            default_account_state, interest_bearing_mint, memo_transfer, transfer_fee,
-            BaseStateWithExtensions, ExtensionType, StateWithExtensionsOwned,
+            confidential_transfer, cpi_guard, default_account_state, interest_bearing_mint,
+            memo_transfer, transfer_fee, BaseStateWithExtensions, ExtensionType,
+            StateWithExtensionsOwned,
         },
         instruction,
+        pod::EncryptionPubkey,
         solana_zk_token_sdk::{
             encryption::{auth_encryption::*, elgamal::*},
             errors::ProofError,
@@ -106,8 +107,8 @@ pub enum ExtensionInitializationParams {
     ConfidentialTransferMint {
         authority: Option<Pubkey>,
         auto_approve_new_accounts: bool,
-        auditor_encryption_pubkey: EncryptionPubkey,
-        withdraw_withheld_authority_encryption_pubkey: EncryptionPubkey,
+        auditor_encryption_pubkey: Option<EncryptionPubkey>,
+        withdraw_withheld_authority_encryption_pubkey: Option<EncryptionPubkey>,
     },
     DefaultAccountState {
         state: AccountState,
@@ -158,10 +159,10 @@ impl ExtensionInitializationParams {
             } => confidential_transfer::instruction::initialize_mint(
                 token_program_id,
                 mint,
-                authority.as_ref(),
+                authority,
                 auto_approve_new_accounts,
-                &auditor_encryption_pubkey,
-                &withdraw_withheld_authority_encryption_pubkey,
+                auditor_encryption_pubkey,
+                withdraw_withheld_authority_encryption_pubkey,
             ),
             Self::DefaultAccountState { state } => {
                 default_account_state::instruction::initialize_default_account_state(
@@ -1492,7 +1493,7 @@ where
         authority: &S,
         new_authority: Option<&S>,
         auto_approve_new_account: bool,
-        auditor_encryption_pubkey: &EncryptionPubkey,
+        auditor_encryption_pubkey: Option<EncryptionPubkey>,
     ) -> TokenResult<T::Output> {
         let mut signers = vec![authority];
         let new_authority_pubkey = if let Some(new_authority) = new_authority {
@@ -1789,16 +1790,20 @@ where
     /// Fetch the ElGamal pubkey key of the auditor associated with a confidential token mint
     pub async fn confidential_transfer_get_auditor_encryption_pubkey<S: Signer>(
         &self,
-    ) -> TokenResult<ElGamalPubkey> {
+    ) -> TokenResult<Option<ElGamalPubkey>> {
         let mint_state = self.get_mint_info().await.unwrap();
         let ct_mint =
             mint_state.get_extension::<confidential_transfer::ConfidentialTransferMint>()?;
-        let auditor_pubkey = ct_mint
-            .auditor_encryption_pubkey
-            .try_into()
-            .map_err(TokenError::Proof)?;
+        let auditor_encryption_pubkey: Option<EncryptionPubkey> =
+            ct_mint.auditor_encryption_pubkey.into();
 
-        Ok(auditor_pubkey)
+        if let Some(encryption_pubkey) = auditor_encryption_pubkey {
+            let encryption_pubkey: ElGamalPubkey =
+                encryption_pubkey.try_into().map_err(TokenError::Proof)?;
+            Ok(Some(encryption_pubkey))
+        } else {
+            Ok(None)
+        }
     }
 
     /// Fetch the ElGamal pubkey key of the withdraw withheld authority associated with a
@@ -1807,16 +1812,20 @@ where
         S: Signer,
     >(
         &self,
-    ) -> TokenResult<ElGamalPubkey> {
+    ) -> TokenResult<Option<ElGamalPubkey>> {
         let mint_state = self.get_mint_info().await.unwrap();
         let ct_mint =
             mint_state.get_extension::<confidential_transfer::ConfidentialTransferMint>()?;
-        let auditor_pubkey = ct_mint
-            .withdraw_withheld_authority_encryption_pubkey
-            .try_into()
-            .map_err(TokenError::Proof)?;
+        let withdraw_withheld_authority_encryption_pubkey: Option<EncryptionPubkey> =
+            ct_mint.withdraw_withheld_authority_encryption_pubkey.into();
 
-        Ok(auditor_pubkey)
+        if let Some(encryption_pubkey) = withdraw_withheld_authority_encryption_pubkey {
+            let encryption_pubkey: ElGamalPubkey =
+                encryption_pubkey.try_into().map_err(TokenError::Proof)?;
+            Ok(Some(encryption_pubkey))
+        } else {
+            Ok(None)
+        }
     }
 
     /// Deposit SPL Tokens into the pending balance of a confidential token account
@@ -1932,7 +1941,7 @@ where
         source_available_balance: u64,
         source_available_balance_ciphertext: &ElGamalCiphertext,
         destination_elgamal_pubkey: &ElGamalPubkey,
-        auditor_elgamal_pubkey: &ElGamalPubkey,
+        auditor_elgamal_pubkey: Option<ElGamalPubkey>,
     ) -> TokenResult<T::Output> {
         let source_elgamal_keypair =
             ElGamalKeypair::new(source_token_authority, source_token_account)
@@ -1966,22 +1975,24 @@ where
         source_available_balance: u64,
         source_available_balance_ciphertext: &ElGamalCiphertext,
         destination_elgamal_pubkey: &ElGamalPubkey,
-        auditor_elgamal_pubkey: &ElGamalPubkey,
+        auditor_elgamal_pubkey: Option<ElGamalPubkey>,
         source_elgamal_keypair: &ElGamalKeypair,
         source_authenticated_encryption_key: &AeKey,
     ) -> TokenResult<T::Output> {
         if amount >> confidential_transfer::MAXIMUM_DEPOSIT_TRANSFER_AMOUNT_BIT_LENGTH != 0 {
             return Err(TokenError::MaximumDepositTransferAmountExceeded);
         }
 
+        let auditor_elgamal_pubkey = auditor_elgamal_pubkey.unwrap_or_default();
+
         let proof_data = confidential_transfer::instruction::TransferData::new(
             amount,
             (
                 source_available_balance,
                 source_available_balance_ciphertext,
             ),
             source_elgamal_keypair,
-            (destination_elgamal_pubkey, auditor_elgamal_pubkey),
+            (destination_elgamal_pubkey, &auditor_elgamal_pubkey),
         )
         .map_err(TokenError::Proof)?;
 
@@ -2019,7 +2030,7 @@ where
         source_available_balance: u64,
         source_available_balance_ciphertext: &ElGamalCiphertext,
         destination_elgamal_pubkey: &ElGamalPubkey,
-        auditor_elgamal_pubkey: &ElGamalPubkey,
+        auditor_elgamal_pubkey: Option<ElGamalPubkey>,
         withdraw_withheld_authority_elgamal_pubkey: &ElGamalPubkey,
         epoch_info: &EpochInfo,
     ) -> TokenResult<T::Output> {
@@ -2057,7 +2068,7 @@ where
         source_available_balance: u64,
         source_available_balance_ciphertext: &ElGamalCiphertext,
         destination_elgamal_pubkey: &ElGamalPubkey,
-        auditor_elgamal_pubkey: &ElGamalPubkey,
+        auditor_elgamal_pubkey: Option<ElGamalPubkey>,
         withdraw_withheld_authority_elgamal_pubkey: &ElGamalPubkey,
         source_elgamal_keypair: &ElGamalKeypair,
         source_authenticated_encryption_key: &AeKey,
@@ -2067,7 +2078,8 @@ where
             return Err(TokenError::MaximumDepositTransferAmountExceeded);
         }
 
-        // TODO: take transfer fee params as input
+        let auditor_elgamal_pubkey = auditor_elgamal_pubkey.unwrap_or_default();
+
         let mint_state = self.get_mint_info().await.unwrap();
         let transfer_fee_config = mint_state
             .get_extension::<transfer_fee::TransferFeeConfig>()
@@ -2081,7 +2093,7 @@ where
                 source_available_balance_ciphertext,
             ),
             source_elgamal_keypair,
-            (destination_elgamal_pubkey, auditor_elgamal_pubkey),
+            (destination_elgamal_pubkey, &auditor_elgamal_pubkey),
             FeeParameters {
                 fee_rate_basis_points: u16::from(fee_parameters.transfer_fee_basis_points),
                 maximum_fee: u64::from(fee_parameters.maximum_fee),