token-2022: Add MemoTransfer extension (#2822)

* Style nits

* Initial MemoTransfer extension

* Stub in check for memo

* Add memo-transfer token-client method

* Add MemoTransfer tests

* Add immutable get_extension, and clean up mod

* Update token/program-2022/src/extension/memo_transfer/instruction.rs

Co-authored-by: Jon Cinque <jon.cinque@gmail.com>

* Update token/rust/src/token.rs

Co-authored-by: Jon Cinque <jon.cinque@gmail.com>

Co-authored-by: Jon Cinque <jon.cinque@gmail.com>

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

@@ -34,7 +34,7 @@ pub enum ConfidentialTransferInstruction {
     /// Accounts expected by this instruction:
     ///
     ///   0. `[writable]` The SPL Token mint
-    //
+    ///
     /// Data expected by this instruction:
     ///   `ConfidentialTransferMint`
     ///

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

@@ -28,7 +28,7 @@ pub enum DefaultAccountStateInstruction {
     /// Accounts expected by this instruction:
     ///
     ///   0. `[writable]` The mint to initialize.
-    //
+    ///
     /// Data expected by this instruction:
     ///   `crate::state::AccountState`
     ///
@@ -46,7 +46,7 @@ pub enum DefaultAccountStateInstruction {
     ///   0. `[writable]` The mint.
     ///   1. `[]` The mint's multisignature freeze authority.
     ///   2. ..2+M `[signer]` M signer accounts.
-    //
+    ///
     /// Data expected by this instruction:
     ///   `crate::state::AccountState`
     ///

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

@@ -0,0 +1,113 @@
+use {
+    crate::{check_program_account, error::TokenError, instruction::TokenInstruction},
+    num_enum::{IntoPrimitive, TryFromPrimitive},
+    solana_program::{
+        instruction::{AccountMeta, Instruction},
+        program_error::ProgramError,
+        pubkey::Pubkey,
+    },
+    std::convert::TryFrom,
+};
+
+/// Default Account State extension instructions
+#[derive(Clone, Copy, Debug, PartialEq, IntoPrimitive, TryFromPrimitive)]
+#[repr(u8)]
+pub enum RequiredMemoTransfersInstruction {
+    /// Require memos for transfers into this Account. Adds the MemoTransfer extension to the
+    /// Account, if it doesn't already exist.
+    ///
+    /// Accounts expected by this instruction:
+    ///
+    ///   0. `[writable]` The account to update.
+    ///   1. `[signer]` The account's owner.
+    ///
+    ///   * Multisignature authority
+    ///   0. `[writable]` The account to update.
+    ///   1. `[]` The account's multisignature owner.
+    ///   2. ..2+M `[signer]` M signer accounts.
+    ///
+    Enable,
+    /// Stop requiring memos for transfers into this Account.
+    ///
+    /// Fails if the account does not have the extension present.
+    ///
+    /// Accounts expected by this instruction:
+    ///
+    ///   0. `[writable]` The account to update.
+    ///   1. `[signer]` The account's owner.
+    ///
+    ///   * Multisignature authority
+    ///   0. `[writable]` The account to update.
+    ///   1. `[]`  The account's multisignature owner.
+    ///   2. ..2+M `[signer]` M signer accounts.
+    ///
+    Disable,
+}
+
+pub(crate) fn decode_instruction(
+    input: &[u8],
+) -> Result<RequiredMemoTransfersInstruction, ProgramError> {
+    if input.len() != 1 {
+        return Err(TokenError::InvalidInstruction.into());
+    }
+    RequiredMemoTransfersInstruction::try_from(input[0])
+        .map_err(|_| TokenError::InvalidInstruction.into())
+}
+
+fn encode_instruction(
+    token_program_id: &Pubkey,
+    accounts: Vec<AccountMeta>,
+    instruction_type: RequiredMemoTransfersInstruction,
+) -> Instruction {
+    let mut data = TokenInstruction::MemoTransferExtension.pack();
+    data.push(instruction_type.into());
+    Instruction {
+        program_id: *token_program_id,
+        accounts,
+        data,
+    }
+}
+
+/// Create an `Enable` instruction
+pub fn enable_required_transfer_memos(
+    token_program_id: &Pubkey,
+    account: &Pubkey,
+    owner: &Pubkey,
+    signers: &[&Pubkey],
+) -> Result<Instruction, ProgramError> {
+    check_program_account(token_program_id)?;
+    let mut accounts = vec![
+        AccountMeta::new(*account, false),
+        AccountMeta::new_readonly(*owner, signers.is_empty()),
+    ];
+    for signer_pubkey in signers.iter() {
+        accounts.push(AccountMeta::new_readonly(**signer_pubkey, true));
+    }
+    Ok(encode_instruction(
+        token_program_id,
+        accounts,
+        RequiredMemoTransfersInstruction::Enable,
+    ))
+}
+
+/// Create a `Disable` instruction
+pub fn disable_required_transfer_memos(
+    token_program_id: &Pubkey,
+    account: &Pubkey,
+    owner: &Pubkey,
+    signers: &[&Pubkey],
+) -> Result<Instruction, ProgramError> {
+    check_program_account(token_program_id)?;
+    let mut accounts = vec![
+        AccountMeta::new(*account, false),
+        AccountMeta::new_readonly(*owner, signers.is_empty()),
+    ];
+    for signer_pubkey in signers.iter() {
+        accounts.push(AccountMeta::new_readonly(**signer_pubkey, true));
+    }
+    Ok(encode_instruction(
+        token_program_id,
+        accounts,
+        RequiredMemoTransfersInstruction::Disable,
+    ))
+}

token/program-2022/src/extension/memo_transfer/mod.rs

@@ -0,0 +1,33 @@
+use {
+    crate::{
+        extension::{Extension, ExtensionType, StateWithExtensionsMut},
+        pod::PodBool,
+        state::Account,
+    },
+    bytemuck::{Pod, Zeroable},
+};
+
+/// Memo Transfer extension instructions
+pub mod instruction;
+
+/// Memo Transfer extension processor
+pub mod processor;
+
+/// Memo Transfer extension for Accounts
+#[repr(C)]
+#[derive(Clone, Copy, Debug, Default, PartialEq, Pod, Zeroable)]
+pub struct MemoTransfer {
+    /// Require transfers into this account to be accompanied by a memo
+    pub require_incoming_transfer_memos: PodBool,
+}
+impl Extension for MemoTransfer {
+    const TYPE: ExtensionType = ExtensionType::MemoTransfer;
+}
+
+/// Determine if a memo is required for transfers into this account
+pub fn memo_required(account_state: &StateWithExtensionsMut<Account>) -> bool {
+    if let Ok(extension) = account_state.get_extension::<MemoTransfer>() {
+        return extension.require_incoming_transfer_memos.into();
+    }
+    false
+}

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

@@ -0,0 +1,98 @@
+use {
+    crate::{
+        check_program_account,
+        extension::{
+            memo_transfer::{
+                instruction::{decode_instruction, RequiredMemoTransfersInstruction},
+                MemoTransfer,
+            },
+            StateWithExtensionsMut,
+        },
+        processor::Processor,
+        state::Account,
+    },
+    solana_program::{
+        account_info::{next_account_info, AccountInfo},
+        entrypoint::ProgramResult,
+        msg,
+        pubkey::Pubkey,
+    },
+};
+
+fn process_enable_required_memo_transfers(
+    program_id: &Pubkey,
+    accounts: &[AccountInfo],
+) -> ProgramResult {
+    let account_info_iter = &mut accounts.iter();
+    let token_account_info = next_account_info(account_info_iter)?;
+    let owner_info = next_account_info(account_info_iter)?;
+    let owner_info_data_len = owner_info.data_len();
+
+    let mut account_data = token_account_info.data.borrow_mut();
+    let mut account = StateWithExtensionsMut::<Account>::unpack(&mut account_data)?;
+
+    Processor::validate_owner(
+        program_id,
+        &account.base.owner,
+        owner_info,
+        owner_info_data_len,
+        account_info_iter.as_slice(),
+    )?;
+
+    let extension = if let Ok(extension) = account.get_extension_mut::<MemoTransfer>() {
+        extension
+    } else {
+        account.init_extension::<MemoTransfer>()?
+    };
+    extension.require_incoming_transfer_memos = true.into();
+    Ok(())
+}
+
+fn process_diasble_required_memo_transfers(
+    program_id: &Pubkey,
+    accounts: &[AccountInfo],
+) -> ProgramResult {
+    let account_info_iter = &mut accounts.iter();
+    let token_account_info = next_account_info(account_info_iter)?;
+    let owner_info = next_account_info(account_info_iter)?;
+    let owner_info_data_len = owner_info.data_len();
+
+    let mut account_data = token_account_info.data.borrow_mut();
+    let mut account = StateWithExtensionsMut::<Account>::unpack(&mut account_data)?;
+
+    Processor::validate_owner(
+        program_id,
+        &account.base.owner,
+        owner_info,
+        owner_info_data_len,
+        account_info_iter.as_slice(),
+    )?;
+
+    let extension = if let Ok(extension) = account.get_extension_mut::<MemoTransfer>() {
+        extension
+    } else {
+        account.init_extension::<MemoTransfer>()?
+    };
+    extension.require_incoming_transfer_memos = false.into();
+    Ok(())
+}
+
+pub(crate) fn process_instruction(
+    program_id: &Pubkey,
+    accounts: &[AccountInfo],
+    input: &[u8],
+) -> ProgramResult {
+    check_program_account(program_id)?;
+
+    let instruction = decode_instruction(input)?;
+    match instruction {
+        RequiredMemoTransfersInstruction::Enable => {
+            msg!("RequiredMemoTransfersInstruction::Enable");
+            process_enable_required_memo_transfers(program_id, accounts)
+        }
+        RequiredMemoTransfersInstruction::Disable => {
+            msg!("RequiredMemoTransfersInstruction::Disable");
+            process_diasble_required_memo_transfers(program_id, accounts)
+        }
+    }
+}

token/program-2022/src/extension/mod.rs

@@ -7,6 +7,7 @@ use {
             confidential_transfer::{ConfidentialTransferAccount, ConfidentialTransferMint},
             default_account_state::DefaultAccountState,
             immutable_owner::ImmutableOwner,
+            memo_transfer::MemoTransfer,
             mint_close_authority::MintCloseAuthority,
             transfer_fee::{TransferFeeAmount, TransferFeeConfig},
         },
@@ -31,6 +32,8 @@ pub mod confidential_transfer;
 pub mod default_account_state;
 /// Immutable Owner extension
 pub mod immutable_owner;
+/// Memo Transfer extension
+pub mod memo_transfer;
 /// Mint Close Authority extension
 pub mod mint_close_authority;
 /// Utility to reallocate token accounts
@@ -431,11 +434,30 @@ impl<'data, S: BaseState> StateWithExtensionsMut<'data, S> {
         }
     }
 
-    /// Unpack a portion of the TLV data as the desired type
+    /// Unpack a portion of the TLV data as the desired type that allows modifying the type
     pub fn get_extension_mut<V: Extension>(&mut self) -> Result<&mut V, ProgramError> {
         self.init_or_get_extension(false)
     }
 
+    /// Unpack a portion of the TLV data as the desired type
+    pub fn get_extension<V: Extension>(&self) -> Result<&V, ProgramError> {
+        if V::TYPE.get_account_type() != S::ACCOUNT_TYPE {
+            return Err(ProgramError::InvalidAccountData);
+        }
+        let TlvIndices {
+            type_start,
+            length_start,
+            value_start,
+        } = get_extension_indices::<V>(self.tlv_data, false)?;
+
+        if self.tlv_data[type_start..].len() < V::TYPE.get_tlv_len() {
+            return Err(ProgramError::InvalidAccountData);
+        }
+        let length = pod_from_bytes::<Length>(&self.tlv_data[length_start..value_start])?;
+        let value_end = value_start.saturating_add(usize::from(*length));
+        pod_from_bytes::<V>(&self.tlv_data[value_start..value_end])
+    }
+
     /// Packs base state data into the base data portion
     pub fn pack_base(&mut self) {
         S::pack_into_slice(&self.base, self.base_data);
@@ -561,6 +583,8 @@ pub enum ExtensionType {
     DefaultAccountState,
     /// Indicates that the Account owner authority cannot be changed
     ImmutableOwner,
+    /// Require inbound transfers to have memo
+    MemoTransfer,
     /// Padding extension used to make an account exactly Multisig::LEN, used for testing
     #[cfg(test)]
     AccountPaddingTest = u16::MAX - 1,
@@ -598,6 +622,7 @@ impl ExtensionType {
                 pod_get_packed_len::<ConfidentialTransferAccount>()
             }
             ExtensionType::DefaultAccountState => pod_get_packed_len::<DefaultAccountState>(),
+            ExtensionType::MemoTransfer => pod_get_packed_len::<MemoTransfer>(),
             #[cfg(test)]
             ExtensionType::AccountPaddingTest => pod_get_packed_len::<AccountPaddingTest>(),
             #[cfg(test)]
@@ -655,7 +680,8 @@ impl ExtensionType {
             | ExtensionType::DefaultAccountState => AccountType::Mint,
             ExtensionType::ImmutableOwner
             | ExtensionType::TransferFeeAmount
-            | ExtensionType::ConfidentialTransferAccount => AccountType::Account,
+            | ExtensionType::ConfidentialTransferAccount
+            | ExtensionType::MemoTransfer => AccountType::Account,
             #[cfg(test)]
             ExtensionType::AccountPaddingTest => AccountType::Account,
             #[cfg(test)]

token/program-2022/src/instruction.rs

@@ -522,6 +522,11 @@ pub enum TokenInstruction {
         /// New extension types to include in the reallocated account
         extension_types: Vec<ExtensionType>,
     },
+    /// The common instruction prefix for Memo Transfer account extension instructions.
+    ///
+    /// See `extension::memo_transfer::instruction::RequiredMemoTransfersInstruction` for
+    /// further details about the extended instructions that share this instruction prefix
+    MemoTransferExtension,
 }
 impl TokenInstruction {
     /// Unpacks a byte buffer into a [TokenInstruction](enum.TokenInstruction.html).
@@ -640,6 +645,7 @@ impl TokenInstruction {
                 }
                 Self::Reallocate { extension_types }
             }
+            28 => Self::MemoTransferExtension,
             _ => return Err(TokenError::InvalidInstruction.into()),
         })
     }
@@ -772,6 +778,9 @@ impl TokenInstruction {
                     buf.extend_from_slice(&<[u8; 2]>::from(*extension_type));
                 }
             }
+            &Self::MemoTransferExtension => {
+                buf.push(28);
+            }
         };
         buf
     }

token/program-2022/src/pod.rs

@@ -74,6 +74,12 @@ impl From<&PodBool> for bool {
     }
 }
 
+impl From<PodBool> for bool {
+    fn from(b: PodBool) -> Self {
+        b.0 != 0
+    }
+}
+
 /// The standard `u16` can cause alignment issues when placed in a `Pod`, define a replacement that
 /// is usable in all `Pod`s
 #[derive(Clone, Copy, Debug, Default, PartialEq, Pod, Zeroable)]

token/program-2022/src/processor.rs

@@ -8,6 +8,7 @@ use {
             confidential_transfer::{self, ConfidentialTransferAccount},
             default_account_state::{self, DefaultAccountState},
             immutable_owner::ImmutableOwner,
+            memo_transfer::{self, memo_required},
             mint_close_authority::MintCloseAuthority,
             reallocate,
             transfer_fee::{self, TransferFeeAmount, TransferFeeConfig},
@@ -372,6 +373,10 @@ impl Processor {
             return Err(TokenError::MintMismatch.into());
         }
 
+        if memo_required(&dest_account) {
+            // TODO: use get_processed_instructions syscall to check for memo
+        }
+
         source_account.base.amount = source_account
             .base
             .amount
@@ -1153,6 +1158,9 @@ impl Processor {
                 msg!("Instruction: Reallocate");
                 reallocate::process_reallocate(program_id, accounts, extension_types)
             }
+            TokenInstruction::MemoTransferExtension => {
+                memo_transfer::processor::process_instruction(program_id, accounts, &input[1..])
+            }
         }
     }
 

token/program-2022/tests/memo_transfer.rs

@@ -0,0 +1,117 @@
+#![cfg(feature = "test-bpf")]
+
+mod program_test;
+use {
+    program_test::{TestContext, TokenContext},
+    solana_program_test::tokio,
+    solana_sdk::{pubkey::Pubkey, signature::Signer},
+    spl_token_2022::extension::{memo_transfer::MemoTransfer, ExtensionType},
+};
+
+async fn test_memo_transfers(
+    token_context: TokenContext,
+    alice_account: Pubkey,
+    bob_account: Pubkey,
+) {
+    let TokenContext {
+        mint_authority,
+        token,
+        alice,
+        bob,
+        ..
+    } = token_context;
+
+    // mint tokens
+    token
+        .mint_to(&alice_account, &mint_authority, 4242)
+        .await
+        .unwrap();
+
+    // require memo transfers into bob_account
+    token
+        .enable_required_transfer_memos(&bob_account, &bob)
+        .await
+        .unwrap();
+
+    let bob_state = token.get_account_info(&bob_account).await.unwrap();
+    let extension = bob_state.get_extension::<MemoTransfer>().unwrap();
+    assert!(bool::from(extension.require_incoming_transfer_memos));
+
+    // attempt to transfer from alice to bob without memo
+    // TODO: should fail when token/program-2022/src/processor.rs#L376 is completed
+    token
+        .transfer_unchecked(&alice_account, &bob_account, &alice, 10)
+        .await
+        .unwrap();
+    let bob_state = token.get_account_info(&bob_account).await.unwrap();
+    assert_eq!(bob_state.base.amount, 10);
+
+    // stop requiring memo transfers into bob_account
+    token
+        .disable_required_transfer_memos(&bob_account, &bob)
+        .await
+        .unwrap();
+
+    // transfer from alice to bob without memo
+    token
+        .transfer_unchecked(&alice_account, &bob_account, &alice, 11)
+        .await
+        .unwrap();
+    let bob_state = token.get_account_info(&bob_account).await.unwrap();
+    assert_eq!(bob_state.base.amount, 21);
+}
+
+#[tokio::test]
+async fn require_memo_transfers_without_realloc() {
+    let mut context = TestContext::new().await;
+    context.init_token_with_mint(vec![]).await.unwrap();
+    let token_context = context.token_context.unwrap();
+
+    // create token accounts
+    let alice_account = token_context
+        .token
+        .create_auxiliary_token_account(&token_context.alice, &token_context.alice.pubkey())
+        .await
+        .unwrap();
+    let bob_account = token_context
+        .token
+        .create_auxiliary_token_account_with_extension_space(
+            &token_context.bob,
+            &token_context.bob.pubkey(),
+            vec![ExtensionType::MemoTransfer],
+        )
+        .await
+        .unwrap();
+
+    test_memo_transfers(token_context, alice_account, bob_account).await;
+}
+
+#[tokio::test]
+async fn require_memo_transfers_with_realloc() {
+    let mut context = TestContext::new().await;
+    context.init_token_with_mint(vec![]).await.unwrap();
+    let token_context = context.token_context.unwrap();
+
+    // create token accounts
+    let alice_account = token_context
+        .token
+        .create_auxiliary_token_account(&token_context.alice, &token_context.alice.pubkey())
+        .await
+        .unwrap();
+    let bob_account = token_context
+        .token
+        .create_auxiliary_token_account(&token_context.bob, &token_context.bob.pubkey())
+        .await
+        .unwrap();
+    token_context
+        .token
+        .reallocate(
+            &token_context.bob.pubkey(),
+            &token_context.bob,
+            &[ExtensionType::MemoTransfer],
+        )
+        .await
+        .unwrap();
+
+    test_memo_transfers(token_context, alice_account, bob_account).await;
+}

token/rust/src/token.rs

@@ -12,7 +12,9 @@ use spl_associated_token_account::{
     get_associated_token_address_with_program_id, instruction::create_associated_token_account,
 };
 use spl_token_2022::{
-    extension::{default_account_state, transfer_fee, ExtensionType, StateWithExtensionsOwned},
+    extension::{
+        default_account_state, memo_transfer, transfer_fee, ExtensionType, StateWithExtensionsOwned,
+    },
     instruction,
     state::{Account, AccountState, Mint},
 };
@@ -257,11 +259,28 @@ where
         &self,
         account: &S,
         owner: &Pubkey,
+    ) -> TokenResult<Pubkey> {
+        self.create_auxiliary_token_account_with_extension_space(account, owner, vec![])
+            .await
+    }
+
+    /// Create and initialize a new token account.
+    pub async fn create_auxiliary_token_account_with_extension_space(
+        &self,
+        account: &S,
+        owner: &Pubkey,
+        extensions: Vec<ExtensionType>,
     ) -> TokenResult<Pubkey> {
         let state = self.get_mint_info().await?;
         let mint_extensions: Vec<ExtensionType> = state.get_extension_types()?;
-        let extensions = ExtensionType::get_required_init_account_extensions(&mint_extensions);
+        let mut required_extensions =
-        let space = ExtensionType::get_account_len::<Account>(&extensions);
+            ExtensionType::get_required_init_account_extensions(&mint_extensions);
+        for extension_type in extensions.into_iter() {
+            if !required_extensions.contains(&extension_type) {
+                required_extensions.push(extension_type);
+            }
+        }
+        let space = ExtensionType::get_account_len::<Account>(&required_extensions);
         self.process_ixs(
             &[
                 system_instruction::create_account(
@@ -750,4 +769,40 @@ where
         )
         .await
     }
+
+    /// Require memos on transfers into this account
+    pub async fn enable_required_transfer_memos<S2: Signer>(
+        &self,
+        account: &Pubkey,
+        authority: &S2,
+    ) -> TokenResult<T::Output> {
+        self.process_ixs(
+            &[memo_transfer::instruction::enable_required_transfer_memos(
+                &self.program_id,
+                account,
+                &authority.pubkey(),
+                &[],
+            )?],
+            &[authority],
+        )
+        .await
+    }
+
+    /// Stop requiring memos on transfers into this account
+    pub async fn disable_required_transfer_memos<S2: Signer>(
+        &self,
+        account: &Pubkey,
+        authority: &S2,
+    ) -> TokenResult<T::Output> {
+        self.process_ixs(
+            &[memo_transfer::instruction::disable_required_transfer_memos(
+                &self.program_id,
+                account,
+                &authority.pubkey(),
+                &[],
+            )?],
+            &[authority],
+        )
+        .await
+    }
 }