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

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)