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>

Author: CriesofCarrots

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)]