token-cli: Memo transfer extension (#3525)

token-cli: Add support for required transfer memos

Author: 0xAndoroid

token/cli/src/main.rs

@@ -44,7 +44,8 @@ use spl_associated_token_account::{
 use spl_token_2022::{
     extension::{
         interest_bearing_mint, interest_bearing_mint::InterestBearingConfig,
-        mint_close_authority::MintCloseAuthority, StateWithExtensionsOwned,
+        memo_transfer::MemoTransfer, mint_close_authority::MintCloseAuthority, ExtensionType,
+        StateWithExtensionsOwned,
     },
     instruction::*,
     state::{Account, Mint, Multisig},
@@ -153,6 +154,8 @@ pub enum CommandName {
     Display,
     Gc,
     SyncNative,
+    EnableRequiredTransferMemos,
+    DisableRequiredTransferMemos,
 }
 impl fmt::Display for CommandName {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
@@ -1887,6 +1890,103 @@ async fn command_sync_native(
     })
 }
 
+// Both enable_required_transfer_mesos and disable_required_transfer_mesos
+// Switches with enable_memos bool
+async fn command_required_transfer_memos(
+    config: &Config<'_>,
+    token_account_address: Pubkey,
+    owner: Pubkey,
+    bulk_signers: BulkSigners,
+    enable_memos: bool,
+) -> CommandResult {
+    if config.sign_only {
+        panic!("Config can not be sign only for enabling/disabling required transfer memos.");
+    }
+    let account_fetch = config
+        .rpc_client
+        .get_account(&token_account_address)
+        .await
+        .map_err(|err| {
+            format!(
+                "Token account {} does not exist: {}",
+                token_account_address, err
+            )
+        })?;
+    let program_id = config.program_id;
+    config.get_account_checked(&token_account_address).await?;
+    let mut instructions: Vec<Instruction> = Vec::new();
+    // Reallocation (if needed)
+    let current_account_len = account_fetch.data.len();
+    let state_with_extension = StateWithExtensionsOwned::<Account>::unpack(account_fetch.data)?;
+    let mut existing_extensions: Vec<ExtensionType> = state_with_extension.get_extension_types()?;
+    if existing_extensions.contains(&ExtensionType::MemoTransfer) {
+        let extension_data: bool = state_with_extension
+            .get_extension::<MemoTransfer>()?
+            .require_incoming_transfer_memos
+            .into();
+        if extension_data == enable_memos {
+            return Ok(format!(
+                "Required memo transfer was already {}",
+                if extension_data {
+                    "enabled"
+                } else {
+                    "disabled"
+                }
+            ));
+        }
+    } else {
+        existing_extensions.push(ExtensionType::MemoTransfer);
+        let needed_account_len = ExtensionType::get_account_len::<Account>(&existing_extensions);
+        if needed_account_len > current_account_len {
+            instructions.push(reallocate(
+                &program_id,
+                &token_account_address,
+                &config.fee_payer.pubkey(),
+                &owner,
+                &config.multisigner_pubkeys,
+                &existing_extensions,
+            )?);
+        }
+    }
+    if enable_memos {
+        instructions.push(
+            spl_token_2022::extension::memo_transfer::instruction::enable_required_transfer_memos(
+                &program_id,
+                &token_account_address,
+                &owner,
+                &config.multisigner_pubkeys,
+            )?,
+        );
+    } else {
+        instructions.push(
+            spl_token_2022::extension::memo_transfer::instruction::disable_required_transfer_memos(
+                &program_id,
+                &token_account_address,
+                &owner,
+                &config.multisigner_pubkeys,
+            )?,
+        );
+    }
+    let tx_return = handle_tx(
+        &CliSignerInfo {
+            signers: bulk_signers,
+        },
+        config,
+        false,
+        0,
+        instructions,
+    )
+    .await?;
+    Ok(match tx_return {
+        TransactionReturnData::CliSignature(signature) => {
+            config.output_format.formatted_string(&signature)
+        }
+        TransactionReturnData::CliSignOnlyData(sign_only_data) => {
+            config.output_format.formatted_string(&sign_only_data)
+        }
+    })
+}
+
 struct SignOnlyNeedsFullMintSpec {}
 impl offline::ArgsConfig for SignOnlyNeedsFullMintSpec {
     fn sign_only_arg<'a, 'b>(&self, arg: Arg<'a, 'b>) -> Arg<'a, 'b> {
@@ -2821,6 +2921,44 @@ fn app<'a, 'b>(
                         .help("Specify the specific token account address to sync"),
                 ),
         )
+        .subcommand(
+            SubCommand::with_name(CommandName::EnableRequiredTransferMemos.into())
+                .about("Enable required transfer memos for token account")
+                .arg(
+                    Arg::with_name("account")
+                        .validator(is_valid_pubkey)
+                        .value_name("TOKEN_ACCOUNT_ADDRESS")
+                        .takes_value(true)
+                        .index(1)
+                        .required(true)
+                        .help("The address of the token account to enable required transfer memos")
+                )
+                .arg(
+                    owner_address_arg()
+                )
+                .arg(multisig_signer_arg())
+                .nonce_args(true)
+                .offline_args()
+        )
+        .subcommand(
+            SubCommand::with_name(CommandName::DisableRequiredTransferMemos.into())
+                .about("Disable required transfer memos for token account")
+                .arg(
+                    Arg::with_name("account")
+                        .validator(is_valid_pubkey)
+                        .value_name("TOKEN_ACCOUNT_ADDRESS")
+                        .takes_value(true)
+                        .index(1)
+                        .required(true)
+                        .help("The address of the token account to disable required transfer memos"),
+                )
+                .arg(
+                    owner_address_arg()
+                )
+                .arg(multisig_signer_arg())
+                .nonce_args(true)
+                .offline_args()
+        )
 }
 
 #[tokio::main]
@@ -3356,6 +3494,28 @@ async fn process_command<'a>(
                 .await;
             command_sync_native(address, bulk_signers, config).await
         }
+        (CommandName::EnableRequiredTransferMemos, arg_matches) => {
+            let (owner_signer, owner) =
+                config.signer_or_default(arg_matches, "owner", &mut wallet_manager);
+            if !bulk_signers.contains(&owner_signer) {
+                bulk_signers.push(owner_signer);
+            }
+            // Since account is required argument it will always be present
+            let token_account =
+                config.pubkey_or_default(arg_matches, "account", &mut wallet_manager);
+            command_required_transfer_memos(config, token_account, owner, bulk_signers, true).await
+        }
+        (CommandName::DisableRequiredTransferMemos, arg_matches) => {
+            let (owner_signer, owner) =
+                config.signer_or_default(arg_matches, "owner", &mut wallet_manager);
+            if !bulk_signers.contains(&owner_signer) {
+                bulk_signers.push(owner_signer);
+            }
+            // Since account is required argument it will always be present
+            let token_account =
+                config.pubkey_or_default(arg_matches, "account", &mut wallet_manager);
+            command_required_transfer_memos(config, token_account, owner, bulk_signers, false).await
+        }
     }
 }
 
@@ -4546,4 +4706,60 @@ mod tests {
         let account = config.rpc_client.get_account(&token_pubkey).await;
         assert!(account.is_err());
     }
+
+    #[tokio::test]
+    #[serial]
+    async fn required_transfer_memos() {
+        let (test_validator, payer) = new_validator_for_test().await;
+        let program_id = spl_token_2022::id();
+        let config = test_config(&test_validator, &payer, &program_id);
+        let token = create_token(&config, &payer).await;
+        let token_account = create_associated_account(&config, &payer, token).await;
+        let result = process_test_command(
+            &config,
+            &payer,
+            &[
+                "spl-token",
+                CommandName::EnableRequiredTransferMemos.into(),
+                &token_account.to_string(),
+            ],
+        )
+        .await;
+        result.unwrap();
+        let extensions = StateWithExtensionsOwned::<Account>::unpack(
+            config
+                .rpc_client
+                .get_account(&token_account)
+                .await
+                .unwrap()
+                .data,
+        )
+        .unwrap();
+        let memo_transfer = extensions.get_extension::<MemoTransfer>().unwrap();
+        let enabled: bool = memo_transfer.require_incoming_transfer_memos.into();
+        assert!(enabled);
+        let result = process_test_command(
+            &config,
+            &payer,
+            &[
+                "spl-token",
+                CommandName::DisableRequiredTransferMemos.into(),
+                &token_account.to_string(),
+            ],
+        )
+        .await;
+        result.unwrap();
+        let extensions = StateWithExtensionsOwned::<Account>::unpack(
+            config
+                .rpc_client
+                .get_account(&token_account)
+                .await
+                .unwrap()
+                .data,
+        )
+        .unwrap();
+        let memo_transfer = extensions.get_extension::<MemoTransfer>().unwrap();
+        let enabled: bool = memo_transfer.require_incoming_transfer_memos.into();
+        assert!(!enabled);
+    }
 }