associated-token-account: Add "CreateIdempotent" instruction (#2883)

* associated-token-account: Add "CreateIfNonExistent" instruction

* Address feedback

* Move ownership check

Author: joncinque

Cargo.lock

@@ -12,9 +12,13 @@ no-entrypoint = []
 test-bpf = []
 
 [dependencies]
+assert_matches = "1.5.0"
 borsh = "0.9.1"
+num-derive = "0.3"
+num-traits = "0.2"
 solana-program = "1.9.5"
 spl-token = { version = "0.1", path = "../../token/program-2022", package = "spl-token-2022", features = ["no-entrypoint"] }
+thiserror = "1.0"
 
 [dev-dependencies]
 solana-program-test = "1.9.5"

associated-token-account/program/Cargo.toml

@@ -0,0 +1,26 @@
+//! Error types
+
+use {
+    num_derive::FromPrimitive,
+    solana_program::{decode_error::DecodeError, program_error::ProgramError},
+    thiserror::Error,
+};
+
+/// Errors that may be returned by the program.
+#[derive(Clone, Debug, Eq, Error, FromPrimitive, PartialEq)]
+pub enum AssociatedTokenAccountError {
+    // 0
+    /// Associated token account owner does not match address derivation
+    #[error("Associated token account owner does not match address derivation")]
+    InvalidOwner,
+}
+impl From<AssociatedTokenAccountError> for ProgramError {
+    fn from(e: AssociatedTokenAccountError) -> Self {
+        ProgramError::Custom(e as u32)
+    }
+}
+impl<T> DecodeError<T> for AssociatedTokenAccountError {
+    fn type_of() -> &'static str {
+        "AssociatedTokenAccountError"
+    }
+}

associated-token-account/program/src/error.rs

@@ -1,6 +1,7 @@
 //! Program instructions
 
 use {
+    assert_matches::assert_matches,
     borsh::{BorshDeserialize, BorshSchema, BorshSerialize},
     solana_program::{
         instruction::{AccountMeta, Instruction},
@@ -14,6 +15,7 @@ use crate::{get_associated_token_address, id};
 #[derive(Clone, Debug, PartialEq, BorshDeserialize, BorshSerialize, BorshSchema)]
 pub enum AssociatedTokenAccountInstruction {
     /// Creates an associated token account for the given wallet address and token mint
+    /// Returns an error if the account exists.
     ///
     ///   0. `[writeable,signer]` Funding account (must be a system account)
     ///   1. `[writeable]` Associated token account address to be created
@@ -22,20 +24,34 @@ pub enum AssociatedTokenAccountInstruction {
     ///   4. `[]` System program
     ///   5. `[]` SPL Token program
     Create,
+    /// Creates an associated token account for the given wallet address and token mint,
+    /// if it doesn't already exist.  Returns an error if the account exists,
+    /// but with a different owner.
+    ///
+    ///   0. `[writeable,signer]` Funding account (must be a system account)
+    ///   1. `[writeable]` Associated token account address to be created
+    ///   2. `[]` Wallet address for the new associated token account
+    ///   3. `[]` The token mint for the new associated token account
+    ///   4. `[]` System program
+    ///   5. `[]` SPL Token program
+    CreateIdempotent,
 }
 
-/// Creates CreateAssociatedTokenAccount instruction
-pub fn create_associated_token_account(
+fn build_associated_token_account_instruction(
     funding_address: &Pubkey,
     wallet_address: &Pubkey,
     token_mint_address: &Pubkey,
     token_program_id: &Pubkey,
+    instruction: AssociatedTokenAccountInstruction,
 ) -> Instruction {
     let associated_account_address =
         get_associated_token_address(wallet_address, token_mint_address);
-
-    let instruction_data = AssociatedTokenAccountInstruction::Create {};
-
+    // safety check, assert if not a creation instruction
+    assert_matches!(
+        instruction,
+        AssociatedTokenAccountInstruction::Create
+            | AssociatedTokenAccountInstruction::CreateIdempotent
+    );
     Instruction {
         program_id: id(),
         accounts: vec![
@@ -46,6 +62,38 @@ pub fn create_associated_token_account(
             AccountMeta::new_readonly(solana_program::system_program::id(), false),
             AccountMeta::new_readonly(*token_program_id, false),
         ],
-        data: instruction_data.try_to_vec().unwrap(),
+        data: instruction.try_to_vec().unwrap(),
     }
 }
+
+/// Creates Create instruction
+pub fn create_associated_token_account(
+    funding_address: &Pubkey,
+    wallet_address: &Pubkey,
+    token_mint_address: &Pubkey,
+    token_program_id: &Pubkey,
+) -> Instruction {
+    build_associated_token_account_instruction(
+        funding_address,
+        wallet_address,
+        token_mint_address,
+        token_program_id,
+        AssociatedTokenAccountInstruction::Create,
+    )
+}
+
+/// Creates CreateIdempotent instruction
+pub fn create_associated_token_account_idempotent(
+    funding_address: &Pubkey,
+    wallet_address: &Pubkey,
+    token_mint_address: &Pubkey,
+    token_program_id: &Pubkey,
+) -> Instruction {
+    build_associated_token_account_instruction(
+        funding_address,
+        wallet_address,
+        token_mint_address,
+        token_program_id,
+        AssociatedTokenAccountInstruction::CreateIdempotent,
+    )
+}

associated-token-account/program/src/instruction.rs

@@ -3,6 +3,7 @@
 #![forbid(unsafe_code)]
 
 mod entrypoint;
+pub mod error;
 pub mod instruction;
 pub mod processor;
 pub mod tools;

associated-token-account/program/src/lib.rs

@@ -2,6 +2,7 @@
 
 use {
     crate::{
+        error::AssociatedTokenAccountError,
         instruction::AssociatedTokenAccountInstruction,
         tools::account::{create_pda_account, get_account_len},
         *,
@@ -15,10 +16,21 @@ use {
         program_error::ProgramError,
         pubkey::Pubkey,
         rent::Rent,
+        system_program,
         sysvar::Sysvar,
     },
+    spl_token::{extension::StateWithExtensions, state::Account},
 };
 
+/// Specify when to create the associated token account
+#[derive(PartialEq)]
+enum CreateMode {
+    /// Always try to create the ATA
+    Always,
+    /// Only try to create the ATA if non-existent
+    Idempotent,
+}
+
 /// Instruction processor
 pub fn process_instruction(
     program_id: &Pubkey,
@@ -35,16 +47,20 @@ pub fn process_instruction(
     msg!("{:?}", instruction);
 
     match instruction {
-        AssociatedTokenAccountInstruction::Create {} => {
-            process_create_associated_token_account(program_id, accounts)
+        AssociatedTokenAccountInstruction::Create => {
+            process_create_associated_token_account(program_id, accounts, CreateMode::Always)
+        }
+        AssociatedTokenAccountInstruction::CreateIdempotent => {
+            process_create_associated_token_account(program_id, accounts, CreateMode::Idempotent)
         }
     }
 }
 
 /// Processes CreateAssociatedTokenAccount instruction
-pub fn process_create_associated_token_account(
+fn process_create_associated_token_account(
     program_id: &Pubkey,
     accounts: &[AccountInfo],
+    create_mode: CreateMode,
 ) -> ProgramResult {
     let account_info_iter = &mut accounts.iter();
 
@@ -56,6 +72,26 @@ pub fn process_create_associated_token_account(
     let spl_token_program_info = next_account_info(account_info_iter)?;
     let spl_token_program_id = spl_token_program_info.key;
 
+    if create_mode == CreateMode::Idempotent
+        && associated_token_account_info.owner == spl_token_program_id
+    {
+        let ata_data = associated_token_account_info.data.borrow();
+        if let Ok(associated_token_account) = StateWithExtensions::<Account>::unpack(&ata_data) {
+            if associated_token_account.base.owner != *wallet_account_info.key {
+                let error = AssociatedTokenAccountError::InvalidOwner;
+                msg!("{}", error);
+                return Err(error.into());
+            }
+            if associated_token_account.base.mint != *spl_token_mint_info.key {
+                return Err(ProgramError::InvalidAccountData);
+            }
+            return Ok(());
+        }
+    }
+    if *associated_token_account_info.owner != system_program::id() {
+        return Err(ProgramError::IllegalOwner);
+    }
+
     let rent = Rent::get()?;
 
     let (associated_token_address, bump_seed) = get_associated_token_address_and_bump_seed_internal(