blockworks-foundation
/
solana-with-rpc-optimizations
mirror of https://github.com/blockworks-foundation/solana-with-rpc-optimizations.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

cli: solana-tokens, validate inputs gracefully (#33926) 

* cli: solana-tokens, refactor imports

* cli: solana-tokens, validate inputs gracefully

* change Allocation struct field types to simplify things

* fix typos, apply some review suggestions

* preserve backward compatibility for public APIs

* apply latest review suggestions

* address next batch of review comments

---------

Co-authored-by: norwnd <norwnd>
This commit is contained in:
norwnd 2023-11-03 03:06:00 +03:00 committed by GitHub
parent 808f67aead
commit 1c00d5d81a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 402 additions and 173 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

563
tokens/src/commands.rs
View File

@ -42,6 +42,7 @@ use {
std::{ std::{
cmp::{self}, cmp::{self},
io, io,
str::FromStr,
sync::{ sync::{
atomic::{AtomicBool, Ordering}, atomic::{AtomicBool, Ordering},
Arc, Arc,
@ -51,6 +52,7 @@ use {
}, },
}; };
/// Allocation is a helper (mostly for tests), prefer using TypedAllocation instead when possible.
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)] #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)]
pub struct Allocation { pub struct Allocation {
pub recipient: String, pub recipient: String,
@ -58,6 +60,14 @@ pub struct Allocation {
pub lockup_date: String, pub lockup_date: String,
} }
/// TypedAllocation is same as Allocation but contains typed fields.
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)]
pub struct TypedAllocation {
pub recipient: Pubkey,
pub amount: u64,
pub lockup_date: Option<DateTime<Utc>>,
}
#[derive(Debug, PartialEq, Eq)] #[derive(Debug, PartialEq, Eq)]
pub enum FundingSource { pub enum FundingSource {
FeePayer, FeePayer,
@ -98,8 +108,20 @@ type StakeExtras = Vec<(Keypair, Option<DateTime<Utc>>)>;
pub enum Error { pub enum Error {
#[error("I/O error")] #[error("I/O error")]
IoError(#[from] io::Error), IoError(#[from] io::Error),
#[error("CSV file seems to be empty")]
CsvIsEmptyError,
#[error("CSV error")] #[error("CSV error")]
CsvError(#[from] csv::Error), CsvError(#[from] csv::Error),
#[error("Bad input data for pubkey: {input}, error: {err}")]
BadInputPubkeyError {
input: String,
err: pubkey::ParsePubkeyError,
},
#[error("Bad input data for lockup date: {input}, error: {err}")]
BadInputLockupDate {
input: String,
err: chrono::ParseError,
},
#[error("PickleDb error")] #[error("PickleDb error")]
PickleDbError(#[from] pickledb::error::Error), PickleDbError(#[from] pickledb::error::Error),
#[error("Transport error")] #[error("Transport error")]
@ -118,15 +140,15 @@ pub enum Error {
ExitSignal, ExitSignal,
} }
fn merge_allocations(allocations: &[Allocation]) -> Vec<Allocation> { fn merge_allocations(allocations: &[TypedAllocation]) -> Vec<TypedAllocation> {
let mut allocation_map = IndexMap::new(); let mut allocation_map = IndexMap::new();
for allocation in allocations { for allocation in allocations {
allocation_map allocation_map
.entry(&allocation.recipient) .entry(&allocation.recipient)
.or_insert(Allocation { .or_insert(TypedAllocation {
recipient: allocation.recipient.clone(), recipient: allocation.recipient,
amount: 0, amount: 0,
lockup_date: "".to_string(), lockup_date: None,
}) })
.amount += allocation.amount; .amount += allocation.amount;
} }
@ -134,13 +156,13 @@ fn merge_allocations(allocations: &[Allocation]) -> Vec<Allocation> {
} }
/// Return true if the recipient and lockups are the same /// Return true if the recipient and lockups are the same
fn has_same_recipient(allocation: &Allocation, transaction_info: &TransactionInfo) -> bool { fn has_same_recipient(allocation: &TypedAllocation, transaction_info: &TransactionInfo) -> bool {
allocation.recipient == transaction_info.recipient.to_string() allocation.recipient == transaction_info.recipient
&& allocation.lockup_date.parse().ok() == transaction_info.lockup_date && allocation.lockup_date == transaction_info.lockup_date
} }
fn apply_previous_transactions( fn apply_previous_transactions(
allocations: &mut Vec<Allocation>, allocations: &mut Vec<TypedAllocation>,
transaction_infos: &[TransactionInfo], transaction_infos: &[TransactionInfo],
) { ) {
for transaction_info in transaction_infos { for transaction_info in transaction_infos {
@ -179,7 +201,7 @@ fn transfer<S: Signer>(
} }
fn distribution_instructions( fn distribution_instructions(
allocation: &Allocation, allocation: &TypedAllocation,
new_stake_account_address: &Pubkey, new_stake_account_address: &Pubkey,
args: &DistributeTokensArgs, args: &DistributeTokensArgs,
lockup_date: Option<DateTime<Utc>>, lockup_date: Option<DateTime<Utc>>,
@ -193,7 +215,7 @@ fn distribution_instructions(
// No stake args; a simple token transfer. // No stake args; a simple token transfer.
None => { None => {
let from = args.sender_keypair.pubkey(); let from = args.sender_keypair.pubkey();
let to = allocation.recipient.parse().unwrap(); let to = allocation.recipient;
let lamports = allocation.amount; let lamports = allocation.amount;
let instruction = system_instruction::transfer(&from, &to, lamports); let instruction = system_instruction::transfer(&from, &to, lamports);
vec![instruction] vec![instruction]
@ -203,7 +225,7 @@ fn distribution_instructions(
Some(stake_args) => { Some(stake_args) => {
let unlocked_sol = stake_args.unlocked_sol; let unlocked_sol = stake_args.unlocked_sol;
let sender_pubkey = args.sender_keypair.pubkey(); let sender_pubkey = args.sender_keypair.pubkey();
let recipient = allocation.recipient.parse().unwrap(); let recipient = allocation.recipient;
let mut instructions = match &stake_args.sender_stake_args { let mut instructions = match &stake_args.sender_stake_args {
// No source stake account, so create a recipient stake account directly. // No source stake account, so create a recipient stake account directly.
@ -304,7 +326,7 @@ fn distribution_instructions(
fn build_messages( fn build_messages(
client: &RpcClient, client: &RpcClient,
db: &mut PickleDb, db: &mut PickleDb,
allocations: &[Allocation], allocations: &[TypedAllocation],
args: &DistributeTokensArgs, args: &DistributeTokensArgs,
exit: Arc<AtomicBool>, exit: Arc<AtomicBool>,
messages: &mut Vec<Message>, messages: &mut Vec<Message>,
@ -318,7 +340,7 @@ fn build_messages(
let associated_token_addresses = allocation_chunk let associated_token_addresses = allocation_chunk
.iter() .iter()
.map(|x| { .map(|x| {
let wallet_address = x.recipient.parse().unwrap(); let wallet_address = x.recipient;
get_associated_token_address(&wallet_address, &spl_token_args.mint) get_associated_token_address(&wallet_address, &spl_token_args.mint)
}) })
.collect::<Vec<_>>(); .collect::<Vec<_>>();
@ -333,11 +355,7 @@ fn build_messages(
return Err(Error::ExitSignal); return Err(Error::ExitSignal);
} }
let new_stake_account_keypair = Keypair::new(); let new_stake_account_keypair = Keypair::new();
let lockup_date = if allocation.lockup_date.is_empty() { let lockup_date = allocation.lockup_date;
None
} else {
Some(allocation.lockup_date.parse::<DateTime<Utc>>().unwrap())
};
let do_create_associated_token_account = if let Some(spl_token_args) = &args.spl_token_args let do_create_associated_token_account = if let Some(spl_token_args) = &args.spl_token_args
{ {
@ -382,7 +400,7 @@ fn build_messages(
fn send_messages( fn send_messages(
client: &RpcClient, client: &RpcClient,
db: &mut PickleDb, db: &mut PickleDb,
allocations: &[Allocation], allocations: &[TypedAllocation],
args: &DistributeTokensArgs, args: &DistributeTokensArgs,
exit: Arc<AtomicBool>, exit: Arc<AtomicBool>,
messages: Vec<Message>, messages: Vec<Message>,
@ -404,7 +422,7 @@ fn send_messages(
signers.push(&*sender_stake_args.stake_authority); signers.push(&*sender_stake_args.stake_authority);
signers.push(&*sender_stake_args.withdraw_authority); signers.push(&*sender_stake_args.withdraw_authority);
signers.push(&new_stake_account_keypair); signers.push(&new_stake_account_keypair);
if !allocation.lockup_date.is_empty() { if allocation.lockup_date.is_some() {
if let Some(lockup_authority) = &sender_stake_args.lockup_authority { if let Some(lockup_authority) = &sender_stake_args.lockup_authority {
signers.push(&**lockup_authority); signers.push(&**lockup_authority);
} else { } else {
@ -435,7 +453,7 @@ fn send_messages(
args.stake_args.as_ref().map(|_| &new_stake_account_address); args.stake_args.as_ref().map(|_| &new_stake_account_address);
db::set_transaction_info( db::set_transaction_info(
db, db,
&allocation.recipient.parse().unwrap(), &allocation.recipient,
allocation.amount, allocation.amount,
&transaction, &transaction,
new_stake_account_address_option, new_stake_account_address_option,
@ -455,7 +473,7 @@ fn send_messages(
fn distribute_allocations( fn distribute_allocations(
client: &RpcClient, client: &RpcClient,
db: &mut PickleDb, db: &mut PickleDb,
allocations: &[Allocation], allocations: &[TypedAllocation],
args: &DistributeTokensArgs, args: &DistributeTokensArgs,
exit: Arc<AtomicBool>, exit: Arc<AtomicBool>,
) -> Result<(), Error> { ) -> Result<(), Error> {
@ -490,63 +508,91 @@ fn distribute_allocations(
fn read_allocations( fn read_allocations(
input_csv: &str, input_csv: &str,
transfer_amount: Option<u64>, transfer_amount: Option<u64>,
require_lockup_heading: bool, with_lockup: bool,
raw_amount: bool, raw_amount: bool,
) -> io::Result<Vec<Allocation>> { ) -> Result<Vec<TypedAllocation>, Error> {
let mut rdr = ReaderBuilder::new().trim(Trim::All).from_path(input_csv)?; let mut rdr = ReaderBuilder::new().trim(Trim::All).from_path(input_csv)?;
let allocations = if let Some(amount) = transfer_amount { let allocations = if let Some(amount) = transfer_amount {
let recipients: Vec<String> = rdr rdr.deserialize()
.deserialize() .map(|recipient| {
.map(|recipient| recipient.unwrap()) let recipient: String = recipient?;
.collect(); let recipient =
recipients Pubkey::from_str(&recipient).map_err(|err| Error::BadInputPubkeyError {
.into_iter() input: recipient,
.map(|recipient| Allocation { err,
recipient, })?;
amount, Ok(TypedAllocation {
lockup_date: "".to_string(), recipient,
amount,
lockup_date: None,
})
}) })
.collect() .collect::<Result<Vec<TypedAllocation>, Error>>()?
} else if require_lockup_heading { } else if with_lockup {
let recipients: Vec<(String, f64, String)> = rdr // We only support SOL token in "require lockup" mode.
.deserialize() rdr.deserialize()
.map(|recipient| recipient.unwrap()) .map(|recipient| {
.collect(); let (recipient, amount, lockup_date): (String, f64, String) = recipient?;
recipients let recipient =
.into_iter() Pubkey::from_str(&recipient).map_err(|err| Error::BadInputPubkeyError {
.map(|(recipient, amount, lockup_date)| Allocation { input: recipient,
recipient, err,
amount: sol_to_lamports(amount), })?;
lockup_date, let lockup_date = if !lockup_date.is_empty() {
let lockup_date = lockup_date.parse::<DateTime<Utc>>().map_err(|err| {
Error::BadInputLockupDate {
input: lockup_date,
err,
}
})?;
Some(lockup_date)
} else {
// empty lockup date means no lockup, it's okay to have only some lockups specified
None
};
Ok(TypedAllocation {
recipient,
amount: sol_to_lamports(amount),
lockup_date,
})
}) })
.collect() .collect::<Result<Vec<TypedAllocation>, Error>>()?
} else if raw_amount { } else if raw_amount {
let recipients: Vec<(String, u64)> = rdr rdr.deserialize()
.deserialize() .map(|recipient| {
.map(|recipient| recipient.unwrap()) let (recipient, amount): (String, u64) = recipient?;
.collect(); let recipient =
recipients Pubkey::from_str(&recipient).map_err(|err| Error::BadInputPubkeyError {
.into_iter() input: recipient,
.map(|(recipient, amount)| Allocation { err,
recipient, })?;
amount, Ok(TypedAllocation {
lockup_date: "".to_string(), recipient,
amount,
lockup_date: None,
})
}) })
.collect() .collect::<Result<Vec<TypedAllocation>, Error>>()?
} else { } else {
let recipients: Vec<(String, f64)> = rdr rdr.deserialize()
.deserialize() .map(|recipient| {
.map(|recipient| recipient.unwrap()) let (recipient, amount): (String, f64) = recipient?;
.collect(); let recipient =
recipients Pubkey::from_str(&recipient).map_err(|err| Error::BadInputPubkeyError {
.into_iter() input: recipient,
.map(|(recipient, amount)| Allocation { err,
recipient, })?;
amount: sol_to_lamports(amount), Ok(TypedAllocation {
lockup_date: "".to_string(), recipient,
amount: sol_to_lamports(amount),
lockup_date: None,
})
}) })
.collect() .collect::<Result<Vec<TypedAllocation>, Error>>()?
}; };
if allocations.is_empty() {
return Err(Error::CsvIsEmptyError);
}
Ok(allocations) Ok(allocations)
} }
@ -566,11 +612,11 @@ pub fn process_allocations(
args: &DistributeTokensArgs, args: &DistributeTokensArgs,
exit: Arc<AtomicBool>, exit: Arc<AtomicBool>,
) -> Result<Option<usize>, Error> { ) -> Result<Option<usize>, Error> {
let require_lockup_heading = args.stake_args.is_some(); let with_lockup = args.stake_args.is_some();
let mut allocations: Vec<Allocation> = read_allocations( let mut allocations: Vec<TypedAllocation> = read_allocations(
&args.input_csv, &args.input_csv,
args.transfer_amount, args.transfer_amount,
require_lockup_heading, with_lockup,
args.spl_token_args.is_some(), args.spl_token_args.is_some(),
)?; )?;
@ -773,7 +819,7 @@ pub fn get_fee_estimate_for_messages(
fn check_payer_balances( fn check_payer_balances(
messages: &[Message], messages: &[Message],
allocations: &[Allocation], allocations: &[TypedAllocation],
client: &RpcClient, client: &RpcClient,
args: &DistributeTokensArgs, args: &DistributeTokensArgs,
) -> Result<(), Error> { ) -> Result<(), Error> {
@ -857,7 +903,7 @@ pub fn process_balances(
args: &BalancesArgs, args: &BalancesArgs,
exit: Arc<AtomicBool>, exit: Arc<AtomicBool>,
) -> Result<(), Error> { ) -> Result<(), Error> {
let allocations: Vec<Allocation> = let allocations: Vec<TypedAllocation> =
read_allocations(&args.input_csv, None, false, args.spl_token_args.is_some())?; read_allocations(&args.input_csv, None, false, args.spl_token_args.is_some())?;
let allocations = merge_allocations(&allocations); let allocations = merge_allocations(&allocations);
@ -885,7 +931,7 @@ pub fn process_balances(
if let Some(spl_token_args) = &args.spl_token_args { if let Some(spl_token_args) = &args.spl_token_args {
print_token_balances(client, allocation, spl_token_args)?; print_token_balances(client, allocation, spl_token_args)?;
} else { } else {
let address: Pubkey = allocation.recipient.parse().unwrap(); let address: Pubkey = allocation.recipient;
let expected = lamports_to_sol(allocation.amount); let expected = lamports_to_sol(allocation.amount);
let actual = lamports_to_sol(client.get_balance(&address).unwrap()); let actual = lamports_to_sol(client.get_balance(&address).unwrap());
println!( println!(
@ -909,9 +955,13 @@ pub fn process_transaction_log(args: &TransactionLogArgs) -> Result<(), Error> {
use { use {
crate::db::check_output_file, crate::db::check_output_file,
solana_sdk::{pubkey::Pubkey, signature::Keypair}, solana_sdk::{
pubkey::{self, Pubkey},
signature::Keypair,
},
tempfile::{tempdir, NamedTempFile}, tempfile::{tempdir, NamedTempFile},
}; };
pub fn test_process_distribute_tokens_with_client( pub fn test_process_distribute_tokens_with_client(
client: &RpcClient, client: &RpcClient,
sender_keypair: Keypair, sender_keypair: Keypair,
@ -939,7 +989,7 @@ pub fn test_process_distribute_tokens_with_client(
} else { } else {
sol_to_lamports(1000.0) sol_to_lamports(1000.0)
}; };
let alice_pubkey = solana_sdk::pubkey::new_rand(); let alice_pubkey = pubkey::new_rand();
let allocations_file = NamedTempFile::new().unwrap(); let allocations_file = NamedTempFile::new().unwrap();
let input_csv = allocations_file.path().to_str().unwrap().to_string(); let input_csv = allocations_file.path().to_str().unwrap().to_string();
let mut wtr = csv::WriterBuilder::new().from_writer(allocations_file); let mut wtr = csv::WriterBuilder::new().from_writer(allocations_file);
@ -1039,7 +1089,7 @@ pub fn test_process_create_stake_with_client(client: &RpcClient, sender_keypair:
.unwrap(); .unwrap();
let expected_amount = sol_to_lamports(1000.0); let expected_amount = sol_to_lamports(1000.0);
let alice_pubkey = solana_sdk::pubkey::new_rand(); let alice_pubkey = pubkey::new_rand();
let file = NamedTempFile::new().unwrap(); let file = NamedTempFile::new().unwrap();
let input_csv = file.path().to_str().unwrap().to_string(); let input_csv = file.path().to_str().unwrap().to_string();
let mut wtr = csv::WriterBuilder::new().from_writer(file); let mut wtr = csv::WriterBuilder::new().from_writer(file);
@ -1161,7 +1211,7 @@ pub fn test_process_distribute_stake_with_client(client: &RpcClient, sender_keyp
.unwrap(); .unwrap();
let expected_amount = sol_to_lamports(1000.0); let expected_amount = sol_to_lamports(1000.0);
let alice_pubkey = solana_sdk::pubkey::new_rand(); let alice_pubkey = pubkey::new_rand();
let file = NamedTempFile::new().unwrap(); let file = NamedTempFile::new().unwrap();
let input_csv = file.path().to_str().unwrap().to_string(); let input_csv = file.path().to_str().unwrap().to_string();
let mut wtr = csv::WriterBuilder::new().from_writer(file); let mut wtr = csv::WriterBuilder::new().from_writer(file);
@ -1328,16 +1378,27 @@ mod tests {
#[test] #[test]
fn test_read_allocations() { fn test_read_allocations() {
let alice_pubkey = solana_sdk::pubkey::new_rand(); let alice_pubkey = pubkey::new_rand();
let allocation = Allocation { let allocation = TypedAllocation {
recipient: alice_pubkey.to_string(), recipient: alice_pubkey,
amount: 42, amount: 42,
lockup_date: "".to_string(), lockup_date: None,
}; };
let file = NamedTempFile::new().unwrap(); let file = NamedTempFile::new().unwrap();
let input_csv = file.path().to_str().unwrap().to_string(); let input_csv = file.path().to_str().unwrap().to_string();
let mut wtr = csv::WriterBuilder::new().from_writer(file); let mut wtr = csv::WriterBuilder::new().from_writer(file);
wtr.serialize(&allocation).unwrap(); wtr.serialize((
"recipient".to_string(),
"amount".to_string(),
"require_lockup".to_string(),
))
.unwrap();
wtr.serialize((
allocation.recipient.to_string(),
allocation.amount,
allocation.lockup_date,
))
.unwrap();
wtr.flush().unwrap(); wtr.flush().unwrap();
assert_eq!( assert_eq!(
@ -1345,10 +1406,10 @@ mod tests {
vec![allocation] vec![allocation]
); );
let allocation_sol = Allocation { let allocation_sol = TypedAllocation {
recipient: alice_pubkey.to_string(), recipient: alice_pubkey,
amount: sol_to_lamports(42.0), amount: sol_to_lamports(42.0),
lockup_date: "".to_string(), lockup_date: None,
}; };
assert_eq!( assert_eq!(
@ -1367,8 +1428,8 @@ mod tests {
#[test] #[test]
fn test_read_allocations_no_lockup() { fn test_read_allocations_no_lockup() {
let pubkey0 = solana_sdk::pubkey::new_rand(); let pubkey0 = pubkey::new_rand();
let pubkey1 = solana_sdk::pubkey::new_rand(); let pubkey1 = pubkey::new_rand();
let file = NamedTempFile::new().unwrap(); let file = NamedTempFile::new().unwrap();
let input_csv = file.path().to_str().unwrap().to_string(); let input_csv = file.path().to_str().unwrap().to_string();
let mut wtr = csv::WriterBuilder::new().from_writer(file); let mut wtr = csv::WriterBuilder::new().from_writer(file);
@ -1379,15 +1440,15 @@ mod tests {
wtr.flush().unwrap(); wtr.flush().unwrap();
let expected_allocations = vec![ let expected_allocations = vec![
Allocation { TypedAllocation {
recipient: pubkey0.to_string(), recipient: pubkey0,
amount: sol_to_lamports(42.0), amount: sol_to_lamports(42.0),
lockup_date: "".to_string(), lockup_date: None,
}, },
Allocation { TypedAllocation {
recipient: pubkey1.to_string(), recipient: pubkey1,
amount: sol_to_lamports(43.0), amount: sol_to_lamports(43.0),
lockup_date: "".to_string(), lockup_date: None,
}, },
]; ];
assert_eq!( assert_eq!(
@ -1397,42 +1458,210 @@ mod tests {
} }
#[test] #[test]
#[should_panic]
fn test_read_allocations_malformed() { fn test_read_allocations_malformed() {
let pubkey0 = solana_sdk::pubkey::new_rand(); let pubkey0 = pubkey::new_rand();
let pubkey1 = solana_sdk::pubkey::new_rand(); let pubkey1 = pubkey::new_rand();
// Empty file.
let file = NamedTempFile::new().unwrap(); let file = NamedTempFile::new().unwrap();
let mut wtr = csv::WriterBuilder::new().from_writer(&file);
wtr.flush().unwrap();
let input_csv = file.path().to_str().unwrap().to_string(); let input_csv = file.path().to_str().unwrap().to_string();
let mut wtr = csv::WriterBuilder::new().from_writer(file); let got = read_allocations(&input_csv, None, false, false);
assert!(matches!(got, Err(Error::CsvIsEmptyError)));
// Missing 2nd column.
let file = NamedTempFile::new().unwrap();
let mut wtr = csv::WriterBuilder::new().from_writer(&file);
wtr.serialize("recipient".to_string()).unwrap();
wtr.serialize(pubkey0.to_string()).unwrap();
wtr.serialize(pubkey1.to_string()).unwrap();
wtr.flush().unwrap();
let input_csv = file.path().to_str().unwrap().to_string();
let got = read_allocations(&input_csv, None, false, false);
assert!(matches!(got, Err(Error::CsvError(..))));
// Missing 3rd column.
let file = NamedTempFile::new().unwrap();
let mut wtr = csv::WriterBuilder::new().from_writer(&file);
wtr.serialize(("recipient".to_string(), "amount".to_string())) wtr.serialize(("recipient".to_string(), "amount".to_string()))
.unwrap(); .unwrap();
wtr.serialize((&pubkey0.to_string(), 42.0)).unwrap(); wtr.serialize((pubkey0.to_string(), "42.0".to_string()))
wtr.serialize((&pubkey1.to_string(), 43.0)).unwrap(); .unwrap();
wtr.serialize((pubkey1.to_string(), "43.0".to_string()))
.unwrap();
wtr.flush().unwrap(); wtr.flush().unwrap();
let input_csv = file.path().to_str().unwrap().to_string();
let got = read_allocations(&input_csv, None, true, false);
assert!(matches!(got, Err(Error::CsvError(..))));
let expected_allocations = vec![ let generate_csv_file = |header: (String, String, String),
Allocation { data: Vec<(String, String, String)>,
recipient: pubkey0.to_string(), file: &NamedTempFile| {
amount: sol_to_lamports(42.0), let mut wtr = csv::WriterBuilder::new().from_writer(file);
lockup_date: "".to_string(), wtr.serialize(header).unwrap();
}, wtr.serialize(&data[0]).unwrap();
Allocation { wtr.serialize(&data[1]).unwrap();
recipient: pubkey1.to_string(), wtr.flush().unwrap();
amount: sol_to_lamports(43.0), };
lockup_date: "".to_string(),
}, let default_header = (
]; "recipient".to_string(),
assert_eq!( "amount".to_string(),
read_allocations(&input_csv, None, true, false).unwrap(), "require_lockup".to_string(),
expected_allocations );
// Bad pubkey (default).
let file = NamedTempFile::new().unwrap();
generate_csv_file(
default_header.clone(),
vec![
(pubkey0.to_string(), "42.0".to_string(), "".to_string()),
("bad pubkey".to_string(), "43.0".to_string(), "".to_string()),
],
&file,
);
let input_csv = file.path().to_str().unwrap().to_string();
let got_err = read_allocations(&input_csv, None, false, false).unwrap_err();
assert!(
matches!(got_err, Error::BadInputPubkeyError { input, .. } if input == *"bad pubkey")
);
// Bad pubkey (with transfer amount).
let file = NamedTempFile::new().unwrap();
generate_csv_file(
default_header.clone(),
vec![
(pubkey0.to_string(), "42.0".to_string(), "".to_string()),
("bad pubkey".to_string(), "43.0".to_string(), "".to_string()),
],
&file,
);
let input_csv = file.path().to_str().unwrap().to_string();
let got_err = read_allocations(&input_csv, Some(123), false, false).unwrap_err();
assert!(
matches!(got_err, Error::BadInputPubkeyError { input, .. } if input == *"bad pubkey")
);
// Bad pubkey (with require lockup).
let file = NamedTempFile::new().unwrap();
generate_csv_file(
default_header.clone(),
vec![
(
pubkey0.to_string(),
"42.0".to_string(),
"2021-02-07T00:00:00Z".to_string(),
),
(
"bad pubkey".to_string(),
"43.0".to_string(),
"2021-02-07T00:00:00Z".to_string(),
),
],
&file,
);
let input_csv = file.path().to_str().unwrap().to_string();
let got_err = read_allocations(&input_csv, None, true, false).unwrap_err();
assert!(
matches!(got_err, Error::BadInputPubkeyError { input, .. } if input == *"bad pubkey")
);
// Bad pubkey (with raw amount).
let file = NamedTempFile::new().unwrap();
generate_csv_file(
default_header.clone(),
vec![
(pubkey0.to_string(), "42".to_string(), "".to_string()),
("bad pubkey".to_string(), "43".to_string(), "".to_string()),
],
&file,
);
let input_csv = file.path().to_str().unwrap().to_string();
let got_err = read_allocations(&input_csv, None, false, true).unwrap_err();
assert!(
matches!(got_err, Error::BadInputPubkeyError { input, .. } if input == *"bad pubkey")
);
// Bad value in 2nd column (default).
let file = NamedTempFile::new().unwrap();
generate_csv_file(
default_header.clone(),
vec![
(
pubkey0.to_string(),
"bad amount".to_string(),
"".to_string(),
),
(
pubkey1.to_string(),
"43.0".to_string().to_string(),
"".to_string(),
),
],
&file,
);
let input_csv = file.path().to_str().unwrap().to_string();
let got = read_allocations(&input_csv, None, false, false);
assert!(matches!(got, Err(Error::CsvError(..))));
// Bad value in 2nd column (with require lockup).
let file = NamedTempFile::new().unwrap();
generate_csv_file(
default_header.clone(),
vec![
(
pubkey0.to_string(),
"bad amount".to_string(),
"".to_string(),
),
(pubkey1.to_string(), "43.0".to_string(), "".to_string()),
],
&file,
);
let input_csv = file.path().to_str().unwrap().to_string();
let got = read_allocations(&input_csv, None, true, false);
assert!(matches!(got, Err(Error::CsvError(..))));
// Bad value in 2nd column (with raw amount).
let file = NamedTempFile::new().unwrap();
generate_csv_file(
default_header.clone(),
vec![
(pubkey0.to_string(), "42".to_string(), "".to_string()),
(pubkey1.to_string(), "43.0".to_string(), "".to_string()), // bad raw amount
],
&file,
);
let input_csv = file.path().to_str().unwrap().to_string();
let got = read_allocations(&input_csv, None, false, true);
assert!(matches!(got, Err(Error::CsvError(..))));
// Bad value in 3rd column.
let file = NamedTempFile::new().unwrap();
generate_csv_file(
default_header.clone(),
vec![
(
pubkey0.to_string(),
"42.0".to_string(),
"2021-01-07T00:00:00Z".to_string(),
),
(
pubkey1.to_string(),
"43.0".to_string(),
"bad lockup date".to_string(),
),
],
&file,
);
let input_csv = file.path().to_str().unwrap().to_string();
let got_err = read_allocations(&input_csv, None, true, false).unwrap_err();
assert!(
matches!(got_err, Error::BadInputLockupDate { input, .. } if input == *"bad lockup date")
); );
} }
#[test] #[test]
fn test_read_allocations_transfer_amount() { fn test_read_allocations_transfer_amount() {
let pubkey0 = solana_sdk::pubkey::new_rand(); let pubkey0 = pubkey::new_rand();
let pubkey1 = solana_sdk::pubkey::new_rand(); let pubkey1 = pubkey::new_rand();
let pubkey2 = solana_sdk::pubkey::new_rand(); let pubkey2 = pubkey::new_rand();
let file = NamedTempFile::new().unwrap(); let file = NamedTempFile::new().unwrap();
let input_csv = file.path().to_str().unwrap().to_string(); let input_csv = file.path().to_str().unwrap().to_string();
let mut wtr = csv::WriterBuilder::new().from_writer(file); let mut wtr = csv::WriterBuilder::new().from_writer(file);
@ -1445,20 +1674,20 @@ mod tests {
let amount = sol_to_lamports(1.5); let amount = sol_to_lamports(1.5);
let expected_allocations = vec![ let expected_allocations = vec![
Allocation { TypedAllocation {
recipient: pubkey0.to_string(), recipient: pubkey0,
amount, amount,
lockup_date: "".to_string(), lockup_date: None,
}, },
Allocation { TypedAllocation {
recipient: pubkey1.to_string(), recipient: pubkey1,
amount, amount,
lockup_date: "".to_string(), lockup_date: None,
}, },
Allocation { TypedAllocation {
recipient: pubkey2.to_string(), recipient: pubkey2,
amount, amount,
lockup_date: "".to_string(), lockup_date: None,
}, },
]; ];
assert_eq!( assert_eq!(
@ -1469,18 +1698,18 @@ mod tests {
#[test] #[test]
fn test_apply_previous_transactions() { fn test_apply_previous_transactions() {
let alice = solana_sdk::pubkey::new_rand(); let alice = pubkey::new_rand();
let bob = solana_sdk::pubkey::new_rand(); let bob = pubkey::new_rand();
let mut allocations = vec![ let mut allocations = vec![
Allocation { TypedAllocation {
recipient: alice.to_string(), recipient: alice,
amount: sol_to_lamports(1.0), amount: sol_to_lamports(1.0),
lockup_date: "".to_string(), lockup_date: None,
}, },
Allocation { TypedAllocation {
recipient: bob.to_string(), recipient: bob,
amount: sol_to_lamports(1.0), amount: sol_to_lamports(1.0),
lockup_date: "".to_string(), lockup_date: None,
}, },
]; ];
let transaction_infos = vec![TransactionInfo { let transaction_infos = vec![TransactionInfo {
@ -1493,24 +1722,24 @@ mod tests {
// Ensure that we applied the transaction to the allocation with // Ensure that we applied the transaction to the allocation with
// a matching recipient address (to bob, not alice). // a matching recipient address (to bob, not alice).
assert_eq!(allocations[0].recipient, alice.to_string()); assert_eq!(allocations[0].recipient, alice);
} }
#[test] #[test]
fn test_has_same_recipient() { fn test_has_same_recipient() {
let alice_pubkey = solana_sdk::pubkey::new_rand(); let alice_pubkey = pubkey::new_rand();
let bob_pubkey = solana_sdk::pubkey::new_rand(); let bob_pubkey = pubkey::new_rand();
let lockup0 = "2021-01-07T00:00:00Z".to_string(); let lockup0 = "2021-01-07T00:00:00Z".to_string();
let lockup1 = "9999-12-31T23:59:59Z".to_string(); let lockup1 = "9999-12-31T23:59:59Z".to_string();
let alice_alloc = Allocation { let alice_alloc = TypedAllocation {
recipient: alice_pubkey.to_string(), recipient: alice_pubkey,
amount: sol_to_lamports(1.0), amount: sol_to_lamports(1.0),
lockup_date: "".to_string(), lockup_date: None,
}; };
let alice_alloc_lockup0 = Allocation { let alice_alloc_lockup0 = TypedAllocation {
recipient: alice_pubkey.to_string(), recipient: alice_pubkey,
amount: sol_to_lamports(1.0), amount: sol_to_lamports(1.0),
lockup_date: lockup0.clone(), lockup_date: lockup0.parse().ok(),
}; };
let alice_info = TransactionInfo { let alice_info = TransactionInfo {
recipient: alice_pubkey, recipient: alice_pubkey,
@ -1550,13 +1779,13 @@ mod tests {
#[test] #[test]
fn test_set_split_stake_lockup() { fn test_set_split_stake_lockup() {
let lockup_date_str = "2021-01-07T00:00:00Z"; let lockup_date_str = "2021-01-07T00:00:00Z";
let allocation = Allocation { let allocation = TypedAllocation {
recipient: Pubkey::default().to_string(), recipient: Pubkey::default(),
amount: sol_to_lamports(1.002_282_880), amount: sol_to_lamports(1.002_282_880),
lockup_date: lockup_date_str.to_string(), lockup_date: lockup_date_str.parse().ok(),
}; };
let stake_account_address = solana_sdk::pubkey::new_rand(); let stake_account_address = pubkey::new_rand();
let new_stake_account_address = solana_sdk::pubkey::new_rand(); let new_stake_account_address = pubkey::new_rand();
let lockup_authority = Keypair::new(); let lockup_authority = Keypair::new();
let lockup_authority_address = lockup_authority.pubkey(); let lockup_authority_address = lockup_authority.pubkey();
let sender_stake_args = SenderStakeArgs { let sender_stake_args = SenderStakeArgs {
@ -1613,12 +1842,12 @@ mod tests {
sender_keypair_file: &str, sender_keypair_file: &str,
fee_payer: &str, fee_payer: &str,
stake_args: Option<StakeArgs>, stake_args: Option<StakeArgs>,
) -> (Vec<Allocation>, DistributeTokensArgs) { ) -> (Vec<TypedAllocation>, DistributeTokensArgs) {
let recipient = solana_sdk::pubkey::new_rand(); let recipient = pubkey::new_rand();
let allocations = vec![Allocation { let allocations = vec![TypedAllocation {
recipient: recipient.to_string(), recipient,
amount: allocation_amount, amount: allocation_amount,
lockup_date: "".to_string(), lockup_date: None,
}]; }];
let args = DistributeTokensArgs { let args = DistributeTokensArgs {
sender_keypair: read_keypair_file(sender_keypair_file).unwrap().into(), sender_keypair: read_keypair_file(sender_keypair_file).unwrap().into(),
@ -1890,10 +2119,10 @@ mod tests {
// Underfunded stake-account // Underfunded stake-account
let expensive_allocation_amount = 5000.0; let expensive_allocation_amount = 5000.0;
let expensive_allocations = vec![Allocation { let expensive_allocations = vec![TypedAllocation {
recipient: solana_sdk::pubkey::new_rand().to_string(), recipient: pubkey::new_rand(),
amount: sol_to_lamports(expensive_allocation_amount), amount: sol_to_lamports(expensive_allocation_amount),
lockup_date: "".to_string(), lockup_date: None,
}]; }];
let err_result = check_payer_balances( let err_result = check_payer_balances(
&[one_signer_message(&client)], &[one_signer_message(&client)],
@ -2108,10 +2337,10 @@ mod tests {
spl_token_args: None, spl_token_args: None,
transfer_amount: None, transfer_amount: None,
}; };
let allocation = Allocation { let allocation = TypedAllocation {
recipient: recipient.to_string(), recipient,
amount: sol_to_lamports(1.0), amount: sol_to_lamports(1.0),
lockup_date: "".to_string(), lockup_date: None,
}; };
let mut messages: Vec<Message> = vec![]; let mut messages: Vec<Message> = vec![];
@ -2230,10 +2459,10 @@ mod tests {
spl_token_args: None, spl_token_args: None,
transfer_amount: None, transfer_amount: None,
}; };
let allocation = Allocation { let allocation = TypedAllocation {
recipient: recipient.to_string(), recipient,
amount: sol_to_lamports(1.0), amount: sol_to_lamports(1.0),
lockup_date: "".to_string(), lockup_date: None,
}; };
let message = transaction.message.clone(); let message = transaction.message.clone();
@ -2329,10 +2558,10 @@ mod tests {
.to_string(); .to_string();
let mut db = db::open_db(&db_file, false).unwrap(); let mut db = db::open_db(&db_file, false).unwrap();
let recipient = Pubkey::new_unique(); let recipient = Pubkey::new_unique();
let allocation = Allocation { let allocation = TypedAllocation {
recipient: recipient.to_string(), recipient,
amount: sol_to_lamports(1.0), amount: sol_to_lamports(1.0),
lockup_date: "".to_string(), lockup_date: None,
}; };
// This is just dummy data; Args will not affect messages // This is just dummy data; Args will not affect messages
let args = DistributeTokensArgs { let args = DistributeTokensArgs {

12
tokens/src/spl_token.rs
View File

@ -1,7 +1,7 @@
use { use {
crate::{ crate::{
args::{DistributeTokensArgs, SplTokenArgs}, args::{DistributeTokensArgs, SplTokenArgs},
commands::{get_fee_estimate_for_messages, Allocation, Error, FundingSource}, commands::{get_fee_estimate_for_messages, Error, FundingSource, TypedAllocation},
}, },
console::style, console::style,
solana_account_decoder::parse_token::{real_number_string, real_number_string_trimmed}, solana_account_decoder::parse_token::{real_number_string, real_number_string_trimmed},
@ -37,7 +37,7 @@ pub fn update_decimals(client: &RpcClient, args: &mut Option<SplTokenArgs>) -> R
} }
pub(crate) fn build_spl_token_instructions( pub(crate) fn build_spl_token_instructions(
allocation: &Allocation, allocation: &TypedAllocation,
args: &DistributeTokensArgs, args: &DistributeTokensArgs,
do_create_associated_token_account: bool, do_create_associated_token_account: bool,
) -> Vec<Instruction> { ) -> Vec<Instruction> {
@ -45,7 +45,7 @@ pub(crate) fn build_spl_token_instructions(
.spl_token_args .spl_token_args
.as_ref() .as_ref()
.expect("spl_token_args must be some"); .expect("spl_token_args must be some");
let wallet_address = allocation.recipient.parse().unwrap(); let wallet_address = allocation.recipient;
let associated_token_address = let associated_token_address =
get_associated_token_address(&wallet_address, &spl_token_args.mint); get_associated_token_address(&wallet_address, &spl_token_args.mint);
let mut instructions = vec![]; let mut instructions = vec![];
@ -75,7 +75,7 @@ pub(crate) fn build_spl_token_instructions(
pub(crate) fn check_spl_token_balances( pub(crate) fn check_spl_token_balances(
messages: &[Message], messages: &[Message],
allocations: &[Allocation], allocations: &[TypedAllocation],
client: &RpcClient, client: &RpcClient,
args: &DistributeTokensArgs, args: &DistributeTokensArgs,
created_accounts: u64, created_accounts: u64,
@ -112,10 +112,10 @@ pub(crate) fn check_spl_token_balances(
pub(crate) fn print_token_balances( pub(crate) fn print_token_balances(
client: &RpcClient, client: &RpcClient,
allocation: &Allocation, allocation: &TypedAllocation,
spl_token_args: &SplTokenArgs, spl_token_args: &SplTokenArgs,
) -> Result<(), Error> { ) -> Result<(), Error> {
let address = allocation.recipient.parse().unwrap(); let address = allocation.recipient;
let expected = allocation.amount; let expected = allocation.amount;
let associated_token_address = get_associated_token_address(&address, &spl_token_args.mint); let associated_token_address = get_associated_token_address(&address, &spl_token_args.mint);
let recipient_account = client let recipient_account = client