token-cli: Figure out program_id dynamically (#3354)

* Resolve program id with mint info, use it where possible

* Avoid using `config.program_id` when we can figure it out

* Rename function

Author: joncinque

token/cli/src/bench.rs

@@ -205,14 +205,9 @@ pub(crate) async fn bench_process_command(
             let (owner_signer, owner) =
                 config.signer_or_default(arg_matches, "owner", wallet_manager);
             signers.push(owner_signer);
-            let from = pubkey_of_signer(arg_matches, "from", wallet_manager)
-                .unwrap()
-                .unwrap_or_else(|| {
-                    get_associated_token_address_with_program_id(&owner, &token, &config.program_id)
-                });
-
+            let from = pubkey_of_signer(arg_matches, "from", wallet_manager).unwrap();
             command_deposit_into_or_withdraw_from(
-                config, signers, &token, n, &owner, ui_amount, &from, true,
+                config, signers, &token, n, &owner, ui_amount, from, true,
             )
             .await?;
         }
@@ -225,14 +220,9 @@ pub(crate) async fn bench_process_command(
             let (owner_signer, owner) =
                 config.signer_or_default(arg_matches, "owner", wallet_manager);
             signers.push(owner_signer);
-            let to = pubkey_of_signer(arg_matches, "to", wallet_manager)
-                .unwrap()
-                .unwrap_or_else(|| {
-                    get_associated_token_address_with_program_id(&owner, &token, &config.program_id)
-                });
-
+            let to = pubkey_of_signer(arg_matches, "to", wallet_manager).unwrap();
             command_deposit_into_or_withdraw_from(
-                config, signers, &token, n, &owner, ui_amount, &to, false,
+                config, signers, &token, n, &owner, ui_amount, to, false,
             )
             .await?;
         }
@@ -266,15 +256,18 @@ fn get_token_addresses_with_seed(
         .collect()
 }
 
-async fn is_valid_token(rpc_client: &RpcClient, token: &Pubkey) -> Result<(), Error> {
+async fn get_valid_mint_program_id(
+    rpc_client: &RpcClient,
+    token: &Pubkey,
+) -> Result<Pubkey, Error> {
     let mint_account = rpc_client
-        .get_account_data(token)
+        .get_account(token)
         .await
         .map_err(|err| format!("Token mint {} does not exist: {}", token, err))?;
 
-    Mint::unpack(&mint_account)
-        .map(|_| ())
-        .map_err(|err| format!("Invalid token mint {}: {}", token, err).into())
+    Mint::unpack(&mint_account.data)
+        .map_err(|err| format!("Invalid token mint {}: {}", token, err))?;
+    Ok(mint_account.owner)
 }
 
 async fn command_create_accounts(
@@ -287,16 +280,15 @@ async fn command_create_accounts(
     let rpc_client = &config.rpc_client;
 
     println!("Scanning accounts...");
-    is_valid_token(rpc_client, token).await?;
+    let program_id = get_valid_mint_program_id(rpc_client, token).await?;
 
     let minimum_balance_for_rent_exemption = rpc_client
         .get_minimum_balance_for_rent_exemption(Account::get_packed_len())
         .await?;
 
     let mut lamports_required = 0;
 
-    let token_addresses_with_seed =
-        get_token_addresses_with_seed(&config.program_id, token, owner, n);
+    let token_addresses_with_seed = get_token_addresses_with_seed(&program_id, token, owner, n);
     let mut messages = vec![];
     for address_chunk in token_addresses_with_seed.chunks(100) {
         let accounts_chunk = rpc_client
@@ -315,9 +307,9 @@ async fn command_create_accounts(
                             seed,
                             minimum_balance_for_rent_exemption,
                             Account::get_packed_len() as u64,
-                            &config.program_id,
+                            &program_id,
                         ),
-                        instruction::initialize_account(&config.program_id, address, token, owner)?,
+                        instruction::initialize_account(&program_id, address, token, owner)?,
                     ],
                     Some(&config.fee_payer),
                 ));
@@ -338,10 +330,9 @@ async fn command_close_accounts(
     let rpc_client = &config.rpc_client;
 
     println!("Scanning accounts...");
-    is_valid_token(rpc_client, token).await?;
+    let program_id = get_valid_mint_program_id(rpc_client, token).await?;
 
-    let token_addresses_with_seed =
-        get_token_addresses_with_seed(&config.program_id, token, owner, n);
+    let token_addresses_with_seed = get_token_addresses_with_seed(&program_id, token, owner, n);
     let mut messages = vec![];
     for address_chunk in token_addresses_with_seed.chunks(100) {
         let accounts_chunk = rpc_client
@@ -360,7 +351,7 @@ async fn command_close_accounts(
                         } else {
                             messages.push(Message::new(
                                 &[instruction::close_account(
-                                    &config.program_id,
+                                    &program_id,
                                     address,
                                     owner,
                                     owner,
@@ -389,23 +380,21 @@ async fn command_deposit_into_or_withdraw_from(
     n: usize,
     owner: &Pubkey,
     ui_amount: f64,
-    from_or_to: &Pubkey,
+    from_or_to: Option<Pubkey>,
     deposit_into: bool,
 ) -> Result<(), Error> {
     let rpc_client = &config.rpc_client;
 
     println!("Scanning accounts...");
-    is_valid_token(rpc_client, token).await?;
+    let program_id = get_valid_mint_program_id(rpc_client, token).await?;
 
-    let (mint_pubkey, decimals) =
-        crate::resolve_mint_info(config, from_or_to, Some(*token), None).await?;
-    if mint_pubkey != *token {
-        return Err(format!("Source account {} is not a {} token", from_or_to, token).into());
-    }
-    let amount = spl_token::ui_amount_to_amount(ui_amount, decimals);
+    let mint_info = config.get_mint_info(token, None).await?;
+    let from_or_to = from_or_to
+        .unwrap_or_else(|| get_associated_token_address_with_program_id(owner, token, &program_id));
+    config.check_account(&from_or_to, Some(*token)).await?;
+    let amount = spl_token::ui_amount_to_amount(ui_amount, mint_info.decimals);
 
-    let token_addresses_with_seed =
-        get_token_addresses_with_seed(&config.program_id, token, owner, n);
+    let token_addresses_with_seed = get_token_addresses_with_seed(&program_id, token, owner, n);
     let mut messages = vec![];
     for address_chunk in token_addresses_with_seed.chunks(100) {
         let accounts_chunk = rpc_client
@@ -416,14 +405,14 @@ async fn command_deposit_into_or_withdraw_from(
             if account.is_some() {
                 messages.push(Message::new(
                     &[instruction::transfer_checked(
-                        &config.program_id,
-                        if deposit_into { from_or_to } else { address },
+                        &program_id,
+                        if deposit_into { &from_or_to } else { address },
                         token,
-                        if deposit_into { address } else { from_or_to },
+                        if deposit_into { address } else { &from_or_to },
                         owner,
                         &[],
                         amount,
-                        decimals,
+                        mint_info.decimals,
                     )?],
                     Some(&config.fee_payer),
                 ));

token/cli/src/config.rs

@@ -1,3 +1,4 @@
+use crate::Error;
 use clap::ArgMatches;
 use solana_clap_utils::{
     input_parsers::pubkey_of_signer,
@@ -6,8 +7,9 @@ use solana_clap_utils::{
 use solana_cli_output::OutputFormat;
 use solana_client::nonblocking::rpc_client::RpcClient;
 use solana_remote_wallet::remote_wallet::RemoteWalletManager;
-use solana_sdk::{pubkey::Pubkey, signature::Signer};
+use solana_sdk::{program_pack::Pack, pubkey::Pubkey, signature::Signer};
 use spl_associated_token_account::*;
+use spl_token::state::{Account, Mint};
 use std::{process::exit, sync::Arc};
 
 #[cfg(test)]
@@ -21,6 +23,12 @@ pub(crate) enum KeypairOrPath {
     Path(String),
 }
 
+pub(crate) struct MintInfo {
+    pub program_id: Pubkey,
+    pub address: Pubkey,
+    pub decimals: u8,
+}
+
 pub(crate) struct Config<'a> {
     pub(crate) rpc_client: Arc<RpcClient>,
     pub(crate) websocket_url: String,
@@ -69,13 +77,29 @@ impl<'a> Config<'a> {
         }
 
         let token = token.unwrap();
+        let program_id = self.get_mint_info(&token, None).await.unwrap().program_id;
+        self.associated_token_address_for_token_and_program(
+            arg_matches,
+            wallet_manager,
+            &token,
+            &program_id,
+        )
+    }
+
+    pub(crate) fn associated_token_address_for_token_and_program(
+        &self,
+        arg_matches: &ArgMatches<'_>,
+        wallet_manager: &mut Option<Arc<RemoteWalletManager>>,
+        token: &Pubkey,
+        program_id: &Pubkey,
+    ) -> Pubkey {
         let owner = self
             .default_address(arg_matches, wallet_manager)
             .unwrap_or_else(|e| {
                 eprintln!("error: {}", e);
                 exit(1);
             });
-        get_associated_token_address_with_program_id(&owner, &token, &self.program_id)
+        get_associated_token_address_with_program_id(&owner, token, program_id)
     }
 
     // Checks if an explicit address was provided, otherwise return the default address.
@@ -184,4 +208,75 @@ impl<'a> Config<'a> {
             }
         }
     }
+
+    pub(crate) async fn get_mint_info(
+        &self,
+        mint: &Pubkey,
+        mint_decimals: Option<u8>,
+    ) -> Result<MintInfo, Error> {
+        if self.sign_only {
+            Ok(MintInfo {
+                program_id: self.program_id,
+                address: *mint,
+                decimals: mint_decimals.unwrap_or_default(),
+            })
+        } else {
+            let account = self.rpc_client.get_account(mint).await?;
+            self.check_owner(mint, &account.owner)?;
+            let mint_account = Mint::unpack(&account.data)
+                .map_err(|_| format!("Could not find mint account {}", mint))?;
+            if let Some(decimals) = mint_decimals {
+                if decimals != mint_account.decimals {
+                    return Err(format!(
+                        "Mint {:?} has decimals {}, not configured decimals {}",
+                        mint, mint_account.decimals, decimals
+                    )
+                    .into());
+                }
+            }
+            Ok(MintInfo {
+                program_id: account.owner,
+                address: *mint,
+                decimals: mint_account.decimals,
+            })
+        }
+    }
+
+    pub(crate) fn check_owner(&self, account: &Pubkey, owner: &Pubkey) -> Result<(), Error> {
+        if self.program_id != *owner {
+            Err(format!(
+                "Account {:?} is owned by {}, not configured program id {}",
+                account, owner, self.program_id
+            )
+            .into())
+        } else {
+            Ok(())
+        }
+    }
+
+    pub(crate) async fn check_account(
+        &self,
+        token_account: &Pubkey,
+        mint_address: Option<Pubkey>,
+    ) -> Result<Pubkey, Error> {
+        if !self.sign_only {
+            let account = self.rpc_client.get_account(token_account).await?;
+            let source_account = Account::unpack(&account.data)
+                .map_err(|_| format!("Could not find token account {}", token_account))?;
+            let source_mint = source_account.mint;
+            if let Some(mint) = mint_address {
+                if source_mint != mint {
+                    return Err(format!(
+                        "Source {:?} does not contain {:?} tokens",
+                        token_account, mint
+                    )
+                    .into());
+                }
+            }
+            self.check_owner(token_account, &account.owner)?;
+            Ok(source_mint)
+        } else {
+            Ok(mint_address.unwrap_or_default())
+        }
+    }
 }