From a49235796423a2b0b8b48eeaacfdafc712611372 Mon Sep 17 00:00:00 2001 From: Tyera Eulberg Date: Wed, 16 Sep 2020 11:42:31 -0600 Subject: [PATCH] v1.2: Backports of solana tokens improvements, including transfer to many arg (#12278) * Clean up solana-tokens (#10667) * Use a trait object in solana-tokens' ThinClient * Inline arg resolution Not worth the code complexity Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> * Add lockups via solana-tokens (bp #11782) (#12263) * Add lockups via solana-tokens (#11782) * Allow stake distributions to update lockups * Reorg * Add lockup test * Fix clippy warning (cherry picked from commit 5553732ae2e629e5fc53d5c57b6ad70fff40cc4d) * Fix build Co-authored-by: Greg Fitzgerald Co-authored-by: Tyera Eulberg * Improve solana-tokens UX (#12253) (#12260) * Fix computed banks port * Readme incorrect * Return error if csv cannot be read * Move column headers over columns * Add dry-run check for sender/fee-payer balances * Use clap requires method for paired args * Write transaction-log anytime outfile is specified * Replace campaign-name with required db-path * Remove bids * Exclude new_stake_account_address from logs for non-stake distributions * Fix readme Co-authored-by: Tyera Eulberg * solana-tokens: Add capability to perform the same transfer to a batch of recipients (bp #12259) (#12266) * solana-tokens: Add capability to perform the same transfer to a batch of recipients (#12259) * Add transfer-amount argument, use simplified input-csv * Add transfer-amount to readme (cherry picked from commit a48cc073cfc6cc9a7b18cc52ac8a842d3874e8e2) # Conflicts: # tokens/src/commands.rs # tokens/tests/commands.rs * Fix build Co-authored-by: Tyera Eulberg Co-authored-by: Tyera Eulberg Co-authored-by: Greg Fitzgerald Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- Cargo.lock | 1 + tokens/Cargo.toml | 1 + tokens/README.md | 76 ++++-- tokens/src/arg_parser.rs | 242 +++++++++++------ tokens/src/args.rs | 106 ++------ tokens/src/commands.rs | 550 +++++++++++++++++++++++++++----------- tokens/src/db.rs | 32 +++ tokens/src/main.rs | 13 +- tokens/src/thin_client.rs | 13 +- tokens/tests/commands.rs | 2 +- 10 files changed, 665 insertions(+), 371 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 0435a553c3..e61dae8530 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4462,6 +4462,7 @@ dependencies = [ name = "solana-tokens" version = "1.2.29" dependencies = [ + "bincode", "chrono", "clap", "console 0.10.3", diff --git a/tokens/Cargo.toml b/tokens/Cargo.toml index 298b30496a..61782b862d 100644 --- a/tokens/Cargo.toml +++ b/tokens/Cargo.toml @@ -30,4 +30,5 @@ tempfile = "3.1.0" thiserror = "1.0" [dev-dependencies] +bincode = "1.3.1" solana-core = { path = "../core", version = "1.2.29" } diff --git a/tokens/README.md b/tokens/README.md index 940183dfd7..bb4853f1ed 100644 --- a/tokens/README.md +++ b/tokens/README.md @@ -7,38 +7,38 @@ expected amount are sent. The command-line tool here automates that process. ## Distribute tokens -Send tokens to the recipients in ``. +Send tokens to the recipients in ``. -Example bids.csv: +Example recipients.csv: ```text -primary_address,bid_amount_dollars -6Vo87BaDhp4v4GHwVDhw5huhxVF8CyxSXYtkUwVHbbPv,6.6 +recipient,amount,lockup_date +3ihfUy1n9gaqihM5bJCiTAGLgWc5zo3DqVUS6T736NLM,42.0, +CYRJWqiSjLitBAcRxPvWpgX3s5TvmN2SuRY3eEYypFvT,43.0, ``` ```bash -solana-tokens distribute-tokens --from --dollars-per-sol --from-bids --input-csv --fee-payer +solana-tokens distribute-tokens --from --input-csv --fee-payer ``` Example transaction log before: ```text -recipient,amount,signature -6Vo87BaDhp4v4GHwVDhw5huhxVF8CyxSXYtkUwVHbbPv,30,1111111111111111111111111111111111111111111111111111111111111111 +recipient,amount,finalized_date,signature +6Vo87BaDhp4v4GHwVDhw5huhxVF8CyxSXYtkUwVHbbPv,70.0,2020-09-15T23:29:26.879747Z,UB168XhBhecxzeD1w2ZRUhwTHpPSqv2WNh8NrZHqz1F2EqxxbSW6iFfVtsg3HkU9NX2cD7R92D8VRLSyArZ9xKQ ``` -Send tokens to the recipients in `` if the distribution is -not already recordered in the transaction log. +Send tokens to the recipients in `` if the distribution is +not already recorded in the transaction log. ```bash -solana-tokens distribute-tokens --from --dollars-per-sol --from-bids --input-csv --fee-payer +solana-tokens distribute-tokens --from --input-csv --fee-payer ``` Example output: ```text -Recipient Amount -6Vo87BaDhp4v4GHwVDhw5huhxVF8CyxSXYtkUwVHbbPv 70 +Recipient Expected Balance (◎) 3ihfUy1n9gaqihM5bJCiTAGLgWc5zo3DqVUS6T736NLM 42 UKUcTXgbeTYh65RaVV5gSf6xBHevqHvAXMo3e8Q6np8k 43 ``` @@ -52,10 +52,9 @@ solana-tokens transaction-log --output-path transactions.csv ```text recipient,amount,signature -6Vo87BaDhp4v4GHwVDhw5huhxVF8CyxSXYtkUwVHbbPv,30,1111111111111111111111111111111111111111111111111111111111111111 -6Vo87BaDhp4v4GHwVDhw5huhxVF8CyxSXYtkUwVHbbPv,70,1111111111111111111111111111111111111111111111111111111111111111 -3ihfUy1n9gaqihM5bJCiTAGLgWc5zo3DqVUS6T736NLM,42,1111111111111111111111111111111111111111111111111111111111111111 -UKUcTXgbeTYh65RaVV5gSf6xBHevqHvAXMo3e8Q6np8k,43,1111111111111111111111111111111111111111111111111111111111111111 +6Vo87BaDhp4v4GHwVDhw5huhxVF8CyxSXYtkUwVHbbPv,70.0,2020-09-15T23:29:26.879747Z,UB168XhBhecxzeD1w2ZRUhwTHpPSqv2WNh8NrZHqz1F2EqxxbSW6iFfVtsg3HkU9NX2cD7R92D8VRLSyArZ9xKQ +3ihfUy1n9gaqihM5bJCiTAGLgWc5zo3DqVUS6T736NLM,42.0,2020-09-15T23:31:50.264241Z,53AVNEVpQBteJBRAKp6naxXsgESDjqe1ge9Dg2HeCSpYWTuGTLqHrBpkHTnpvPJURNgKWxkJfihuRa5STVRjL2hy +CYRJWqiSjLitBAcRxPvWpgX3s5TvmN2SuRY3eEYypFvT,43.0,2020-09-15T23:33:53.680821Z,4XsMfLx9D2ZxVpdJ5xdkV2w4X4SKEQ5zbQhcH4NcRwgZDkdRNiZjvnMFaWaWHUh5eF1LwFPpQdjn6mzSsiCVj3L7 ``` ### Calculate what tokens should be sent @@ -64,26 +63,49 @@ List the differences between a list of expected distributions and the record of transactions have already been sent. ```bash -solana-tokens distribute-tokens --dollars-per-sol --dry-run --from-bids --input-csv +solana-tokens distribute-tokens --dry-run --input-csv ``` -Example bids.csv: +Example recipients.csv: ```text -primary_address,bid_amount_dollars -6Vo87BaDhp4v4GHwVDhw5huhxVF8CyxSXYtkUwVHbbPv,6.6 -6Vo87BaDhp4v4GHwVDhw5huhxVF8CyxSXYtkUwVHbbPv,15.4 -3ihfUy1n9gaqihM5bJCiTAGLgWc5zo3DqVUS6T736NLM,9.24 -UKUcTXgbeTYh65RaVV5gSf6xBHevqHvAXMo3e8Q6np8k,9.46 +recipient,amount,lockup_date +6Vo87BaDhp4v4GHwVDhw5huhxVF8CyxSXYtkUwVHbbPv,80, +7aHDubg5FBYj1SgmyBgU3ZJdtfuqYCQsJQK2pTR5JUqr,42, ``` Example output: ```text -Recipient Amount -6Vo87BaDhp4v4GHwVDhw5huhxVF8CyxSXYtkUwVHbbPv 70 -3ihfUy1n9gaqihM5bJCiTAGLgWc5zo3DqVUS6T736NLM 42 -UKUcTXgbeTYh65RaVV5gSf6xBHevqHvAXMo3e8Q6np8k 43 +Recipient Expected Balance (◎) +6Vo87BaDhp4v4GHwVDhw5huhxVF8CyxSXYtkUwVHbbPv 10 +7aHDubg5FBYj1SgmyBgU3ZJdtfuqYCQsJQK2pTR5JUqr 42 +``` + +## Distribute tokens: transfer-amount + +This tool also makes it straightforward to transfer the same amount of tokens to a simple list of recipients. Just add the `--transfer-amount` arg to specify the amount: + +Example recipients.csv: + +```text +recipient +6Vo87BaDhp4v4GHwVDhw5huhxVF8CyxSXYtkUwVHbbPv +7aHDubg5FBYj1SgmyBgU3ZJdtfuqYCQsJQK2pTR5JUqr +CYRJWqiSjLitBAcRxPvWpgX3s5TvmN2SuRY3eEYypFvT +``` + +```bash +solana-tokens distribute-tokens --transfer-amount 10 --from --input-csv --fee-payer +``` + +Example output: + +```text +Recipient Expected Balance (◎) +6Vo87BaDhp4v4GHwVDhw5huhxVF8CyxSXYtkUwVHbbPv 10 +7aHDubg5FBYj1SgmyBgU3ZJdtfuqYCQsJQK2pTR5JUqr 10 +CYRJWqiSjLitBAcRxPvWpgX3s5TvmN2SuRY3eEYypFvT 10 ``` ## Distribute stake accounts diff --git a/tokens/src/arg_parser.rs b/tokens/src/arg_parser.rs index e61bf8047f..cbf9922346 100644 --- a/tokens/src/arg_parser.rs +++ b/tokens/src/arg_parser.rs @@ -2,8 +2,14 @@ use crate::args::{ Args, BalancesArgs, Command, DistributeTokensArgs, StakeArgs, TransactionLogArgs, }; use clap::{value_t, value_t_or_exit, App, Arg, ArgMatches, SubCommand}; -use solana_clap_utils::input_validators::{is_valid_pubkey, is_valid_signer}; +use solana_clap_utils::{ + input_parsers::value_of, + input_validators::{is_amount, is_valid_pubkey, is_valid_signer}, + keypair::{pubkey_from_path, signer_from_path}, +}; use solana_cli_config::CONFIG_FILE; +use solana_remote_wallet::remote_wallet::maybe_wallet_manager; +use std::error::Error; use std::ffi::OsString; use std::process::exit; @@ -36,16 +42,16 @@ where SubCommand::with_name("distribute-tokens") .about("Distribute tokens") .arg( - Arg::with_name("campaign_name") - .long("campaign-name") + Arg::with_name("db_path") + .long("db-path") + .required(true) .takes_value(true) - .value_name("NAME") - .help("Campaign name for storing transaction data"), - ) - .arg( - Arg::with_name("from_bids") - .long("from-bids") - .help("Input CSV contains bids in dollars, not allocations in SOL"), + .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") @@ -56,17 +62,26 @@ where .help("Input CSV file"), ) .arg( - Arg::with_name("dollars_per_sol") - .long("dollars-per-sol") + Arg::with_name("transfer_amount") + .long("transfer-amount") .takes_value(true) - .value_name("NUMBER") - .help("Dollars per SOL, if input CSV contains bids"), + .value_name("AMOUNT") + .validator(is_amount) + .help("The amount to send to each recipient, in SOL"), ) .arg( Arg::with_name("dry_run") .long("dry-run") .help("Do not execute any transfers"), ) + .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("sender_keypair") .long("from") @@ -90,11 +105,16 @@ where SubCommand::with_name("distribute-stake") .about("Distribute stake accounts") .arg( - Arg::with_name("campaign_name") - .long("campaign-name") + Arg::with_name("db_path") + .long("db-path") + .required(true) .takes_value(true) - .value_name("NAME") - .help("Campaign name for storing transaction data"), + .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") @@ -109,6 +129,14 @@ where .long("dry-run") .help("Do not execute any transfers"), ) + .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("sender_keypair") .long("from") @@ -153,6 +181,14 @@ where .validator(is_valid_signer) .help("Withdraw Authority Keypair"), ) + .arg( + Arg::with_name("lockup_authority") + .long("lockup-authority") + .takes_value(true) + .value_name("KEYPAIR") + .validator(is_valid_signer) + .help("Lockup Authority Keypair"), + ) .arg( Arg::with_name("fee_payer") .long("fee-payer") @@ -173,29 +209,18 @@ where .takes_value(true) .value_name("FILE") .help("Bids CSV file"), - ) - .arg( - Arg::with_name("from_bids") - .long("from-bids") - .help("Input CSV contains bids in dollars, not allocations in SOL"), - ) - .arg( - Arg::with_name("dollars_per_sol") - .long("dollars-per-sol") - .takes_value(true) - .value_name("NUMBER") - .help("Dollars per SOL"), ), ) .subcommand( SubCommand::with_name("transaction-log") .about("Print the database to a CSV file") .arg( - Arg::with_name("campaign_name") - .long("campaign-name") + Arg::with_name("db_path") + .long("db-path") + .required(true) .takes_value(true) - .value_name("NAME") - .help("Campaign name for storing transaction data"), + .value_name("FILE") + .help("Location of database to query"), ) .arg( Arg::with_name("output_path") @@ -209,70 +234,130 @@ where .get_matches_from(args) } -fn create_db_path(campaign_name: Option) -> String { - let (prefix, hyphen) = if let Some(name) = campaign_name { - (name, "-") - } else { - ("".to_string(), "") - }; - let path = dirs::home_dir().unwrap(); - let filename = format!("{}{}transactions.db", prefix, hyphen); - path.join(".config") - .join("solana-tokens") - .join(filename) - .to_str() - .unwrap() - .to_string() -} +fn parse_distribute_tokens_args( + matches: &ArgMatches<'_>, +) -> Result> { + let mut wallet_manager = maybe_wallet_manager()?; + let signer_matches = ArgMatches::default(); // No default signer -fn parse_distribute_tokens_args(matches: &ArgMatches<'_>) -> DistributeTokensArgs { - DistributeTokensArgs { + let sender_keypair_str = value_t_or_exit!(matches, "sender_keypair", String); + let sender_keypair = signer_from_path( + &signer_matches, + &sender_keypair_str, + "sender", + &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, + )?; + + Ok(DistributeTokensArgs { input_csv: value_t_or_exit!(matches, "input_csv", String), - from_bids: matches.is_present("from_bids"), - transaction_db: create_db_path(value_t!(matches, "campaign_name", String).ok()), - dollars_per_sol: value_t!(matches, "dollars_per_sol", f64).ok(), + 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: value_t_or_exit!(matches, "sender_keypair", String), - fee_payer: value_t_or_exit!(matches, "fee_payer", String), + sender_keypair, + fee_payer, stake_args: None, - } + transfer_amount: value_of(matches, "transfer_amount"), + }) } -fn parse_distribute_stake_args(matches: &ArgMatches<'_>) -> DistributeTokensArgs { - let stake_args = StakeArgs { - stake_account_address: value_t_or_exit!(matches, "stake_account_address", String), - sol_for_fees: value_t_or_exit!(matches, "sol_for_fees", f64), - stake_authority: value_t_or_exit!(matches, "stake_authority", String), - withdraw_authority: value_t_or_exit!(matches, "withdraw_authority", String), +fn parse_distribute_stake_args( + matches: &ArgMatches<'_>, +) -> Result> { + let mut wallet_manager = maybe_wallet_manager()?; + let signer_matches = ArgMatches::default(); // No default signer + + let sender_keypair_str = value_t_or_exit!(matches, "sender_keypair", String); + let sender_keypair = signer_from_path( + &signer_matches, + &sender_keypair_str, + "sender", + &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 stake_account_address_str = value_t_or_exit!(matches, "stake_account_address", String); + let stake_account_address = pubkey_from_path( + &signer_matches, + &stake_account_address_str, + "stake account address", + &mut wallet_manager, + )?; + + let stake_authority_str = value_t_or_exit!(matches, "stake_authority", String); + let stake_authority = signer_from_path( + &signer_matches, + &stake_authority_str, + "stake authority", + &mut wallet_manager, + )?; + + let withdraw_authority_str = value_t_or_exit!(matches, "withdraw_authority", String); + let withdraw_authority = signer_from_path( + &signer_matches, + &withdraw_authority_str, + "withdraw authority", + &mut wallet_manager, + )?; + + let lockup_authority_str = value_t!(matches, "lockup_authority", String).ok(); + let lockup_authority = match lockup_authority_str { + Some(path) => Some(signer_from_path( + &signer_matches, + &path, + "lockup authority", + &mut wallet_manager, + )?), + None => None, }; - DistributeTokensArgs { + + let stake_args = StakeArgs { + stake_account_address, + sol_for_fees: value_t_or_exit!(matches, "sol_for_fees", f64), + stake_authority, + withdraw_authority, + lockup_authority, + }; + Ok(DistributeTokensArgs { input_csv: value_t_or_exit!(matches, "input_csv", String), - from_bids: false, - transaction_db: create_db_path(value_t!(matches, "campaign_name", String).ok()), - dollars_per_sol: None, + 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: value_t_or_exit!(matches, "sender_keypair", String), - fee_payer: value_t_or_exit!(matches, "fee_payer", String), + sender_keypair, + fee_payer, stake_args: Some(stake_args), - } + transfer_amount: None, + }) } fn parse_balances_args(matches: &ArgMatches<'_>) -> BalancesArgs { BalancesArgs { input_csv: value_t_or_exit!(matches, "input_csv", String), - from_bids: matches.is_present("from_bids"), - dollars_per_sol: value_t!(matches, "dollars_per_sol", f64).ok(), } } fn parse_transaction_log_args(matches: &ArgMatches<'_>) -> TransactionLogArgs { TransactionLogArgs { - transaction_db: create_db_path(value_t!(matches, "campaign_name", String).ok()), + transaction_db: value_t_or_exit!(matches, "db_path", String), output_path: value_t_or_exit!(matches, "output_path", String), } } -pub fn parse_args(args: I) -> Args +pub fn parse_args(args: I) -> Result> where I: IntoIterator, T: Into + Clone, @@ -283,10 +368,10 @@ where let command = match matches.subcommand() { ("distribute-tokens", Some(matches)) => { - Command::DistributeTokens(parse_distribute_tokens_args(matches)) + Command::DistributeTokens(parse_distribute_tokens_args(matches)?) } ("distribute-stake", Some(matches)) => { - Command::DistributeTokens(parse_distribute_stake_args(matches)) + Command::DistributeTokens(parse_distribute_stake_args(matches)?) } ("balances", Some(matches)) => Command::Balances(parse_balances_args(matches)), ("transaction-log", Some(matches)) => { @@ -297,9 +382,10 @@ where exit(1); } }; - Args { + let args = Args { config_file, url, command, - } + }; + Ok(args) } diff --git a/tokens/src/args.rs b/tokens/src/args.rs index 0eed694638..8c28f99044 100644 --- a/tokens/src/args.rs +++ b/tokens/src/args.rs @@ -1,31 +1,26 @@ -use clap::ArgMatches; -use solana_clap_utils::keypair::{pubkey_from_path, signer_from_path}; -use solana_remote_wallet::remote_wallet::{maybe_wallet_manager, RemoteWalletManager}; use solana_sdk::{pubkey::Pubkey, signature::Signer}; -use std::{error::Error, sync::Arc}; -pub struct DistributeTokensArgs { +pub struct DistributeTokensArgs { pub input_csv: String, - pub from_bids: bool, pub transaction_db: String, - pub dollars_per_sol: Option, + pub output_path: Option, pub dry_run: bool, - pub sender_keypair: K, - pub fee_payer: K, - pub stake_args: Option>, + pub sender_keypair: Box, + pub fee_payer: Box, + pub stake_args: Option, + pub transfer_amount: Option, } -pub struct StakeArgs { +pub struct StakeArgs { pub sol_for_fees: f64, - pub stake_account_address: P, - pub stake_authority: K, - pub withdraw_authority: K, + pub stake_account_address: Pubkey, + pub stake_authority: Box, + pub withdraw_authority: Box, + pub lockup_authority: Option>, } pub struct BalancesArgs { pub input_csv: String, - pub from_bids: bool, - pub dollars_per_sol: Option, } pub struct TransactionLogArgs { @@ -33,85 +28,14 @@ pub struct TransactionLogArgs { pub output_path: String, } -pub enum Command { - DistributeTokens(DistributeTokensArgs), +pub enum Command { + DistributeTokens(DistributeTokensArgs), Balances(BalancesArgs), TransactionLog(TransactionLogArgs), } -pub struct Args { +pub struct Args { pub config_file: String, pub url: Option, - pub command: Command, -} - -pub fn resolve_stake_args( - wallet_manager: &mut Option>, - args: StakeArgs, -) -> Result>, Box> { - let matches = ArgMatches::default(); - let resolved_args = StakeArgs { - stake_account_address: pubkey_from_path( - &matches, - &args.stake_account_address, - "stake account address", - wallet_manager, - ) - .unwrap(), - sol_for_fees: args.sol_for_fees, - stake_authority: signer_from_path( - &matches, - &args.stake_authority, - "stake authority", - wallet_manager, - ) - .unwrap(), - withdraw_authority: signer_from_path( - &matches, - &args.withdraw_authority, - "withdraw authority", - wallet_manager, - ) - .unwrap(), - }; - Ok(resolved_args) -} - -pub fn resolve_command( - command: Command, -) -> Result>, Box> { - match command { - Command::DistributeTokens(args) => { - let mut wallet_manager = maybe_wallet_manager()?; - let matches = ArgMatches::default(); - let resolved_stake_args = args - .stake_args - .map(|args| resolve_stake_args(&mut wallet_manager, args)); - let resolved_args = DistributeTokensArgs { - input_csv: args.input_csv, - from_bids: args.from_bids, - transaction_db: args.transaction_db, - dollars_per_sol: args.dollars_per_sol, - dry_run: args.dry_run, - sender_keypair: signer_from_path( - &matches, - &args.sender_keypair, - "sender", - &mut wallet_manager, - ) - .unwrap(), - fee_payer: signer_from_path( - &matches, - &args.fee_payer, - "fee-payer", - &mut wallet_manager, - ) - .unwrap(), - stake_args: resolved_stake_args.map_or(Ok(None), |r| r.map(Some))?, - }; - Ok(Command::DistributeTokens(resolved_args)) - } - Command::Balances(args) => Ok(Command::Balances(args)), - Command::TransactionLog(args) => Ok(Command::TransactionLog(args)), - } + pub command: Command, } diff --git a/tokens/src/commands.rs b/tokens/src/commands.rs index 217762f22d..95e32c94d7 100644 --- a/tokens/src/commands.rs +++ b/tokens/src/commands.rs @@ -1,6 +1,9 @@ -use crate::args::{BalancesArgs, DistributeTokensArgs, StakeArgs, TransactionLogArgs}; -use crate::db::{self, TransactionInfo}; -use crate::thin_client::{Client, ThinClient}; +use crate::{ + args::{BalancesArgs, DistributeTokensArgs, StakeArgs, TransactionLogArgs}, + db::{self, TransactionInfo}, + thin_client::{Client, ThinClient}, +}; +use chrono::prelude::*; use console::style; use csv::{ReaderBuilder, Trim}; use indexmap::IndexMap; @@ -8,6 +11,7 @@ use indicatif::{ProgressBar, ProgressStyle}; use pickledb::PickleDb; use serde::{Deserialize, Serialize}; use solana_sdk::{ + instruction::Instruction, message::Message, native_token::{lamports_to_sol, sol_to_lamports}, signature::{unique_signers, Signature, Signer}, @@ -15,7 +19,7 @@ use solana_sdk::{ transport::TransportError, }; use solana_stake_program::{ - stake_instruction, + stake_instruction::{self, LockupArgs}, stake_state::{Authorized, Lockup, StakeAuthorize}, }; use std::{ @@ -35,6 +39,7 @@ struct Bid { struct Allocation { recipient: String, amount: f64, + lockup_date: String, } #[derive(thiserror::Error, Debug)] @@ -47,6 +52,14 @@ pub enum Error { PickleDbError(#[from] pickledb::error::Error), #[error("Transport error")] TransportError(#[from] TransportError), + #[error("Missing lockup authority")] + MissingLockupAuthority, + #[error("insufficient funds for fee ({0} SOL)")] + InsufficientFundsForFees(f64), + #[error("insufficient funds for distribution ({0} SOL)")] + InsufficientFundsForDistribution(f64), + #[error("insufficient funds for distribution ({0} SOL) and fee ({1} SOL)")] + InsufficientFundsForDistributionAndFees(f64, f64), } fn merge_allocations(allocations: &[Allocation]) -> Vec { @@ -57,12 +70,19 @@ fn merge_allocations(allocations: &[Allocation]) -> Vec { .or_insert(Allocation { recipient: allocation.recipient.clone(), amount: 0.0, + lockup_date: "".to_string(), }) .amount += allocation.amount; } allocation_map.values().cloned().collect() } +/// 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 apply_previous_transactions( allocations: &mut Vec, transaction_infos: &[TransactionInfo], @@ -70,7 +90,7 @@ fn apply_previous_transactions( for transaction_info in transaction_infos { let mut amount = transaction_info.amount; for allocation in allocations.iter_mut() { - if allocation.recipient != transaction_info.recipient.to_string() { + if !has_same_recipient(&allocation, &transaction_info) { continue; } if allocation.amount >= amount { @@ -85,19 +105,86 @@ fn apply_previous_transactions( allocations.retain(|x| x.amount > 0.5); } -fn create_allocation(bid: &Bid, dollars_per_sol: f64) -> Allocation { - Allocation { - recipient: bid.primary_address.clone(), - amount: bid.accepted_amount_dollars / dollars_per_sol, +fn distribution_instructions( + allocation: &Allocation, + new_stake_account_address: &Pubkey, + args: &DistributeTokensArgs, + lockup_date: Option>, +) -> Vec { + if args.stake_args.is_none() { + let from = args.sender_keypair.pubkey(); + let to = allocation.recipient.parse().unwrap(); + let lamports = sol_to_lamports(allocation.amount); + let instruction = system_instruction::transfer(&from, &to, lamports); + return vec![instruction]; } + + let stake_args = args.stake_args.as_ref().unwrap(); + let sol_for_fees = stake_args.sol_for_fees; + let sender_pubkey = args.sender_keypair.pubkey(); + let stake_authority = stake_args.stake_authority.pubkey(); + let withdraw_authority = stake_args.withdraw_authority.pubkey(); + + let mut instructions = stake_instruction::split( + &stake_args.stake_account_address, + &stake_authority, + sol_to_lamports(allocation.amount - sol_for_fees), + &new_stake_account_address, + ); + + let recipient = allocation.recipient.parse().unwrap(); + + // Make the recipient the new stake authority + instructions.push(stake_instruction::authorize( + &new_stake_account_address, + &stake_authority, + &recipient, + StakeAuthorize::Staker, + )); + + // Make the recipient the new withdraw authority + instructions.push(stake_instruction::authorize( + &new_stake_account_address, + &withdraw_authority, + &recipient, + StakeAuthorize::Withdrawer, + )); + + // Add lockup + if let Some(lockup_date) = lockup_date { + let lockup_authority = stake_args + .lockup_authority + .as_ref() + .map(|signer| signer.pubkey()) + .unwrap(); + let lockup = LockupArgs { + unix_timestamp: Some(lockup_date.timestamp()), + epoch: None, + custodian: None, + }; + instructions.push(stake_instruction::set_lockup( + &new_stake_account_address, + &lockup, + &lockup_authority, + )); + } + + instructions.push(system_instruction::transfer( + &sender_pubkey, + &recipient, + sol_to_lamports(sol_for_fees), + )); + + instructions } -fn distribute_tokens( - client: &ThinClient, +fn distribute_allocations( + client: &ThinClient, db: &mut PickleDb, allocations: &[Allocation], - args: &DistributeTokensArgs>, + args: &DistributeTokensArgs, ) -> Result<(), Error> { + let mut num_signatures = 0; for allocation in allocations { let new_stake_account_keypair = Keypair::new(); let new_stake_account_address = new_stake_account_keypair.pubkey(); @@ -107,56 +194,26 @@ fn distribute_tokens( signers.push(&*stake_args.stake_authority); signers.push(&*stake_args.withdraw_authority); signers.push(&new_stake_account_keypair); + if allocation.lockup_date != "" { + if let Some(lockup_authority) = &stake_args.lockup_authority { + signers.push(&**lockup_authority); + } else { + return Err(Error::MissingLockupAuthority); + } + } } let signers = unique_signers(signers); + num_signatures += signers.len(); - println!("{:<44} {:>24.9}", allocation.recipient, allocation.amount); - let instructions = if let Some(stake_args) = &args.stake_args { - let sol_for_fees = stake_args.sol_for_fees; - let sender_pubkey = args.sender_keypair.pubkey(); - let stake_authority = stake_args.stake_authority.pubkey(); - let withdraw_authority = stake_args.withdraw_authority.pubkey(); - - let mut instructions = stake_instruction::split( - &stake_args.stake_account_address, - &stake_authority, - sol_to_lamports(allocation.amount - sol_for_fees), - &new_stake_account_address, - ); - - let recipient = allocation.recipient.parse().unwrap(); - - // Make the recipient the new stake authority - instructions.push(stake_instruction::authorize( - &new_stake_account_address, - &stake_authority, - &recipient, - StakeAuthorize::Staker, - )); - - // Make the recipient the new withdraw authority - instructions.push(stake_instruction::authorize( - &new_stake_account_address, - &withdraw_authority, - &recipient, - StakeAuthorize::Withdrawer, - )); - - instructions.push(system_instruction::transfer( - &sender_pubkey, - &recipient, - sol_to_lamports(sol_for_fees), - )); - - instructions + let lockup_date = if allocation.lockup_date == "" { + None } else { - let from = args.sender_keypair.pubkey(); - let to = allocation.recipient.parse().unwrap(); - let lamports = sol_to_lamports(allocation.amount); - let instruction = system_instruction::transfer(&from, &to, lamports); - vec![instruction] + Some(allocation.lockup_date.parse::>().unwrap()) }; + println!("{:<44} {:>24.9}", allocation.recipient, allocation.amount); + let instructions = + distribution_instructions(allocation, &new_stake_account_address, args, lockup_date); let fee_payer_pubkey = args.fee_payer.pubkey(); let message = Message::new(&instructions, Some(&fee_payer_pubkey)); match client.send_and_confirm_message(message, &signers) { @@ -166,9 +223,10 @@ fn distribute_tokens( &allocation.recipient.parse().unwrap(), allocation.amount, &transaction, - Some(&new_stake_account_address), + args.stake_args.as_ref().map(|_| &new_stake_account_address), false, last_valid_slot, + lockup_date, )?; } Err(e) => { @@ -176,26 +234,37 @@ fn distribute_tokens( } }; } + if args.dry_run { + let undistributed_tokens: f64 = allocations.iter().map(|x| x.amount).sum(); + check_payer_balances( + num_signatures, + sol_to_lamports(undistributed_tokens), + client, + args, + )?; + } Ok(()) } -fn read_allocations( - input_csv: &str, - from_bids: bool, - dollars_per_sol: Option, -) -> Vec { - let rdr = ReaderBuilder::new().trim(Trim::All).from_path(input_csv); - if from_bids { - let bids: Vec = rdr.unwrap().deserialize().map(|bid| bid.unwrap()).collect(); - bids.into_iter() - .map(|bid| create_allocation(&bid, dollars_per_sol.unwrap())) +fn read_allocations(input_csv: &str, transfer_amount: Option) -> io::Result> { + 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(), + }) .collect() } else { - rdr.unwrap() - .deserialize() - .map(|entry| entry.unwrap()) - .collect() - } + rdr.deserialize().map(|entry| entry.unwrap()).collect() + }; + Ok(allocations) } fn new_spinner_progress_bar() -> ProgressBar { @@ -206,12 +275,11 @@ fn new_spinner_progress_bar() -> ProgressBar { progress_bar } -pub fn process_distribute_tokens( - client: &ThinClient, - args: &DistributeTokensArgs>, +pub fn process_allocations( + client: &ThinClient, + args: &DistributeTokensArgs, ) -> Result, Error> { - let mut allocations: Vec = - read_allocations(&args.input_csv, args.from_bids, args.dollars_per_sol); + let mut allocations: Vec = read_allocations(&args.input_csv, args.transfer_amount)?; let starting_total_tokens: f64 = allocations.iter().map(|x| x.amount).sum(); println!( @@ -219,13 +287,6 @@ pub fn process_distribute_tokens( style("Total in input_csv:").bold(), starting_total_tokens, ); - if let Some(dollars_per_sol) = args.dollars_per_sol { - println!( - "{} ${}", - style("Total in input_csv:").bold(), - starting_total_tokens * dollars_per_sol, - ); - } let mut db = db::open_db(&args.transaction_db, args.dry_run)?; @@ -240,6 +301,20 @@ pub fn process_distribute_tokens( 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,); + println!( + "{} ◎{}", + style("Undistributed:").bold(), + undistributed_tokens, + ); + println!( + "{} ◎{}", + style("Total:").bold(), + distributed_tokens + undistributed_tokens, + ); + println!( "{}", style(format!( @@ -249,49 +324,20 @@ pub fn process_distribute_tokens( .bold() ); - 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,); - if let Some(dollars_per_sol) = args.dollars_per_sol { - println!( - "{} ${}", - style("Distributed:").bold(), - distributed_tokens * dollars_per_sol, - ); - } - println!( - "{} ◎{}", - style("Undistributed:").bold(), - undistributed_tokens, - ); - if let Some(dollars_per_sol) = args.dollars_per_sol { - println!( - "{} ${}", - style("Undistributed:").bold(), - undistributed_tokens * dollars_per_sol, - ); - } - println!( - "{} ◎{}", - style("Total:").bold(), - distributed_tokens + undistributed_tokens, - ); - if let Some(dollars_per_sol) = args.dollars_per_sol { - println!( - "{} ${}", - style("Total:").bold(), - (distributed_tokens + undistributed_tokens) * dollars_per_sol, - ); - } - - distribute_tokens(client, &mut db, &allocations, args)?; + distribute_allocations(client, &mut db, &allocations, args)?; let opt_confirmations = finalize_transactions(client, &mut db, args.dry_run)?; + + if !args.dry_run { + if let Some(output_path) = &args.output_path { + db::write_transaction_log(&db, &output_path)?; + } + } Ok(opt_confirmations) } -fn finalize_transactions( - client: &ThinClient, +fn finalize_transactions( + client: &ThinClient, db: &mut PickleDb, dry_run: bool, ) -> Result, Error> { @@ -322,8 +368,8 @@ fn finalize_transactions( // Update the finalized bit on any transactions that are now rooted // Return the lowest number of confirmations on the unfinalized transactions or None if all are finalized. -fn update_finalized_transactions( - client: &ThinClient, +fn update_finalized_transactions( + client: &ThinClient, db: &mut PickleDb, ) -> Result, Error> { let transaction_infos = db::read_transaction_infos(db); @@ -368,12 +414,42 @@ fn update_finalized_transactions( Ok(confirmations) } -pub fn process_balances( - client: &ThinClient, - args: &BalancesArgs, -) -> Result<(), csv::Error> { - let allocations: Vec = - read_allocations(&args.input_csv, args.from_bids, args.dollars_per_sol); +fn check_payer_balances( + num_signatures: usize, + allocation_lamports: u64, + client: &ThinClient, + args: &DistributeTokensArgs, +) -> Result<(), Error> { + let (_blockhash, fee_calculator, _last_valid_slot) = client.get_fees()?; + let fees = fee_calculator + .lamports_per_signature + .checked_mul(num_signatures as u64) + .unwrap(); + if args.fee_payer.pubkey() == args.sender_keypair.pubkey() { + let balance = client.get_balance(&args.fee_payer.pubkey())?; + if balance < fees + allocation_lamports { + return Err(Error::InsufficientFundsForDistributionAndFees( + lamports_to_sol(allocation_lamports), + lamports_to_sol(fees), + )); + } + } else { + let fee_payer_balance = client.get_balance(&args.fee_payer.pubkey())?; + if fee_payer_balance < fees { + return Err(Error::InsufficientFundsForFees(lamports_to_sol(fees))); + } + let sender_balance = client.get_balance(&args.sender_keypair.pubkey())?; + if sender_balance < allocation_lamports { + return Err(Error::InsufficientFundsForDistribution(lamports_to_sol( + allocation_lamports, + ))); + } + } + Ok(()) +} + +pub fn process_balances(client: &ThinClient, args: &BalancesArgs) -> Result<(), csv::Error> { + let allocations: Vec = read_allocations(&args.input_csv, None)?; let allocations = merge_allocations(&allocations); println!( @@ -407,9 +483,14 @@ pub fn process_transaction_log(args: &TransactionLogArgs) -> Result<(), Error> { Ok(()) } +use crate::db::check_output_file; use solana_sdk::{pubkey::Pubkey, signature::Keypair}; use tempfile::{tempdir, NamedTempFile}; -pub fn test_process_distribute_tokens_with_client(client: C, sender_keypair: Keypair) { +pub fn test_process_distribute_tokens_with_client( + client: C, + sender_keypair: Keypair, + transfer_amount: Option, +) { let thin_client = ThinClient::new(client, false); let fee_payer = Keypair::new(); let (transaction, _last_valid_slot) = thin_client @@ -422,7 +503,12 @@ pub fn test_process_distribute_tokens_with_client(client: C, sender_k let alice_pubkey = Pubkey::new_rand(); let allocation = Allocation { recipient: alice_pubkey.to_string(), - amount: 1000.0, + amount: if let Some(amount) = transfer_amount { + amount + } else { + 1000.0 + }, + lockup_date: "".to_string(), }; let allocations_file = NamedTempFile::new().unwrap(); let input_csv = allocations_file.path().to_str().unwrap().to_string(); @@ -438,17 +524,20 @@ pub fn test_process_distribute_tokens_with_client(client: C, sender_k .unwrap() .to_string(); - let args: DistributeTokensArgs> = DistributeTokensArgs { + let output_file = NamedTempFile::new().unwrap(); + let output_path = output_file.path().to_str().unwrap().to_string(); + + let args = DistributeTokensArgs { sender_keypair: Box::new(sender_keypair), fee_payer: Box::new(fee_payer), dry_run: false, input_csv, - from_bids: false, transaction_db: transaction_db.clone(), - dollars_per_sol: None, + output_path: Some(output_path.clone()), stake_args: None, + transfer_amount, }; - let confirmations = process_distribute_tokens(&thin_client, &args).unwrap(); + let confirmations = process_allocations(&thin_client, &args).unwrap(); assert_eq!(confirmations, None); let transaction_infos = @@ -466,8 +555,10 @@ pub fn test_process_distribute_tokens_with_client(client: C, sender_k expected_amount, ); + check_output_file(&output_path, &db::open_db(&transaction_db, true).unwrap()); + // Now, run it again, and check there's no double-spend. - process_distribute_tokens(&thin_client, &args).unwrap(); + process_allocations(&thin_client, &args).unwrap(); let transaction_infos = db::read_transaction_infos(&db::open_db(&transaction_db, true).unwrap()); assert_eq!(transaction_infos.len(), 1); @@ -482,6 +573,8 @@ pub fn test_process_distribute_tokens_with_client(client: C, sender_k thin_client.get_balance(&alice_pubkey).unwrap(), expected_amount, ); + + check_output_file(&output_path, &db::open_db(&transaction_db, true).unwrap()); } pub fn test_process_distribute_stake_with_client(client: C, sender_keypair: Keypair) { @@ -521,6 +614,7 @@ pub fn test_process_distribute_stake_with_client(client: C, sender_ke let allocation = Allocation { recipient: alice_pubkey.to_string(), amount: 1000.0, + lockup_date: "".to_string(), }; let file = NamedTempFile::new().unwrap(); let input_csv = file.path().to_str().unwrap().to_string(); @@ -536,23 +630,27 @@ pub fn test_process_distribute_stake_with_client(client: C, sender_ke .unwrap() .to_string(); - let stake_args: StakeArgs> = StakeArgs { + let output_file = NamedTempFile::new().unwrap(); + let output_path = output_file.path().to_str().unwrap().to_string(); + + let stake_args = StakeArgs { stake_account_address, stake_authority: Box::new(stake_authority), withdraw_authority: Box::new(withdraw_authority), + lockup_authority: None, sol_for_fees: 1.0, }; - let args: DistributeTokensArgs> = DistributeTokensArgs { + let args = DistributeTokensArgs { fee_payer: Box::new(fee_payer), dry_run: false, input_csv, transaction_db: transaction_db.clone(), + output_path: Some(output_path.clone()), stake_args: Some(stake_args), - from_bids: false, sender_keypair: Box::new(sender_keypair), - dollars_per_sol: None, + transfer_amount: None, }; - let confirmations = process_distribute_tokens(&thin_client, &args).unwrap(); + let confirmations = process_allocations(&thin_client, &args).unwrap(); assert_eq!(confirmations, None); let transaction_infos = @@ -575,8 +673,10 @@ pub fn test_process_distribute_stake_with_client(client: C, sender_ke expected_amount - sol_to_lamports(1.0), ); + check_output_file(&output_path, &db::open_db(&transaction_db, true).unwrap()); + // Now, run it again, and check there's no double-spend. - process_distribute_tokens(&thin_client, &args).unwrap(); + process_allocations(&thin_client, &args).unwrap(); let transaction_infos = db::read_transaction_infos(&db::open_db(&transaction_db, true).unwrap()); assert_eq!(transaction_infos.len(), 1); @@ -595,6 +695,8 @@ pub fn test_process_distribute_stake_with_client(client: C, sender_ke thin_client.get_balance(&new_stake_account_address).unwrap(), expected_amount - sol_to_lamports(1.0), ); + + check_output_file(&output_path, &db::open_db(&transaction_db, true).unwrap()); } #[cfg(test)] @@ -602,17 +704,26 @@ mod tests { use super::*; use solana_runtime::{bank::Bank, bank_client::BankClient}; use solana_sdk::genesis_config::create_genesis_config; + use solana_stake_program::stake_instruction::StakeInstruction; #[test] - fn test_process_distribute_tokens() { + fn test_process_token_allocations() { let (genesis_config, sender_keypair) = create_genesis_config(sol_to_lamports(9_000_000.0)); let bank = Bank::new(&genesis_config); let bank_client = BankClient::new(bank); - test_process_distribute_tokens_with_client(bank_client, sender_keypair); + test_process_distribute_tokens_with_client(bank_client, sender_keypair, None); } #[test] - fn test_process_distribute_stake() { + fn test_process_transfer_amount_allocations() { + let (genesis_config, sender_keypair) = create_genesis_config(sol_to_lamports(9_000_000.0)); + let bank = Bank::new(&genesis_config); + let bank_client = BankClient::new(bank); + test_process_distribute_tokens_with_client(bank_client, sender_keypair, Some(1.5)); + } + + #[test] + fn test_process_stake_allocations() { let (genesis_config, sender_keypair) = create_genesis_config(sol_to_lamports(9_000_000.0)); let bank = Bank::new(&genesis_config); let bank_client = BankClient::new(bank); @@ -625,6 +736,7 @@ mod tests { let allocation = Allocation { recipient: alice_pubkey.to_string(), amount: 42.0, + lockup_date: "".to_string(), }; let file = NamedTempFile::new().unwrap(); let input_csv = file.path().to_str().unwrap().to_string(); @@ -632,29 +744,48 @@ mod tests { wtr.serialize(&allocation).unwrap(); wtr.flush().unwrap(); - assert_eq!(read_allocations(&input_csv, false, None), vec![allocation]); + assert_eq!( + read_allocations(&input_csv, None).unwrap(), + vec![allocation] + ); } #[test] - fn test_read_allocations_from_bids() { - let alice_pubkey = Pubkey::new_rand(); - let bid = Bid { - primary_address: alice_pubkey.to_string(), - accepted_amount_dollars: 42.0, - }; + fn test_read_allocations_transfer_amount() { + 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); - wtr.serialize(&bid).unwrap(); + wtr.serialize("recipient".to_string()).unwrap(); + wtr.serialize(&pubkey0.to_string()).unwrap(); + wtr.serialize(&pubkey1.to_string()).unwrap(); + wtr.serialize(&pubkey2.to_string()).unwrap(); wtr.flush().unwrap(); - let allocation = Allocation { - recipient: bid.primary_address, - amount: 84.0, - }; + let amount = 1.5; + + let expected_allocations = vec![ + Allocation { + recipient: pubkey0.to_string(), + amount, + lockup_date: "".to_string(), + }, + Allocation { + recipient: pubkey1.to_string(), + amount, + lockup_date: "".to_string(), + }, + Allocation { + recipient: pubkey2.to_string(), + amount, + lockup_date: "".to_string(), + }, + ]; assert_eq!( - read_allocations(&input_csv, true, Some(0.5)), - vec![allocation] + read_allocations(&input_csv, Some(amount)).unwrap(), + expected_allocations ); } @@ -666,10 +797,12 @@ mod tests { Allocation { recipient: alice.to_string(), amount: 1.0, + lockup_date: "".to_string(), }, Allocation { recipient: bob.to_string(), amount: 1.0, + lockup_date: "".to_string(), }, ]; let transaction_infos = vec![TransactionInfo { @@ -684,4 +817,101 @@ mod tests { // a matching recipient address (to bob, not alice). assert_eq!(allocations[0].recipient, alice.to_string()); } + + #[test] + fn test_has_same_recipient() { + 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(), + amount: 1.0, + lockup_date: "".to_string(), + }; + let alice_alloc_lockup0 = Allocation { + recipient: alice_pubkey.to_string(), + amount: 1.0, + lockup_date: lockup0.clone(), + }; + let alice_info = TransactionInfo { + recipient: alice_pubkey, + lockup_date: None, + ..TransactionInfo::default() + }; + let alice_info_lockup0 = TransactionInfo { + recipient: alice_pubkey, + lockup_date: lockup0.parse().ok(), + ..TransactionInfo::default() + }; + let alice_info_lockup1 = TransactionInfo { + recipient: alice_pubkey, + lockup_date: lockup1.parse().ok(), + ..TransactionInfo::default() + }; + let bob_info = TransactionInfo { + recipient: bob_pubkey, + lockup_date: None, + ..TransactionInfo::default() + }; + assert!(!has_same_recipient(&alice_alloc, &bob_info)); // Different recipient, no lockup + assert!(!has_same_recipient(&alice_alloc, &alice_info_lockup0)); // One with no lockup, one locked up + assert!(!has_same_recipient( + &alice_alloc_lockup0, + &alice_info_lockup1 + )); // Different lockups + assert!(has_same_recipient(&alice_alloc, &alice_info)); // Same recipient, no lockups + assert!(has_same_recipient( + &alice_alloc_lockup0, + &alice_info_lockup0 + )); // Same recipient, same lockups + } + + const SET_LOCKUP_INDEX: usize = 4; + + #[test] + fn test_set_stake_lockup() { + let lockup_date_str = "2021-01-07T00:00:00Z"; + let allocation = Allocation { + recipient: Pubkey::default().to_string(), + amount: 1.0, + lockup_date: lockup_date_str.to_string(), + }; + let stake_account_address = Pubkey::new_rand(); + let new_stake_account_address = Pubkey::new_rand(); + let lockup_authority = Keypair::new(); + let stake_args = StakeArgs { + stake_account_address, + stake_authority: Box::new(Keypair::new()), + withdraw_authority: Box::new(Keypair::new()), + lockup_authority: Some(Box::new(lockup_authority)), + sol_for_fees: 1.0, + }; + let args = DistributeTokensArgs { + fee_payer: Box::new(Keypair::new()), + dry_run: false, + input_csv: "".to_string(), + transaction_db: "".to_string(), + output_path: None, + stake_args: Some(stake_args), + sender_keypair: Box::new(Keypair::new()), + transfer_amount: None, + }; + let lockup_date = lockup_date_str.parse().unwrap(); + let instructions = distribution_instructions( + &allocation, + &new_stake_account_address, + &args, + Some(lockup_date), + ); + let lockup_instruction = + bincode::deserialize(&instructions[SET_LOCKUP_INDEX].data).unwrap(); + if let StakeInstruction::SetLockup(lockup_args) = lockup_instruction { + assert_eq!(lockup_args.unix_timestamp, Some(lockup_date.timestamp())); + assert_eq!(lockup_args.epoch, None); // Don't change the epoch + assert_eq!(lockup_args.custodian, None); // Don't change the lockup authority + } else { + panic!("expected SetLockup instruction"); + } + } } diff --git a/tokens/src/db.rs b/tokens/src/db.rs index d792bb3463..5e89e746bc 100644 --- a/tokens/src/db.rs +++ b/tokens/src/db.rs @@ -13,12 +13,14 @@ pub struct TransactionInfo { pub finalized_date: Option>, pub transaction: Transaction, pub last_valid_slot: Slot, + pub lockup_date: Option>, } #[derive(Serialize, Deserialize, Debug, Default, PartialEq)] struct SignedTransactionInfo { recipient: String, amount: f64, + #[serde(skip_serializing_if = "String::is_empty", default)] new_stake_account_address: String, finalized_date: Option>, signature: String, @@ -37,6 +39,7 @@ impl Default for TransactionInfo { finalized_date: None, transaction, last_valid_slot: 0, + lockup_date: None, } } } @@ -106,6 +109,7 @@ pub fn set_transaction_info( new_stake_account_address: Option<&Pubkey>, finalized: bool, last_valid_slot: Slot, + lockup_date: Option>, ) -> Result<(), Error> { let finalized_date = if finalized { Some(Utc::now()) } else { None }; let transaction_info = TransactionInfo { @@ -115,6 +119,7 @@ pub fn set_transaction_info( finalized_date, transaction: transaction.clone(), last_valid_slot, + lockup_date, }; let signature = transaction.signatures[0]; db.set(&signature.to_string(), &transaction_info)?; @@ -174,6 +179,33 @@ pub fn update_finalized_transaction( Ok(None) } +use csv::{ReaderBuilder, Trim}; +pub(crate) fn check_output_file(path: &str, db: &PickleDb) { + let mut rdr = ReaderBuilder::new() + .trim(Trim::All) + .from_path(path) + .unwrap(); + let logged_infos: Vec = + rdr.deserialize().map(|entry| entry.unwrap()).collect(); + + let mut transaction_infos = read_transaction_infos(db); + transaction_infos.sort_by(compare_transaction_infos); + let transaction_infos: Vec = transaction_infos + .iter() + .map(|info| SignedTransactionInfo { + recipient: info.recipient.to_string(), + amount: info.amount, + new_stake_account_address: info + .new_stake_account_address + .map(|x| x.to_string()) + .unwrap_or_else(|| "".to_string()), + finalized_date: info.finalized_date, + signature: info.transaction.signatures[0].to_string(), + }) + .collect(); + assert_eq!(logged_infos, transaction_infos); +} + #[cfg(test)] mod tests { use super::*; diff --git a/tokens/src/main.rs b/tokens/src/main.rs index 3ea9db57dd..6a9c4984fd 100644 --- a/tokens/src/main.rs +++ b/tokens/src/main.rs @@ -1,19 +1,14 @@ use solana_cli_config::Config; use solana_cli_config::CONFIG_FILE; use solana_client::rpc_client::RpcClient; -use solana_tokens::{ - arg_parser::parse_args, - args::{resolve_command, Command}, - commands, - thin_client::ThinClient, -}; +use solana_tokens::{arg_parser::parse_args, args::Command, commands, thin_client::ThinClient}; use std::env; use std::error::Error; use std::path::Path; use std::process; fn main() -> Result<(), Box> { - let command_args = parse_args(env::args_os()); + let command_args = parse_args(env::args_os())?; let config = if Path::new(&command_args.config_file).exists() { Config::load(&command_args.config_file)? } else { @@ -27,10 +22,10 @@ fn main() -> Result<(), Box> { let json_rpc_url = command_args.url.unwrap_or(config.json_rpc_url); let client = RpcClient::new(json_rpc_url); - match resolve_command(command_args.command)? { + match command_args.command { Command::DistributeTokens(args) => { let thin_client = ThinClient::new(client, args.dry_run); - commands::process_distribute_tokens(&thin_client, &args)?; + commands::process_allocations(&thin_client, &args)?; } Command::Balances(args) => { let thin_client = ThinClient::new(client, false); diff --git a/tokens/src/thin_client.rs b/tokens/src/thin_client.rs index bb87cce926..03403825ab 100644 --- a/tokens/src/thin_client.rs +++ b/tokens/src/thin_client.rs @@ -115,14 +115,17 @@ impl Client for BankClient { } } -pub struct ThinClient { - client: C, +pub struct ThinClient<'a> { + client: Box, dry_run: bool, } -impl ThinClient { - pub fn new(client: C, dry_run: bool) -> Self { - Self { client, dry_run } +impl<'a> ThinClient<'a> { + pub fn new(client: C, dry_run: bool) -> Self { + Self { + client: Box::new(client), + dry_run, + } } pub fn send_transaction(&self, transaction: Transaction) -> Result { diff --git a/tokens/tests/commands.rs b/tokens/tests/commands.rs index 3bdab39c44..c58302f0a7 100644 --- a/tokens/tests/commands.rs +++ b/tokens/tests/commands.rs @@ -11,7 +11,7 @@ fn test_process_distribute_with_rpc_client() { ..TestValidatorOptions::default() }); let rpc_client = RpcClient::new_socket(validator.leader_data.rpc); - test_process_distribute_tokens_with_client(rpc_client, validator.alice); + test_process_distribute_tokens_with_client(rpc_client, validator.alice, None); validator.server.close().unwrap(); remove_dir_all(validator.ledger_path).unwrap();