token-2022: add fee withdraw and harvest instructions (#2907)

* token-2022: add fee harvest instruction for confidential extension

* token-2022: cargo fmt

* token-2022: add `WithdrawWithheldTokensMint` processor

* token-2022: add `WithdrawWithheldTokensAccounts` processor

* token-2022: incorporate VerifyWithdrawWithheldTokens instruction

* token-2022: minor variable name changes

* token-2022: additional comments + renaming variables

* token-2022: check if confidential extension was approved for fee harvest

Co-authored-by: Michael Vines <[email protected]>

* token-2022: fix error type for overflow

* token-2022: remove unnecessary fee extension check for confidential token harvest

* token-2022: add extra detail on front-running for harvesting fees

* token-2022: fix pending balance

* token-2022: bump zk-token-sdk to 0.8.0

* token-2022: fix minor variable name issue

* token-2022: minor spelling

Co-authored-by: Tyera Eulberg <[email protected]>

* token-2022: addressing smaller review feedback

* token-2022: change fee_mint to transfer_fee_config in confidential transfer

* Update token/program-2022/Cargo.toml

* token-2022: fix cargo bpf-test fail

* token-2022: fixing a previous comment mistake

Co-authored-by: Michael Vines <[email protected]>
Co-authored-by: Tyera Eulberg <[email protected]>

Author: 3 people

Cargo.lock

@@ -10,10 +10,12 @@ use {
     },
     spl_token_2022::{
         extension::{
-            confidential_transfer::{ConfidentialTransferAccount, ConfidentialTransferMint},
+            confidential_transfer::{
+                ConfidentialTransferAccount, ConfidentialTransferMint, EncryptedWithheldAmount,
+            },
             ExtensionType,
         },
-        solana_zk_token_sdk::encryption::elgamal::*,
+        solana_zk_token_sdk::{encryption::elgamal::*, zk_token_elgamal::pod::Zeroable},
     },
     spl_token_client::token::{ExtensionInitializationParams, TokenError as TokenClientError},
     std::convert::TryInto,
@@ -36,8 +38,9 @@ impl ConfidentialTransferMintWithKeypairs {
         let ct_mint = ConfidentialTransferMint {
             authority: ct_mint_authority.pubkey().into(),
             auto_approve_new_accounts: true.into(),
-            transfer_auditor: ct_mint_transfer_auditor.public.into(),
-            withdraw_withheld_authority: ct_mint_withdraw_withheld_authority.public.into(),
+            pubkey_auditor: ct_mint_transfer_auditor.public.into(),
+            pubkey_withdraw_withheld_authority: ct_mint_withdraw_withheld_authority.public.into(),
+            withheld_amount: EncryptedWithheldAmount::zeroed(),
         };
         Self {
             ct_mint,
@@ -162,7 +165,7 @@ async fn ct_configure_token_account() {
     assert!(!bool::from(&extension.approved));
     assert!(bool::from(&extension.allow_balance_credits));
     assert_eq!(
-        extension.elgamal_pubkey,
+        extension.pubkey_elgamal,
         alice_elgamal_keypair.public.into()
     );
     assert_eq!(

token/program-2022-test/tests/confidential_transfer.rs

@@ -19,7 +19,7 @@ num-derive = "0.3"
 num-traits = "0.2"
 num_enum = "0.5.4"
 solana-program = "1.9.9"
-solana-zk-token-sdk = "0.6.0"
+solana-zk-token-sdk = "0.8.1"
 spl-memo = { version = "3.0.1", path = "../../memo/program", features = [ "no-entrypoint" ] }
 spl-token = { version = "3.3",  path = "../program", features = ["no-entrypoint"] }
 thiserror = "1.0"

token/program-2022/Cargo.toml

@@ -101,9 +101,9 @@ pub enum TokenError {
     /// ElGamal public key mismatch
     #[error("ElGamal public key mismatch")]
     ConfidentialTransferElGamalPubkeyMismatch,
-    /// Available balance mismatch
-    #[error("Available balance mismatch")]
-    ConfidentialTransferAvailableBalanceMismatch,
+    /// Balance mismatch
+    #[error("Balance mismatch")]
+    ConfidentialTransferBalanceMismatch,
     /// Mint has non-zero supply. Burn all tokens before closing the mint.
     #[error("Mint has non-zero supply. Burn all tokens before closing the mint")]
     MintHasSupply,

token/program-2022/src/error.rs

@@ -3,8 +3,7 @@ use solana_zk_token_sdk::encryption::{auth_encryption::AeCiphertext, elgamal::El
 pub use solana_zk_token_sdk::zk_token_proof_instruction::*;
 use {
     crate::{
-        check_program_account, extension::confidential_transfer::ConfidentialTransferMint,
-        instruction::TokenInstruction, pod::*,
+        check_program_account, extension::confidential_transfer::*, instruction::TokenInstruction,
     },
     bytemuck::{Pod, Zeroable},
     num_derive::{FromPrimitive, ToPrimitive},
@@ -16,6 +15,7 @@ use {
         sysvar,
     },
     solana_zk_token_sdk::zk_token_elgamal::pod,
+    std::convert::TryFrom,
 };
 
 /// Confidential Transfer extension instructions
@@ -228,16 +228,89 @@ pub enum ConfidentialTransferInstruction {
     ///   None
     ///
     DisableBalanceCredits,
+
+    /// Transfer all withheld confidential tokens in the mint to an account. Signed by the mint's
+    /// withdraw withheld tokens authority.
+    ///
+    /// Accounts expected by this instruction:
+    ///
+    ///   0. `[writable]` The token mint. Must include the `TransferFeeConfig` extension.
+    ///   1. `[writable]` The fee receiver account. Must include the `TransferFeeAmount` and
+    ///      `ConfidentialTransferAccount` extensions.
+    ///   2. `[]` Instructions sysvar
+    ///   3. `[signer]` The mint's `withdraw_withheld_authority`.
+    /// or:
+    ///   3. `[]` The mint's `withdraw_withheld_authority`'s multisignature owner/delegate.
+    ///   4. ..3+M `[signer]` M signer accounts.
+    ///
+    /// Data expected by this instruction:
+    ///   WithdrawWithheldTokensFromMintData
+    ///
+    WithdrawWithheldTokensFromMint,
+
+    /// Transfer all withheld tokens to an account. Signed by the mint's withdraw withheld tokens
+    /// authority. This instruction is susceptible to front-running. Use
+    /// `HarvestWithheldTokensToMint` and `WithdrawWithheldTokensFromMint` as an alternative.
+    ///
+    /// Note on front-running: This instruction requires a zero-knowledge proof verification
+    /// instruction that is checked with respect to the account state (the currently withheld
+    /// fees). Suppose that a withdraw withheld authority generates the
+    /// `WithdrawWithheldTokensFromAccounts` instruction along with a corresponding zero-knowledge
+    /// proof for a specified set of accounts, and submits it on chain. If the withheld fees at any
+    /// of the specified accounts change before the `WithdrawWithheldTokensFromAccounts` is
+    /// executed on chain, the zero-knowledge proof will not verify with respect to the new state,
+    /// forcing the transaction to fail.
+    ///
+    /// If front-running occurs, then users can look up the updated states of the accounts,
+    /// generate a new zero-knowledge proof and try again. Alternatively, withdraw withheld
+    /// authority can first move the withheld amount to the mint using
+    /// `HarvestWithheldTokensToMint` and then move the withheld fees from mint to a specified
+    /// destination account using `WithdrawWithheldTokensFromMint`.
+    ///
+    /// Accounts expected by this instruction:
+    ///
+    ///   0. `[]` The token mint. Must include the `TransferFeeConfig` extension.
+    ///   1. `[writable]` The fee receiver account. Must include the `TransferFeeAmount` and
+    ///      `ConfidentialTransferAccount` extensions.
+    ///   2. `[]` Instructions sysvar
+    ///   3. `[signer]` The mint's `withdraw_withheld_authority`.
+    ///   4. ..3+N `[writable]` The source accounts to withdraw from.
+    /// or:
+    ///   3. `[]` The mint's `withdraw_withheld_authority`'s multisignature owner/delegate.
+    ///   4. ..4+M `[signer]` M signer accounts.
+    ///   4+M+1. ..3+M+N `[writable]` The source accounts to withdraw from.
+    ///
+    /// Data expected by this instruction:
+    ///   WithdrawWithheldTokensFromAccountsData
+    ///
+    WithdrawWithheldTokensFromAccounts,
+
+    /// Permissionless instruction to transfer all withheld confidential tokens to the mint.
+    ///
+    /// Succeeds for frozen accounts.
+    ///
+    /// Accounts provided should include both the `TransferFeeAmount` and
+    /// `ConfidentialTransferAccount` extension. If not, the account is skipped.
+    ///
+    /// Accounts expected by this instruction:
+    ///
+    ///   0. `[writable]` The mint.
+    ///   1. ..1+N `[writable]` The source accounts to harvest from.
+    ///
+    /// Data expected by this instruction:
+    ///   None
+    ///
+    HarvestWithheldTokensToMint,
 }
 
 /// Data expected by `ConfidentialTransferInstruction::ConfigureAccount`
 #[derive(Clone, Copy, Pod, Zeroable)]
 #[repr(C)]
 pub struct ConfigureAccountInstructionData {
     /// The public key associated with the account
-    pub elgamal_pubkey: pod::ElGamalPubkey,
+    pub elgamal_pubkey: EncryptionPubkey,
     /// The decryptable balance (always 0) once the configure account succeeds
-    pub decryptable_zero_balance: pod::AeCiphertext,
+    pub decryptable_zero_balance: DecryptableBalance,
 }
 
 /// Data expected by `ConfidentialTransferInstruction::EmptyAccount`
@@ -268,7 +341,7 @@ pub struct WithdrawInstructionData {
     /// Expected number of base 10 digits to the right of the decimal place
     pub decimals: u8,
     /// The new decryptable balance if the withrawal succeeds
-    pub new_decryptable_available_balance: pod::AeCiphertext,
+    pub new_decryptable_available_balance: DecryptableBalance,
     /// Relative location of the `ProofInstruction::VerifyWithdraw` instruction to the `Withdraw`
     /// instruction in the transaction
     pub proof_instruction_offset: i8,
@@ -279,7 +352,7 @@ pub struct WithdrawInstructionData {
 #[repr(C)]
 pub struct TransferInstructionData {
     /// The new source decryptable balance if the transfer succeeds
-    pub new_source_decryptable_available_balance: pod::AeCiphertext,
+    pub new_source_decryptable_available_balance: DecryptableBalance,
     /// Relative location of the `ProofInstruction::VerifyTransfer` instruction to the
     /// `Transfer` instruction in the transaction
     pub proof_instruction_offset: i8,
@@ -296,6 +369,26 @@ pub struct ApplyPendingBalanceData {
     pub new_decryptable_available_balance: pod::AeCiphertext,
 }
 
+/// Data expected by `ConfidentialTransferInstruction::WithdrawWithheldTokensFromMint`
+#[derive(Clone, Copy, Pod, Zeroable)]
+#[repr(C)]
+pub struct WithdrawWithheldTokensFromMintData {
+    /// Relative location of the `ProofInstruction::VerifyWithdrawWithheld` instruction to the
+    /// `WithdrawWithheldTokensFromMint` instruction in the transaction
+    pub proof_instruction_offset: i8,
+}
+
+/// Data expected by `ConfidentialTransferInstruction::WithdrawWithheldTokensFromAccounts`
+#[derive(Clone, Copy, Pod, Zeroable)]
+#[repr(C)]
+pub struct WithdrawWithheldTokensFromAccountsData {
+    /// Number of token accounts harvested
+    pub num_token_accounts: u8,
+    /// Relative location of the `ProofInstruction::VerifyWithdrawWithheld` instruction to the
+    /// `VerifyWithdrawWithheldTokensFromAccounts` instruction in the transaction
+    pub proof_instruction_offset: i8,
+}
+
 pub(crate) fn decode_instruction_type(
     input: &[u8],
 ) -> Result<ConfidentialTransferInstruction, ProgramError> {
@@ -521,7 +614,7 @@ pub fn inner_withdraw(
     mint: &Pubkey,
     amount: u64,
     decimals: u8,
-    new_decryptable_available_balance: pod::AeCiphertext,
+    new_decryptable_available_balance: DecryptableBalance,
     authority: &Pubkey,
     multisig_signers: &[&Pubkey],
     proof_instruction_offset: i8,
@@ -593,7 +686,7 @@ pub fn inner_transfer(
     source_token_account: &Pubkey,
     destination_token_account: &Pubkey,
     mint: &Pubkey,
-    new_source_decryptable_available_balance: pod::AeCiphertext,
+    new_source_decryptable_available_balance: DecryptableBalance,
     authority: &Pubkey,
     multisig_signers: &[&Pubkey],
     proof_instruction_offset: i8,
@@ -657,7 +750,7 @@ pub fn inner_apply_pending_balance(
     token_program_id: &Pubkey,
     token_account: &Pubkey,
     expected_pending_balance_credit_counter: u64,
-    new_decryptable_available_balance: pod::AeCiphertext,
+    new_decryptable_available_balance: DecryptableBalance,
     authority: &Pubkey,
     multisig_signers: &[&Pubkey],
 ) -> Result<Instruction, ProgramError> {
@@ -760,3 +853,144 @@ pub fn disable_balance_credits(
         multisig_signers,
     )
 }
+
+/// Create a inner `WithdrawWithheldTokensFromMint` instruction
+///
+/// This instruction is suitable for use with a cross-program `invoke`
+pub fn inner_withdraw_withheld_tokens_from_mint(
+    token_program_id: &Pubkey,
+    mint: &Pubkey,
+    destination: &Pubkey,
+    authority: &Pubkey,
+    multisig_signers: &[&Pubkey],
+    proof_instruction_offset: i8,
+) -> Result<Instruction, ProgramError> {
+    check_program_account(token_program_id)?;
+    let mut accounts = vec![
+        AccountMeta::new(*mint, false),
+        AccountMeta::new(*destination, false),
+        AccountMeta::new_readonly(sysvar::instructions::id(), false),
+        AccountMeta::new_readonly(*authority, multisig_signers.is_empty()),
+    ];
+
+    for multisig_signer in multisig_signers.iter() {
+        accounts.push(AccountMeta::new(**multisig_signer, false));
+    }
+
+    Ok(encode_instruction(
+        token_program_id,
+        accounts,
+        ConfidentialTransferInstruction::WithdrawWithheldTokensFromMint,
+        &WithdrawWithheldTokensFromMintData {
+            proof_instruction_offset,
+        },
+    ))
+}
+
+/// Create a `WithdrawWithheldTokensFromMint` instruction
+pub fn withdraw_withheld_tokens_from_mint(
+    token_program_id: &Pubkey,
+    mint: &Pubkey,
+    destination: &Pubkey,
+    authority: &Pubkey,
+    multisig_signers: &[&Pubkey],
+    proof_data: &WithdrawWithheldTokensData,
+) -> Result<Vec<Instruction>, ProgramError> {
+    Ok(vec![
+        verify_withdraw_withheld_tokens(proof_data),
+        inner_withdraw_withheld_tokens_from_mint(
+            token_program_id,
+            mint,
+            destination,
+            authority,
+            multisig_signers,
+            -1,
+        )?,
+    ])
+}
+
+/// Create a inner `WithdrawWithheldTokensFromMint` instruction
+///
+/// This instruction is suitable for use with a cross-program `invoke`
+pub fn inner_withdraw_withheld_tokens_from_accounts(
+    token_program_id: &Pubkey,
+    mint: &Pubkey,
+    destination: &Pubkey,
+    authority: &Pubkey,
+    multisig_signers: &[&Pubkey],
+    sources: &[&Pubkey],
+    proof_instruction_offset: i8,
+) -> Result<Instruction, ProgramError> {
+    check_program_account(token_program_id)?;
+    let num_token_accounts =
+        u8::try_from(sources.len()).map_err(|_| ProgramError::InvalidInstructionData)?;
+    let mut accounts = vec![
+        AccountMeta::new(*mint, false),
+        AccountMeta::new(*destination, false),
+        AccountMeta::new_readonly(sysvar::instructions::id(), false),
+        AccountMeta::new_readonly(*authority, multisig_signers.is_empty()),
+    ];
+
+    for multisig_signer in multisig_signers.iter() {
+        accounts.push(AccountMeta::new(**multisig_signer, false));
+    }
+
+    for source in sources.iter() {
+        accounts.push(AccountMeta::new(**source, false));
+    }
+
+    Ok(encode_instruction(
+        token_program_id,
+        accounts,
+        ConfidentialTransferInstruction::WithdrawWithheldTokensFromAccounts,
+        &WithdrawWithheldTokensFromAccountsData {
+            proof_instruction_offset,
+            num_token_accounts,
+        },
+    ))
+}
+
+/// Create a `WithdrawWithheldTokensFromAccounts` instruction
+pub fn withdraw_withheld_tokens_from_accounts(
+    token_program_id: &Pubkey,
+    mint: &Pubkey,
+    destination: &Pubkey,
+    authority: &Pubkey,
+    multisig_signers: &[&Pubkey],
+    sources: &[&Pubkey],
+    proof_data: &WithdrawWithheldTokensData,
+) -> Result<Vec<Instruction>, ProgramError> {
+    Ok(vec![
+        verify_withdraw_withheld_tokens(proof_data),
+        inner_withdraw_withheld_tokens_from_accounts(
+            token_program_id,
+            mint,
+            destination,
+            authority,
+            multisig_signers,
+            sources,
+            -1,
+        )?,
+    ])
+}
+
+/// Creates a `HarvestWithheldTokensToMint` instruction
+pub fn harvest_withheld_tokens_to_mint(
+    token_program_id: &Pubkey,
+    mint: &Pubkey,
+    sources: &[&Pubkey],
+) -> Result<Instruction, ProgramError> {
+    check_program_account(token_program_id)?;
+    let mut accounts = vec![AccountMeta::new(*mint, false)];
+
+    for source in sources.iter() {
+        accounts.push(AccountMeta::new(**source, false));
+    }
+
+    Ok(encode_instruction(
+        token_program_id,
+        accounts,
+        ConfidentialTransferInstruction::HarvestWithheldTokensToMint,
+        &(),
+    ))
+}