gaspersic/solana
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Distribute spl tokens (#13559)

Browse Source 
* Add helpers to covert between sdk types

* Add distribute-spl-tokens to args and arg-parsing

* Build spl-token transfer-checked instructions

* Check spl-token balances properly

* Add display handling to support spl-token

* Small refactor to allow failures in allocation iter

* Use Associated Token Account for spl-token distributions

* Add spl token support to balances command

* Update readme

* Add spl-token tests

* Rename spl-tokens file

* Move a couple more things out of commands

* Stop requiring lockup_date heading for non-stake distributions

* Use epsilon for allocation retention
This commit is contained in:
Tyera Eulberg
2020-11-19 10:32:31 -07:00
committed by GitHub
parent 1ffab5de77
commit 2ef4369237
12 changed files with 793 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files

166
tokens/src/arg_parser.rs
View File

@ -1,11 +1,11 @@
use crate::args::{
Args, BalancesArgs, Command, DistributeTokensArgs, StakeArgs, TransactionLogArgs,
Args, BalancesArgs, Command, DistributeTokensArgs, SplTokenArgs, StakeArgs, TransactionLogArgs,
};
use clap::{
crate_description, crate_name, value_t, value_t_or_exit, App, Arg, ArgMatches, SubCommand,
};
use solana_clap_utils::{
input_parsers::value_of,
input_parsers::{pubkey_of_signer, value_of},
input_validators::{is_amount, is_valid_pubkey, is_valid_signer},
keypair::{pubkey_from_path, signer_from_path},
};
@ -42,7 +42,7 @@ where
)
.subcommand(
SubCommand::with_name("distribute-tokens")
.about("Distribute tokens")
.about("Distribute SOL")
.arg(
Arg::with_name("db_path")
.long("db-path")
@ -201,6 +201,78 @@ where
.help("Fee payer"),
),
)
.subcommand(
SubCommand::with_name("distribute-spl-tokens")
.about("Distribute SPL tokens")
.arg(
Arg::with_name("db_path")
.long("db-path")
.required(true)
.takes_value(true)
.value_name("FILE")
.help(
"Location for storing distribution database. \
The database is used for tracking transactions as they are finalized \
and preventing double spends.",
),
)
.arg(
Arg::with_name("input_csv")
.long("input-csv")
.required(true)
.takes_value(true)
.value_name("FILE")
.help("Allocations CSV file"),
)
.arg(
Arg::with_name("dry_run")
.long("dry-run")
.help("Do not execute any transfers"),
)
.arg(
Arg::with_name("transfer_amount")
.long("transfer-amount")
.takes_value(true)
.value_name("AMOUNT")
.validator(is_amount)
.help("The amount of SPL tokens to send to each recipient"),
)
.arg(
Arg::with_name("output_path")
.long("output-path")
.short("o")
.value_name("FILE")
.takes_value(true)
.help("Write the transaction log to this file"),
)
.arg(
Arg::with_name("token_account_address")
.long("from")
.required(true)
.takes_value(true)
.value_name("TOKEN_ACCOUNT_ADDRESS")
.validator(is_valid_pubkey)
.help("SPL token account to send from"),
)
.arg(
Arg::with_name("token_owner")
.long("owner")
.required(true)
.takes_value(true)
.value_name("TOKEN_ACCOUNT_OWNER_KEYPAIR")
.validator(is_valid_signer)
.help("SPL token account owner"),
)
.arg(
Arg::with_name("fee_payer")
.long("fee-payer")
.required(true)
.takes_value(true)
.value_name("KEYPAIR")
.validator(is_valid_signer)
.help("Fee payer"),
),
)
.subcommand(
SubCommand::with_name("balances")
.about("Balance of each account")
@ -213,6 +285,27 @@ where
.help("Allocations CSV file"),
),
)
.subcommand(
SubCommand::with_name("spl-token-balances")
.about("Balance of SPL token associated accounts")
.arg(
Arg::with_name("input_csv")
.long("input-csv")
.required(true)
.takes_value(true)
.value_name("FILE")
.help("Allocations CSV file"),
)
.arg(
Arg::with_name("mint_address")
.long("mint")
.required(true)
.takes_value(true)
.value_name("MINT_ADDRESS")
.validator(is_valid_pubkey)
.help("SPL token mint of distribution"),
),
)
.subcommand(
SubCommand::with_name("transaction-log")
.about("Print the database to a CSV file")
@ -266,6 +359,7 @@ fn parse_distribute_tokens_args(
sender_keypair,
fee_payer,
stake_args: None,
spl_token_args: None,
transfer_amount: value_of(matches, "transfer_amount"),
})
}
@ -342,14 +436,68 @@ fn parse_distribute_stake_args(
sender_keypair,
fee_payer,
stake_args: Some(stake_args),
spl_token_args: None,
transfer_amount: None,
})
}
fn parse_balances_args(matches: &ArgMatches<'_>) -> BalancesArgs {
BalancesArgs {
fn parse_distribute_spl_tokens_args(
matches: &ArgMatches<'_>,
) -> Result<DistributeTokensArgs, Box<dyn Error>> {
let mut wallet_manager = maybe_wallet_manager()?;
let signer_matches = ArgMatches::default(); // No default signer
let token_owner_str = value_t_or_exit!(matches, "token_owner", String);
let token_owner = signer_from_path(
&signer_matches,
&token_owner_str,
"owner",
&mut wallet_manager,
)?;
let fee_payer_str = value_t_or_exit!(matches, "fee_payer", String);
let fee_payer = signer_from_path(
&signer_matches,
&fee_payer_str,
"fee-payer",
&mut wallet_manager,
)?;
let token_account_address_str = value_t_or_exit!(matches, "token_account_address", String);
let token_account_address = pubkey_from_path(
&signer_matches,
&token_account_address_str,
"token account address",
&mut wallet_manager,
)?;
Ok(DistributeTokensArgs {
input_csv: value_t_or_exit!(matches, "input_csv", String),
}
transaction_db: value_t_or_exit!(matches, "db_path", String),
output_path: matches.value_of("output_path").map(|path| path.to_string()),
dry_run: matches.is_present("dry_run"),
sender_keypair: token_owner,
fee_payer,
stake_args: None,
spl_token_args: Some(SplTokenArgs {
token_account_address,
..SplTokenArgs::default()
}),
transfer_amount: value_of(matches, "transfer_amount"),
})
}
fn parse_balances_args(matches: &ArgMatches<'_>) -> Result<BalancesArgs, Box<dyn Error>> {
let mut wallet_manager = maybe_wallet_manager()?;
let spl_token_args =
pubkey_of_signer(matches, "mint_address", &mut wallet_manager)?.map(|mint| SplTokenArgs {
mint,
..SplTokenArgs::default()
});
Ok(BalancesArgs {
input_csv: value_t_or_exit!(matches, "input_csv", String),
spl_token_args,
})
}
fn parse_transaction_log_args(matches: &ArgMatches<'_>) -> TransactionLogArgs {
@ -375,7 +523,11 @@ where
("distribute-stake", Some(matches)) => {
Command::DistributeTokens(parse_distribute_stake_args(matches)?)
}
("balances", Some(matches)) => Command::Balances(parse_balances_args(matches)),
("distribute-spl-tokens", Some(matches)) => {
Command::DistributeTokens(parse_distribute_spl_tokens_args(matches)?)
}
("balances", Some(matches)) => Command::Balances(parse_balances_args(matches)?),
("spl-token-balances", Some(matches)) => Command::Balances(parse_balances_args(matches)?),
("transaction-log", Some(matches)) => {
Command::TransactionLog(parse_transaction_log_args(matches))
}

9
tokens/src/args.rs
View File

@ -8,6 +8,7 @@ pub struct DistributeTokensArgs {
pub sender_keypair: Box<dyn Signer>,
pub fee_payer: Box<dyn Signer>,
pub stake_args: Option<StakeArgs>,
pub spl_token_args: Option<SplTokenArgs>,
pub transfer_amount: Option<f64>,
}
@ -19,8 +20,16 @@ pub struct StakeArgs {
pub lockup_authority: Option<Box<dyn Signer>>,
}
#[derive(Default)]
pub struct SplTokenArgs {
pub token_account_address: Pubkey,
pub mint: Pubkey,
pub decimals: u8,
}
pub struct BalancesArgs {
pub input_csv: String,
pub spl_token_args: Option<SplTokenArgs>,
}
pub struct TransactionLogArgs {

266
tokens/src/commands.rs
View File

@ -1,5 +1,9 @@
use crate::args::{BalancesArgs, DistributeTokensArgs, StakeArgs, TransactionLogArgs};
use crate::db::{self, TransactionInfo};
use crate::{
args::{BalancesArgs, DistributeTokensArgs, StakeArgs, TransactionLogArgs},
db::{self, TransactionInfo},
spl_token::*,
token_display::Token,
};
use chrono::prelude::*;
use console::style;
use csv::{ReaderBuilder, Trim};
@ -7,6 +11,7 @@ use indexmap::IndexMap;
use indicatif::{ProgressBar, ProgressStyle};
use pickledb::PickleDb;
use serde::{Deserialize, Serialize};
use solana_account_decoder::parse_token::{pubkey_from_spl_token_v2_0, spl_token_v2_0_pubkey};
use solana_client::{
client_error::{ClientError, Result as ClientResult},
rpc_client::RpcClient,
@ -25,6 +30,8 @@ use solana_stake_program::{
stake_instruction::{self, LockupArgs},
stake_state::{Authorized, Lockup, StakeAuthorize},
};
use spl_associated_token_account_v1_0::get_associated_token_address;
use spl_token_v2_0::solana_program::program_error::ProgramError;
use std::{
cmp::{self},
io,
@ -33,15 +40,16 @@ use std::{
};
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
struct Allocation {
recipient: String,
amount: f64,
lockup_date: String,
pub struct Allocation {
pub recipient: String,
pub amount: f64,
pub lockup_date: String,
}
#[derive(Debug, PartialEq)]
pub enum FundingSource {
FeePayer,
SplTokenAccount,
StakeAccount,
SystemAccount,
}
@ -86,6 +94,8 @@ pub enum Error {
MissingLockupAuthority,
#[error("insufficient funds in {0:?}, requires {1} SOL")]
InsufficientFunds(FundingSources, f64),
#[error("Program error")]
ProgramError(#[from] ProgramError),
}
fn merge_allocations(allocations: &[Allocation]) -> Vec<Allocation> {
@ -128,7 +138,7 @@ fn apply_previous_transactions(
}
}
}
allocations.retain(|x| x.amount > 0.5);
allocations.retain(|x| x.amount > f64::EPSILON);
}
fn transfer<S: Signer>(
@ -153,8 +163,9 @@ fn distribution_instructions(
new_stake_account_address: &Pubkey,
args: &DistributeTokensArgs,
lockup_date: Option<DateTime<Utc>>,
do_create_associated_token_account: bool,
) -> Vec<Instruction> {
if args.stake_args.is_none() {
if args.stake_args.is_none() && args.spl_token_args.is_none() {
let from = args.sender_keypair.pubkey();
let to = allocation.recipient.parse().unwrap();
let lamports = sol_to_lamports(allocation.amount);
@ -162,6 +173,10 @@ fn distribution_instructions(
return vec![instruction];
}
if args.spl_token_args.is_some() {
return build_spl_token_instructions(allocation, args, do_create_associated_token_account);
}
let stake_args = args.stake_args.as_ref().unwrap();
let unlocked_sol = stake_args.unlocked_sol;
let sender_pubkey = args.sender_keypair.pubkey();
@ -228,34 +243,65 @@ fn distribute_allocations(
args: &DistributeTokensArgs,
) -> Result<(), Error> {
type StakeExtras = Vec<(Keypair, Option<DateTime<Utc>>)>;
let (messages, stake_extras): (Vec<Message>, StakeExtras) = allocations
.iter()
.map(|allocation| {
let new_stake_account_keypair = Keypair::new();
let lockup_date = if allocation.lockup_date == "" {
None
} else {
Some(allocation.lockup_date.parse::<DateTime<Utc>>().unwrap())
};
let mut messages: Vec<Message> = vec![];
let mut stake_extras: StakeExtras = vec![];
let mut created_accounts = 0;
for allocation in allocations.iter() {
let new_stake_account_keypair = Keypair::new();
let lockup_date = if allocation.lockup_date == "" {
None
} else {
Some(allocation.lockup_date.parse::<DateTime<Utc>>().unwrap())
};
println!("{:<44} {:>24.9}", allocation.recipient, allocation.amount);
let instructions = distribution_instructions(
allocation,
&new_stake_account_keypair.pubkey(),
args,
lockup_date,
let (decimals, do_create_associated_token_account) = if let Some(spl_token_args) =
&args.spl_token_args
{
let wallet_address = allocation.recipient.parse().unwrap();
let associated_token_address = get_associated_token_address(
&wallet_address,
&spl_token_v2_0_pubkey(&spl_token_args.mint),
);
let fee_payer_pubkey = args.fee_payer.pubkey();
let message = Message::new(&instructions, Some(&fee_payer_pubkey));
(message, (new_stake_account_keypair, lockup_date))
})
.unzip();
let do_create_associated_token_account = client
.get_multiple_accounts(&[pubkey_from_spl_token_v2_0(&associated_token_address)])?
[0]
.is_none();
if do_create_associated_token_account {
created_accounts += 1;
}
(
spl_token_args.decimals as usize,
do_create_associated_token_account,
)
} else {
(9, false)
};
println!(
"{:<44} {:>24.2$}",
allocation.recipient, allocation.amount, decimals
);
let instructions = distribution_instructions(
allocation,
&new_stake_account_keypair.pubkey(),
args,
lockup_date,
do_create_associated_token_account,
);
let fee_payer_pubkey = args.fee_payer.pubkey();
let message = Message::new(&instructions, Some(&fee_payer_pubkey));
messages.push(message);
stake_extras.push((new_stake_account_keypair, lockup_date));
}
let num_signatures = messages
.iter()
.map(|message| message.header.num_required_signatures as usize)
.sum();
check_payer_balances(num_signatures, allocations, client, args)?;
if args.spl_token_args.is_some() {
check_spl_token_balances(num_signatures, allocations, client, args, created_accounts)?;
} else {
check_payer_balances(num_signatures, allocations, client, args)?;
}
for ((allocation, message), (new_stake_account_keypair, lockup_date)) in
allocations.iter().zip(messages).zip(stake_extras)
@ -313,7 +359,11 @@ fn distribute_allocations(
Ok(())
}
fn read_allocations(input_csv: &str, transfer_amount: Option<f64>) -> io::Result<Vec<Allocation>> {
fn read_allocations(
input_csv: &str,
transfer_amount: Option<f64>,
require_lockup_heading: bool,
) -> io::Result<Vec<Allocation>> {
let mut rdr = ReaderBuilder::new().trim(Trim::All).from_path(input_csv)?;
let allocations = if let Some(amount) = transfer_amount {
let recipients: Vec<String> = rdr
@ -328,8 +378,21 @@ fn read_allocations(input_csv: &str, transfer_amount: Option<f64>) -> io::Result
lockup_date: "".to_string(),
})
.collect()
} else {
} else if require_lockup_heading {
rdr.deserialize().map(|entry| entry.unwrap()).collect()
} else {
let recipients: Vec<(String, f64)> = rdr
.deserialize()
.map(|recipient| recipient.unwrap())
.collect();
recipients
.into_iter()
.map(|(recipient, amount)| Allocation {
recipient,
amount,
lockup_date: "".to_string(),
})
.collect()
};
Ok(allocations)
}
@ -346,11 +409,17 @@ pub fn process_allocations(
client: &RpcClient,
args: &DistributeTokensArgs,
) -> Result<Option<usize>, Error> {
let mut allocations: Vec<Allocation> = read_allocations(&args.input_csv, args.transfer_amount)?;
let require_lockup_heading = args.stake_args.is_some();
let mut allocations: Vec<Allocation> = read_allocations(
&args.input_csv,
args.transfer_amount,
require_lockup_heading,
)?;
let is_sol = args.spl_token_args.is_none();
let starting_total_tokens: f64 = allocations.iter().map(|x| x.amount).sum();
let starting_total_tokens = Token::from(allocations.iter().map(|x| x.amount).sum(), is_sol);
println!(
"{} ◎{}",
"{} {}",
style("Total in input_csv:").bold(),
starting_total_tokens,
);
@ -368,27 +437,23 @@ pub fn process_allocations(
return Ok(confirmations);
}
let distributed_tokens: f64 = transaction_infos.iter().map(|x| x.amount).sum();
let undistributed_tokens: f64 = allocations.iter().map(|x| x.amount).sum();
println!("{} ◎{}", style("Distributed:").bold(), distributed_tokens,);
let distributed_tokens = Token::from(transaction_infos.iter().map(|x| x.amount).sum(), is_sol);
let undistributed_tokens = Token::from(allocations.iter().map(|x| x.amount).sum(), is_sol);
println!("{} {}", style("Distributed:").bold(), distributed_tokens,);
println!(
"{} ◎{}",
"{} {}",
style("Undistributed:").bold(),
undistributed_tokens,
);
println!(
"{} ◎{}",
"{} {}",
style("Total:").bold(),
distributed_tokens + undistributed_tokens,
);
println!(
"{}",
style(format!(
"{:<44} {:>24}",
"Recipient", "Expected Balance (◎)"
))
.bold()
style(format!("{:<44} {:>24}", "Recipient", "Expected Balance",)).bold()
);
distribute_allocations(client, &mut db, &allocations, args)?;
@ -572,30 +637,41 @@ fn check_payer_balances(
Ok(())
}
pub fn process_balances(client: &RpcClient, args: &BalancesArgs) -> Result<(), csv::Error> {
let allocations: Vec<Allocation> = read_allocations(&args.input_csv, None)?;
pub fn process_balances(client: &RpcClient, args: &BalancesArgs) -> Result<(), Error> {
let allocations: Vec<Allocation> = read_allocations(&args.input_csv, None, false)?;
let allocations = merge_allocations(&allocations);
let token = if let Some(spl_token_args) = &args.spl_token_args {
spl_token_args.mint.to_string()
} else {
"◎".to_string()
};
println!("{} {}", style("Token:").bold(), token);
println!(
"{}",
style(format!(
"{:<44} {:>24} {:>24} {:>24}",
"Recipient", "Expected Balance (◎)", "Actual Balance (◎)", "Difference (◎)"
"Recipient", "Expected Balance", "Actual Balance", "Difference"
))
.bold()
);
for allocation in &allocations {
let address = allocation.recipient.parse().unwrap();
let expected = lamports_to_sol(sol_to_lamports(allocation.amount));
let actual = lamports_to_sol(client.get_balance(&address).unwrap());
println!(
"{:<44} {:>24.9} {:>24.9} {:>24.9}",
allocation.recipient,
expected,
actual,
actual - expected
);
if let Some(spl_token_args) = &args.spl_token_args {
print_token_balances(client, allocation, spl_token_args)?;
} else {
let address: Pubkey = allocation.recipient.parse().unwrap();
let expected = lamports_to_sol(sol_to_lamports(allocation.amount));
let actual = lamports_to_sol(client.get_balance(&address).unwrap());
println!(
"{:<44} {:>24.9} {:>24.9} {:>24.9}",
allocation.recipient,
expected,
actual,
actual - expected,
);
}
}
Ok(())
@ -666,6 +742,7 @@ pub fn test_process_distribute_tokens_with_client(
transaction_db: transaction_db.clone(),
output_path: Some(output_path.clone()),
stake_args: None,
spl_token_args: None,
transfer_amount,
};
let confirmations = process_allocations(client, &args).unwrap();
@ -777,6 +854,7 @@ pub fn test_process_distribute_stake_with_client(client: &RpcClient, sender_keyp
transaction_db: transaction_db.clone(),
output_path: Some(output_path.clone()),
stake_args: Some(stake_args),
spl_token_args: None,
sender_keypair: Box::new(sender_keypair),
transfer_amount: None,
};
@ -910,11 +988,78 @@ mod tests {
wtr.flush().unwrap();
assert_eq!(
read_allocations(&input_csv, None).unwrap(),
read_allocations(&input_csv, None, false).unwrap(),
vec![allocation.clone()]
);
assert_eq!(
read_allocations(&input_csv, None, true).unwrap(),
vec![allocation]
);
}
#[test]
fn test_read_allocations_no_lockup() {
let pubkey0 = solana_sdk::pubkey::new_rand();
let pubkey1 = solana_sdk::pubkey::new_rand();
let file = NamedTempFile::new().unwrap();
let input_csv = file.path().to_str().unwrap().to_string();
let mut wtr = csv::WriterBuilder::new().from_writer(file);
wtr.serialize(("recipient".to_string(), "amount".to_string()))
.unwrap();
wtr.serialize((&pubkey0.to_string(), 42.0)).unwrap();
wtr.serialize((&pubkey1.to_string(), 43.0)).unwrap();
wtr.flush().unwrap();
let expected_allocations = vec![
Allocation {
recipient: pubkey0.to_string(),
amount: 42.0,
lockup_date: "".to_string(),
},
Allocation {
recipient: pubkey1.to_string(),
amount: 43.0,
lockup_date: "".to_string(),
},
];
assert_eq!(
read_allocations(&input_csv, None, false).unwrap(),
expected_allocations
);
}
#[test]
#[should_panic]
fn test_read_allocations_malformed() {
let pubkey0 = solana_sdk::pubkey::new_rand();
let pubkey1 = solana_sdk::pubkey::new_rand();
let file = NamedTempFile::new().unwrap();
let input_csv = file.path().to_str().unwrap().to_string();
let mut wtr = csv::WriterBuilder::new().from_writer(file);
wtr.serialize(("recipient".to_string(), "amount".to_string()))
.unwrap();
wtr.serialize((&pubkey0.to_string(), 42.0)).unwrap();
wtr.serialize((&pubkey1.to_string(), 43.0)).unwrap();
wtr.flush().unwrap();
let expected_allocations = vec![
Allocation {
recipient: pubkey0.to_string(),
amount: 42.0,
lockup_date: "".to_string(),
},
Allocation {
recipient: pubkey1.to_string(),
amount: 43.0,
lockup_date: "".to_string(),
},
];
assert_eq!(
read_allocations(&input_csv, None, true).unwrap(),
expected_allocations
);
}
#[test]
fn test_read_allocations_transfer_amount() {
let pubkey0 = solana_sdk::pubkey::new_rand();
@ -949,7 +1094,7 @@ mod tests {
},
];
assert_eq!(
read_allocations(&input_csv, Some(amount)).unwrap(),
read_allocations(&input_csv, Some(amount), false).unwrap(),
expected_allocations
);
}
@ -1059,6 +1204,7 @@ mod tests {
transaction_db: "".to_string(),
output_path: None,
stake_args: Some(stake_args),
spl_token_args: None,
sender_keypair: Box::new(Keypair::new()),
transfer_amount: None,
};
@ -1068,6 +1214,7 @@ mod tests {
&new_stake_account_address,
&args,
Some(lockup_date),
false,
);
let lockup_instruction =
bincode::deserialize(&instructions[SET_LOCKUP_INDEX].data).unwrap();
@ -1107,6 +1254,7 @@ mod tests {
transaction_db: "".to_string(),
output_path: None,
stake_args,
spl_token_args: None,
transfer_amount: None,
};
(allocations, args)

2
tokens/src/lib.rs
View File

@ -2,3 +2,5 @@ pub mod arg_parser;
pub mod args;
pub mod commands;
mod db;
pub mod spl_token;
pub mod token_display;

8
tokens/src/main.rs
View File

@ -1,6 +1,6 @@
use solana_cli_config::{Config, CONFIG_FILE};
use solana_client::rpc_client::RpcClient;
use solana_tokens::{arg_parser::parse_args, args::Command, commands};
use solana_tokens::{arg_parser::parse_args, args::Command, commands, spl_token};
use std::{env, error::Error, path::Path, process};
fn main() -> Result<(), Box<dyn Error>> {
@ -19,10 +19,12 @@ fn main() -> Result<(), Box<dyn Error>> {
let client = RpcClient::new(json_rpc_url);
match command_args.command {
Command::DistributeTokens(args) => {
Command::DistributeTokens(mut args) => {
spl_token::update_token_args(&client, &mut args.spl_token_args)?;
commands::process_allocations(&client, &args)?;
}
Command::Balances(args) => {
Command::Balances(mut args) => {
spl_token::update_decimals(&client, &mut args.spl_token_args)?;
commands::process_balances(&client, &args)?;
}
Command::TransactionLog(args) => {

184
tokens/src/spl_token.rs Normal file
View File

@ -0,0 +1,184 @@
use crate::{
args::{DistributeTokensArgs, SplTokenArgs},
commands::{Allocation, Error, FundingSource},
};
use console::style;
use solana_account_decoder::parse_token::{
pubkey_from_spl_token_v2_0, spl_token_v2_0_pubkey, token_amount_to_ui_amount,
};
use solana_client::rpc_client::RpcClient;
use solana_sdk::{instruction::Instruction, native_token::lamports_to_sol};
use solana_transaction_status::parse_token::spl_token_v2_0_instruction;
use spl_associated_token_account_v1_0::{
create_associated_token_account, get_associated_token_address,
};
use spl_token_v2_0::{
solana_program::program_pack::Pack,
state::{Account as SplTokenAccount, Mint},
};
pub fn update_token_args(client: &RpcClient, args: &mut Option<SplTokenArgs>) -> Result<(), Error> {
if let Some(spl_token_args) = args {
let sender_account = client
.get_account(&spl_token_args.token_account_address)
.unwrap_or_default();
let mint_address =
pubkey_from_spl_token_v2_0(&SplTokenAccount::unpack(&sender_account.data)?.mint);
spl_token_args.mint = mint_address;
update_decimals(client, args)?;
}
Ok(())
}
pub fn update_decimals(client: &RpcClient, args: &mut Option<SplTokenArgs>) -> Result<(), Error> {
if let Some(spl_token_args) = args {
let mint_account = client.get_account(&spl_token_args.mint).unwrap_or_default();
let mint = Mint::unpack(&mint_account.data)?;
spl_token_args.decimals = mint.decimals;
}
Ok(())
}
pub fn spl_token_amount(amount: f64, decimals: u8) -> u64 {
(amount * 10_usize.pow(decimals as u32) as f64) as u64
}
pub fn build_spl_token_instructions(
allocation: &Allocation,
args: &DistributeTokensArgs,
do_create_associated_token_account: bool,
) -> Vec<Instruction> {
let spl_token_args = args
.spl_token_args
.as_ref()
.expect("spl_token_args must be some");
let wallet_address = allocation.recipient.parse().unwrap();
let associated_token_address = get_associated_token_address(
&wallet_address,
&spl_token_v2_0_pubkey(&spl_token_args.mint),
);
let mut instructions = vec![];
if do_create_associated_token_account {
let create_associated_token_account_instruction = create_associated_token_account(
&spl_token_v2_0_pubkey(&args.fee_payer.pubkey()),
&wallet_address,
&spl_token_v2_0_pubkey(&spl_token_args.mint),
);
instructions.push(spl_token_v2_0_instruction(
create_associated_token_account_instruction,
));
}
let spl_instruction = spl_token_v2_0::instruction::transfer_checked(
&spl_token_v2_0::id(),
&spl_token_v2_0_pubkey(&spl_token_args.token_account_address),
&spl_token_v2_0_pubkey(&spl_token_args.mint),
&associated_token_address,
&spl_token_v2_0_pubkey(&args.sender_keypair.pubkey()),
&[],
spl_token_amount(allocation.amount, spl_token_args.decimals),
spl_token_args.decimals,
)
.unwrap();
instructions.push(spl_token_v2_0_instruction(spl_instruction));
instructions
}
pub fn check_spl_token_balances(
num_signatures: usize,
allocations: &[Allocation],
client: &RpcClient,
args: &DistributeTokensArgs,
created_accounts: u64,
) -> Result<(), Error> {
let spl_token_args = args
.spl_token_args
.as_ref()
.expect("spl_token_args must be some");
let undistributed_tokens: f64 = allocations.iter().map(|x| x.amount).sum();
let allocation_amount = spl_token_amount(undistributed_tokens, spl_token_args.decimals);
let fee_calculator = client.get_recent_blockhash()?.1;
let fees = fee_calculator
.lamports_per_signature
.checked_mul(num_signatures as u64)
.unwrap();
let token_account_rent_exempt_balance =
client.get_minimum_balance_for_rent_exemption(SplTokenAccount::LEN)?;
let account_creation_amount = created_accounts * token_account_rent_exempt_balance;
let fee_payer_balance = client.get_balance(&args.fee_payer.pubkey())?;
if fee_payer_balance < fees + account_creation_amount {
return Err(Error::InsufficientFunds(
vec![FundingSource::FeePayer].into(),
lamports_to_sol(fees + account_creation_amount),
));
}
let source_token_account = client
.get_account(&spl_token_args.token_account_address)
.unwrap_or_default();
let source_token = SplTokenAccount::unpack(&source_token_account.data)?;
if source_token.amount < allocation_amount {
return Err(Error::InsufficientFunds(
vec![FundingSource::SplTokenAccount].into(),
token_amount_to_ui_amount(allocation_amount, spl_token_args.decimals).ui_amount,
));
}
Ok(())
}
pub fn print_token_balances(
client: &RpcClient,
allocation: &Allocation,
spl_token_args: &SplTokenArgs,
) -> Result<(), Error> {
let address = allocation.recipient.parse().unwrap();
let expected = allocation.amount;
let associated_token_address = get_associated_token_address(
&spl_token_v2_0_pubkey(&address),
&spl_token_v2_0_pubkey(&spl_token_args.mint),
);
let recipient_account = client
.get_account(&pubkey_from_spl_token_v2_0(&associated_token_address))
.unwrap_or_default();
let (actual, difference) =
if let Ok(recipient_token) = SplTokenAccount::unpack(&recipient_account.data) {
let actual = token_amount_to_ui_amount(recipient_token.amount, spl_token_args.decimals)
.ui_amount;
(
style(format!(
"{:>24.1$}",
actual, spl_token_args.decimals as usize
)),
format!(
"{:>24.1$}",
actual - expected,
spl_token_args.decimals as usize
),
)
} else {
(
style("Associated token account not yet created".to_string()).yellow(),
"".to_string(),
)
};
println!(
"{:<44} {:>24.4$} {:>24} {:>24}",
allocation.recipient, expected, actual, difference, spl_token_args.decimals as usize
);
Ok(())
}
#[cfg(test)]
mod tests {
// The following unit tests were written for v1.4 using the ProgramTest framework, passing its
// BanksClient into the `solana-tokens` methods. With the revert to RpcClient in this module
// (https://github.com/solana-labs/solana/pull/13623), that approach was no longer viable.
// These tests were removed rather than rewritten to avoid accruing technical debt. Once a new
// rpc/client framework is implemented, they should be restored.
//
// async fn test_process_spl_token_allocations()
// async fn test_process_spl_token_transfer_amount_allocations()
// async fn test_check_spl_token_balances()
//
// TODO: link to v1.4 tests
}

62
tokens/src/token_display.rs Normal file
View File

@ -0,0 +1,62 @@
use std::{
fmt::{Debug, Display, Formatter, Result},
ops::Add,
};
const SOL_SYMBOL: &str = "◎";
#[derive(PartialEq)]
pub enum TokenType {
Sol,
SplToken,
}
pub struct Token {
amount: f64,
token_type: TokenType,
}
impl Token {
fn write_with_symbol(&self, f: &mut Formatter) -> Result {
match &self.token_type {
TokenType::Sol => write!(f, "{}{}", SOL_SYMBOL, self.amount,),
TokenType::SplToken => write!(f, "{} tokens", self.amount,),
}
}
pub fn from(amount: f64, is_sol: bool) -> Self {
let token_type = if is_sol {
TokenType::Sol
} else {
TokenType::SplToken
};
Self { amount, token_type }
}
}
impl Display for Token {
fn fmt(&self, f: &mut Formatter) -> Result {
self.write_with_symbol(f)
}
}
impl Debug for Token {
fn fmt(&self, f: &mut Formatter) -> Result {
self.write_with_symbol(f)
}
}
impl Add for Token {
type Output = Token;
fn add(self, other: Self) -> Self {
if self.token_type == other.token_type {
Self {
amount: self.amount + other.amount,
token_type: self.token_type,
}
} else {
self
}
}
}