Require pubkey validity proof when initializing confidential account (#3784)

Author: samkim-crypto

token/client/src/token.rs

@@ -1501,9 +1501,8 @@ where
         authority: &S,
         maximum_pending_balance_credit_counter: u64,
     ) -> TokenResult<T::Output> {
-        let elgamal_pubkey = ElGamalKeypair::new(authority, token_account)
-            .map_err(TokenError::Key)?
-            .public;
+        let elgamal_keypair =
+            ElGamalKeypair::new(authority, token_account).map_err(TokenError::Key)?;
         let decryptable_zero_balance = AeKey::new(authority, token_account)
             .map_err(TokenError::Key)?
             .encrypt(0);
@@ -1512,7 +1511,7 @@ where
             token_account,
             authority,
             maximum_pending_balance_credit_counter,
-            elgamal_pubkey,
+            &elgamal_keypair,
             decryptable_zero_balance,
         )
         .await
@@ -1525,20 +1524,25 @@ where
         token_account: &Pubkey,
         authority: &S,
         maximum_pending_balance_credit_counter: u64,
-        elgamal_pubkey: ElGamalPubkey,
+        elgamal_keypair: &ElGamalKeypair,
         decryptable_zero_balance: AeCiphertext,
     ) -> TokenResult<T::Output> {
+        let proof_data =
+            confidential_transfer::instruction::PubkeyValidityData::new(elgamal_keypair)
+                .map_err(TokenError::Proof)?;
+
         self.process_ixs(
-            &[confidential_transfer::instruction::configure_account(
+            &confidential_transfer::instruction::configure_account(
                 &self.program_id,
                 token_account,
                 &self.pubkey,
-                elgamal_pubkey.into(),
+                elgamal_keypair.public.into(),
                 decryptable_zero_balance,
                 maximum_pending_balance_credit_counter,
                 &authority.pubkey(),
                 &[],
-            )?],
+                &proof_data,
+            )?,
             &[authority],
         )
         .await

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

@@ -65,18 +65,23 @@ pub enum ConfidentialTransferInstruction {
     /// Upon success confidential deposits and transfers are enabled, use the
     /// `DisableBalanceCredits` instruction to disable.
     ///
+    /// In order for this instruction to be successfully processed, it must be accompanied by the
+    /// `VerifyPubkey` instruction of the `zk_token_proof` program in the same transaction.
+    ///
     /// Accounts expected by this instruction:
     ///
     ///   * Single owner/delegate
     ///   0. `[writeable]` The SPL Token account.
     ///   1. `[]` The corresponding SPL Token mint.
-    ///   2. `[signer]` The single source account owner.
+    ///   2. `[]` Instructions sysvar.
+    ///   3. `[signer]` The single source account owner.
     ///
     ///   * Multisignature owner/delegate
     ///   0. `[writeable]` The SPL Token account.
     ///   1. `[]` The corresponding SPL Token mint.
     ///   2. `[]` The multisig source account owner.
-    ///   3.. `[signer]` Required M signer accounts for the SPL Token Multisig account.
+    ///   3. `[]` Instructions sysvar.
+    ///   4.. `[signer]` Required M signer accounts for the SPL Token Multisig account.
     ///
     /// Data expected by this instruction:
     ///   `ConfigureAccountInstructionData`
@@ -378,6 +383,9 @@ pub struct ConfigureAccountInstructionData {
     /// The maximum number of despots and transfers that an account can receiver before the
     /// `ApplyPendingBalance` is executed
     pub maximum_pending_balance_credit_counter: PodU64,
+    /// Relative location of the `ProofInstruction::VerifyPubkey` instruction to the
+    /// `ConfigureAccount` instruction in the transaction
+    pub proof_instruction_offset: i8,
 }
 
 /// Data expected by `ConfidentialTransferInstruction::EmptyAccount`
@@ -499,9 +507,11 @@ pub fn update_mint(
 }
 
 /// Create a `ConfigureAccount` instruction
+///
+/// This instruction is suitable for use with a cross-program `invoke`
 #[allow(clippy::too_many_arguments)]
 #[cfg(not(target_os = "solana"))]
-pub fn configure_account(
+pub fn inner_configure_account(
     token_program_id: &Pubkey,
     token_account: &Pubkey,
     mint: &Pubkey,
@@ -510,11 +520,13 @@ pub fn configure_account(
     maximum_pending_balance_credit_counter: u64,
     authority: &Pubkey,
     multisig_signers: &[&Pubkey],
+    proof_instruction_offset: i8,
 ) -> Result<Instruction, ProgramError> {
     check_program_account(token_program_id)?;
     let mut accounts = vec![
         AccountMeta::new(*token_account, false),
         AccountMeta::new_readonly(*mint, false),
+        AccountMeta::new_readonly(sysvar::instructions::id(), false),
         AccountMeta::new_readonly(*authority, multisig_signers.is_empty()),
     ];
 
@@ -531,10 +543,41 @@ pub fn configure_account(
             encryption_pubkey,
             decryptable_zero_balance: decryptable_zero_balance.into(),
             maximum_pending_balance_credit_counter: maximum_pending_balance_credit_counter.into(),
+            proof_instruction_offset,
         },
     ))
 }
 
+/// Create a `ConfigureAccount` instruction
+#[allow(clippy::too_many_arguments)]
+#[cfg(not(target_os = "solana"))]
+pub fn configure_account(
+    token_program_id: &Pubkey,
+    token_account: &Pubkey,
+    mint: &Pubkey,
+    encryption_pubkey: EncryptionPubkey,
+    decryptable_zero_balance: AeCiphertext,
+    maximum_pending_balance_credit_counter: u64,
+    authority: &Pubkey,
+    multisig_signers: &[&Pubkey],
+    proof_data: &PubkeyValidityData,
+) -> Result<Vec<Instruction>, ProgramError> {
+    Ok(vec![
+        inner_configure_account(
+            token_program_id,
+            token_account,
+            mint,
+            encryption_pubkey,
+            decryptable_zero_balance,
+            maximum_pending_balance_credit_counter,
+            authority,
+            multisig_signers,
+            1,
+        )?,
+        verify_pubkey_validity(proof_data),
+    ])
+}
+
 /// Create an `ApproveAccount` instruction
 pub fn approve_account(
     token_program_id: &Pubkey,

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

@@ -124,15 +124,15 @@ fn process_update_mint(
 fn process_configure_account(
     program_id: &Pubkey,
     accounts: &[AccountInfo],
-    ConfigureAccountInstructionData {
-        encryption_pubkey,
-        decryptable_zero_balance,
-        maximum_pending_balance_credit_counter,
-    }: &ConfigureAccountInstructionData,
+    encryption_pubkey: &EncryptionPubkey,
+    decryptable_zero_balance: &DecryptableBalance,
+    maximum_pending_balance_credit_counter: &PodU64,
+    proof_instruction_offset: i64,
 ) -> ProgramResult {
     let account_info_iter = &mut accounts.iter();
     let token_account_info = next_account_info(account_info_iter)?;
     let mint_info = next_account_info(account_info_iter)?;
+    let instructions_sysvar_info = next_account_info(account_info_iter)?;
     let authority_info = next_account_info(account_info_iter)?;
     let authority_info_data_len = authority_info.data_len();
 
@@ -157,6 +157,17 @@ fn process_configure_account(
     let mint = StateWithExtensions::<Mint>::unpack(mint_data)?;
     let confidential_transfer_mint = mint.get_extension::<ConfidentialTransferMint>()?;
 
+    let previous_instruction =
+        get_instruction_relative(proof_instruction_offset, instructions_sysvar_info)?;
+    let proof_data = decode_proof_instruction::<PubkeyValidityData>(
+        ProofInstruction::VerifyPubkeyValidity,
+        &previous_instruction,
+    )?;
+
+    if proof_data.pubkey != *encryption_pubkey {
+        return Err(TokenError::ConfidentialTransferElGamalPubkeyMismatch.into());
+    }
+
     // Note: The caller is expected to use the `Reallocate` instruction to ensure there is
     // sufficient room in their token account for the new `ConfidentialTransferAccount` extension
     let mut confidential_transfer_account =
@@ -1186,10 +1197,14 @@ pub(crate) fn process_instruction(
         }
         ConfidentialTransferInstruction::ConfigureAccount => {
             msg!("ConfidentialTransferInstruction::ConfigureAccount");
+            let data = decode_instruction_data::<ConfigureAccountInstructionData>(input)?;
             process_configure_account(
                 program_id,
                 accounts,
-                decode_instruction_data::<ConfigureAccountInstructionData>(input)?,
+                &data.encryption_pubkey,
+                &data.decryptable_zero_balance,
+                &data.maximum_pending_balance_credit_counter,
+                data.proof_instruction_offset as i64,
             )
         }
         ConfidentialTransferInstruction::ApproveAccount => {