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