Add create-stake command to solana-tokens CLI (#17550)

* Add create-stake command to solana-tokens CLI

* Add --unlocked-sol arg

Thanks @CriesofCarrots!

tokens/src/arg_parser.rs

@@ -1,5 +1,6 @@
 use crate::args::{
-    Args, BalancesArgs, Command, DistributeTokensArgs, SplTokenArgs, StakeArgs, TransactionLogArgs,
+    Args, BalancesArgs, Command, DistributeTokensArgs, SenderStakeArgs, SplTokenArgs, StakeArgs,
+    TransactionLogArgs,
 };
 use clap::{
     crate_description, crate_name, value_t, value_t_or_exit, App, Arg, ArgMatches, SubCommand,
@@ -102,9 +103,80 @@ where
                         .help("Fee payer"),
                 ),
         )
+        .subcommand(
+            SubCommand::with_name("create-stake")
+                .about("Create stake accounts")
+                .arg(
+                    Arg::with_name("db_path")
+                        .long("db-path")
+                        .required(true)
+                        .takes_value(true)
+                        .value_name("FILE")
+                        .help(
+                            "Location for storing distribution database. \
+                            The database is used for tracking transactions as they are finalized \
+                            and preventing double spends.",
+                        ),
+                )
+                .arg(
+                    Arg::with_name("input_csv")
+                        .long("input-csv")
+                        .required(true)
+                        .takes_value(true)
+                        .value_name("FILE")
+                        .help("Allocations CSV file"),
+                )
+                .arg(
+                    Arg::with_name("dry_run")
+                        .long("dry-run")
+                        .help("Do not execute any transfers"),
+                )
+                .arg(
+                    Arg::with_name("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")
+                        .required(true)
+                        .takes_value(true)
+                        .value_name("SENDING_KEYPAIR")
+                        .validator(is_valid_signer)
+                        .help("Keypair to fund accounts"),
+                )
+                .arg(
+                    Arg::with_name("unlocked_sol")
+                        .default_value("1.0")
+                        .long("unlocked-sol")
+                        .takes_value(true)
+                        .value_name("SOL_AMOUNT")
+                        .help("Amount of SOL to put in system account to pay for fees"),
+                )
+                .arg(
+                    Arg::with_name("lockup_authority")
+                        .long("lockup-authority")
+                        .takes_value(true)
+                        .value_name("PUBKEY")
+                        .validator(is_valid_pubkey)
+                        .help("Lockup Authority Address"),
+                )
+                .arg(
+                    Arg::with_name("fee_payer")
+                        .long("fee-payer")
+                        .required(true)
+                        .takes_value(true)
+                        .value_name("KEYPAIR")
+                        .validator(is_valid_signer)
+                        .help("Fee payer"),
+                ),
+        )
         .subcommand(
             SubCommand::with_name("distribute-stake")
-                .about("Distribute stake accounts")
+                .about("Split to stake accounts")
                 .arg(
                     Arg::with_name("db_path")
                         .long("db-path")
@@ -363,6 +435,58 @@ fn parse_distribute_tokens_args(
     })
 }
 
+fn parse_create_stake_args(
+    matches: &ArgMatches<'_>,
+) -> Result<DistributeTokensArgs, Box<dyn Error>> {
+    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 lockup_authority_str = value_t!(matches, "lockup_authority", String).ok();
+    let lockup_authority = lockup_authority_str
+        .map(|path| {
+            pubkey_from_path(
+                &signer_matches,
+                &path,
+                "lockup authority",
+                &mut wallet_manager,
+            )
+        })
+        .transpose()?;
+
+    let stake_args = StakeArgs {
+        unlocked_sol: sol_to_lamports(value_t_or_exit!(matches, "unlocked_sol", f64)),
+        lockup_authority,
+        sender_stake_args: None,
+    };
+    Ok(DistributeTokensArgs {
+        input_csv: value_t_or_exit!(matches, "input_csv", String),
+        transaction_db: value_t_or_exit!(matches, "db_path", String),
+        output_path: matches.value_of("output_path").map(|path| path.to_string()),
+        dry_run: matches.is_present("dry_run"),
+        sender_keypair,
+        fee_payer,
+        stake_args: Some(stake_args),
+        spl_token_args: None,
+        transfer_amount: None,
+    })
+}
+
 fn parse_distribute_stake_args(
     matches: &ArgMatches<'_>,
 ) -> Result<DistributeTokensArgs, Box<dyn Error>> {
@@ -421,13 +545,18 @@ fn parse_distribute_stake_args(
         })
         .transpose()?;
 
-    let stake_args = StakeArgs {
+    let lockup_authority_address = lockup_authority.as_ref().map(|keypair| keypair.pubkey());
+    let sender_stake_args = SenderStakeArgs {
         stake_account_address,
-        unlocked_sol: sol_to_lamports(value_t_or_exit!(matches, "unlocked_sol", f64)),
         stake_authority,
         withdraw_authority,
         lockup_authority,
     };
+    let stake_args = StakeArgs {
+        unlocked_sol: sol_to_lamports(value_t_or_exit!(matches, "unlocked_sol", f64)),
+        lockup_authority: lockup_authority_address,
+        sender_stake_args: Some(sender_stake_args),
+    };
     Ok(DistributeTokensArgs {
         input_csv: value_t_or_exit!(matches, "input_csv", String),
         transaction_db: value_t_or_exit!(matches, "db_path", String),
@@ -520,6 +649,9 @@ where
         ("distribute-tokens", Some(matches)) => {
             Command::DistributeTokens(parse_distribute_tokens_args(matches)?)
         }
+        ("create-stake", Some(matches)) => {
+            Command::DistributeTokens(parse_create_stake_args(matches)?)
+        }
         ("distribute-stake", Some(matches)) => {
             Command::DistributeTokens(parse_distribute_stake_args(matches)?)
         }

tokens/src/args.rs

@@ -1,5 +1,18 @@
 use solana_sdk::{pubkey::Pubkey, signature::Signer};
 
+pub struct SenderStakeArgs {
+    pub stake_account_address: Pubkey,
+    pub stake_authority: Box<dyn Signer>,
+    pub withdraw_authority: Box<dyn Signer>,
+    pub lockup_authority: Option<Box<dyn Signer>>,
+}
+
+pub struct StakeArgs {
+    pub unlocked_sol: u64,
+    pub lockup_authority: Option<Pubkey>,
+    pub sender_stake_args: Option<SenderStakeArgs>,
+}
+
 pub struct DistributeTokensArgs {
     pub input_csv: String,
     pub transaction_db: String,
@@ -12,14 +25,6 @@ pub struct DistributeTokensArgs {
     pub transfer_amount: Option<u64>,
 }
 
-pub struct StakeArgs {
-    pub unlocked_sol: u64,
-    pub stake_account_address: Pubkey,
-    pub stake_authority: Box<dyn Signer>,
-    pub withdraw_authority: Box<dyn Signer>,
-    pub lockup_authority: Option<Box<dyn Signer>>,
-}
-
 #[derive(Default)]
 pub struct SplTokenArgs {
     pub token_account_address: Pubkey,

tokens/src/commands.rs

@@ -1,5 +1,5 @@
 use crate::{
-    args::{BalancesArgs, DistributeTokensArgs, StakeArgs, TransactionLogArgs},
+    args::{BalancesArgs, DistributeTokensArgs, SenderStakeArgs, StakeArgs, TransactionLogArgs},
     db::{self, TransactionInfo},
     spl_token::*,
     token_display::Token,
@@ -178,33 +178,62 @@ fn distribution_instructions(
     lockup_date: Option<DateTime<Utc>>,
     do_create_associated_token_account: bool,
 ) -> Vec<Instruction> {
-    if args.stake_args.is_none() && args.spl_token_args.is_none() {
-        let from = args.sender_keypair.pubkey();
-        let to = allocation.recipient.parse().unwrap();
-        let lamports = allocation.amount;
-        let instruction = system_instruction::transfer(&from, &to, lamports);
-        return vec![instruction];
-    }
-
     if args.spl_token_args.is_some() {
         return build_spl_token_instructions(allocation, args, do_create_associated_token_account);
     }
 
-    let stake_args = args.stake_args.as_ref().unwrap();
+    match &args.stake_args {
+        // No stake args; a simple token transfer.
+        None => {
+            let from = args.sender_keypair.pubkey();
+            let to = allocation.recipient.parse().unwrap();
+            let lamports = allocation.amount;
+            let instruction = system_instruction::transfer(&from, &to, lamports);
+            vec![instruction]
+        }
+
+        // Stake args provided, so create a recipient stake account.
+        Some(stake_args) => {
             let unlocked_sol = stake_args.unlocked_sol;
             let sender_pubkey = args.sender_keypair.pubkey();
-    let stake_authority = stake_args.stake_authority.pubkey();
+            let recipient = allocation.recipient.parse().unwrap();
-    let withdraw_authority = stake_args.withdraw_authority.pubkey();
 
+            let mut instructions = match &stake_args.sender_stake_args {
+                // No source stake account, so create a recipient stake account directly.
+                None => {
+                    // Make the recipient both the new stake and withdraw authority
+                    let authorized = Authorized {
+                        staker: recipient,
+                        withdrawer: recipient,
+                    };
+                    let mut lockup = Lockup::default();
+                    if let Some(lockup_date) = lockup_date {
+                        lockup.unix_timestamp = lockup_date.timestamp();
+                    }
+                    if let Some(lockup_authority) = stake_args.lockup_authority {
+                        lockup.custodian = lockup_authority;
+                    }
+                    stake_instruction::create_account(
+                        &sender_pubkey,
+                        &new_stake_account_address,
+                        &authorized,
+                        &lockup,
+                        allocation.amount - unlocked_sol,
+                    )
+                }
+
+                // A sender stake account was provided, so create a recipient stake account by
+                // splitting the sender account.
+                Some(sender_stake_args) => {
+                    let stake_authority = sender_stake_args.stake_authority.pubkey();
+                    let withdraw_authority = sender_stake_args.withdraw_authority.pubkey();
                     let mut instructions = stake_instruction::split(
-        &stake_args.stake_account_address,
+                        &sender_stake_args.stake_account_address,
                         &stake_authority,
                         allocation.amount - unlocked_sol,
                         &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,
@@ -225,11 +254,6 @@ fn distribution_instructions(
 
                     // 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,
@@ -238,10 +262,15 @@ fn distribution_instructions(
                         instructions.push(stake_instruction::set_lockup(
                             &new_stake_account_address,
                             &lockup,
-            &lockup_authority,
+                            &stake_args.lockup_authority.unwrap(),
                         ));
                     }
 
+                    instructions
+                }
+            };
+
+            // Transfer some unlocked tokens to recipient, which they can use for transaction fees.
             instructions.push(system_instruction::transfer(
                 &sender_pubkey,
                 &recipient,
@@ -249,6 +278,8 @@ fn distribution_instructions(
             ));
 
             instructions
+        }
+    }
 }
 
 fn build_messages(
@@ -336,17 +367,20 @@ fn send_messages(
 
         let mut signers = vec![&*args.fee_payer, &*args.sender_keypair];
         if let Some(stake_args) = &args.stake_args {
-            signers.push(&*stake_args.stake_authority);
+            signers.push(&new_stake_account_keypair);
-            signers.push(&*stake_args.withdraw_authority);
+            if let Some(sender_stake_args) = &stake_args.sender_stake_args {
+                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 let Some(lockup_authority) = &stake_args.lockup_authority {
+                    if let Some(lockup_authority) = &sender_stake_args.lockup_authority {
                         signers.push(&**lockup_authority);
                     } else {
                         return Err(Error::MissingLockupAuthority);
                     }
                 }
             }
+        }
         let signers = unique_signers(signers);
         let result: ClientResult<(Transaction, u64)> = {
             if args.dry_run {
@@ -366,12 +400,14 @@ fn send_messages(
         };
         match result {
             Ok((transaction, last_valid_slot)) => {
+                let new_stake_account_address_option =
+                    args.stake_args.as_ref().map(|_| &new_stake_account_address);
                 db::set_transaction_info(
                     db,
                     &allocation.recipient.parse().unwrap(),
                     allocation.amount,
                     &transaction,
-                    args.stake_args.as_ref().map(|_| &new_stake_account_address),
+                    new_stake_account_address_option,
                     false,
                     last_valid_slot,
                     lockup_date,
@@ -708,8 +744,13 @@ fn check_payer_balances(
     let (distribution_source, unlocked_sol_source) = if let Some(stake_args) = &args.stake_args {
         let total_unlocked_sol = allocations.len() as u64 * stake_args.unlocked_sol;
         undistributed_tokens -= total_unlocked_sol;
+        let from_pubkey = if let Some(sender_stake_args) = &stake_args.sender_stake_args {
+            sender_stake_args.stake_account_address
+        } else {
+            args.sender_keypair.pubkey()
+        };
         (
-            stake_args.stake_account_address,
+            from_pubkey,
             Some((args.sender_keypair.pubkey(), total_unlocked_sol)),
         )
     } else {
@@ -909,7 +950,7 @@ pub fn test_process_distribute_tokens_with_client(
     check_output_file(&output_path, &db::open_db(&transaction_db, true).unwrap());
 }
 
-pub fn test_process_distribute_stake_with_client(client: &RpcClient, sender_keypair: Keypair) {
+pub fn test_process_create_stake_with_client(client: &RpcClient, sender_keypair: Keypair) {
     let exit = Arc::new(AtomicBool::default());
     let fee_payer = Keypair::new();
     let transaction = transfer(
@@ -975,11 +1016,137 @@ pub fn test_process_distribute_stake_with_client(client: &RpcClient, sender_keyp
     let output_path = output_file.path().to_str().unwrap().to_string();
 
     let stake_args = StakeArgs {
+        lockup_authority: None,
+        unlocked_sol: sol_to_lamports(1.0),
+        sender_stake_args: None,
+    };
+    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),
+        spl_token_args: None,
+        sender_keypair: Box::new(sender_keypair),
+        transfer_amount: None,
+    };
+    let confirmations = process_allocations(client, &args, exit.clone()).unwrap();
+    assert_eq!(confirmations, None);
+
+    let transaction_infos =
+        db::read_transaction_infos(&db::open_db(&transaction_db, true).unwrap());
+    assert_eq!(transaction_infos.len(), 1);
+    assert_eq!(transaction_infos[0].recipient, alice_pubkey);
+    assert_eq!(transaction_infos[0].amount, expected_amount);
+
+    assert_eq!(
+        client.get_balance(&alice_pubkey).unwrap(),
+        sol_to_lamports(1.0),
+    );
+    let new_stake_account_address = transaction_infos[0].new_stake_account_address.unwrap();
+    assert_eq!(
+        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());
+
+    // Now, run it again, and check there's no double-spend.
+    process_allocations(client, &args, exit).unwrap();
+    let transaction_infos =
+        db::read_transaction_infos(&db::open_db(&transaction_db, true).unwrap());
+    assert_eq!(transaction_infos.len(), 1);
+    assert_eq!(transaction_infos[0].recipient, alice_pubkey);
+    assert_eq!(transaction_infos[0].amount, expected_amount);
+
+    assert_eq!(
+        client.get_balance(&alice_pubkey).unwrap(),
+        sol_to_lamports(1.0),
+    );
+    assert_eq!(
+        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());
+}
+
+pub fn test_process_distribute_stake_with_client(client: &RpcClient, sender_keypair: Keypair) {
+    let exit = Arc::new(AtomicBool::default());
+    let fee_payer = Keypair::new();
+    let transaction = transfer(
+        client,
+        sol_to_lamports(1.0),
+        &sender_keypair,
+        &fee_payer.pubkey(),
+    )
+    .unwrap();
+    client
+        .send_and_confirm_transaction_with_spinner(&transaction)
+        .unwrap();
+
+    let stake_account_keypair = Keypair::new();
+    let stake_account_address = stake_account_keypair.pubkey();
+    let stake_authority = Keypair::new();
+    let withdraw_authority = Keypair::new();
+
+    let authorized = Authorized {
+        staker: stake_authority.pubkey(),
+        withdrawer: withdraw_authority.pubkey(),
+    };
+    let lockup = Lockup::default();
+    let instructions = stake_instruction::create_account(
+        &sender_keypair.pubkey(),
+        &stake_account_address,
+        &authorized,
+        &lockup,
+        sol_to_lamports(3000.0),
+    );
+    let message = Message::new(&instructions, Some(&sender_keypair.pubkey()));
+    let signers = [&sender_keypair, &stake_account_keypair];
+    let (blockhash, _fees) = client.get_recent_blockhash().unwrap();
+    let transaction = Transaction::new(&signers, message, blockhash);
+    client
+        .send_and_confirm_transaction_with_spinner(&transaction)
+        .unwrap();
+
+    let expected_amount = sol_to_lamports(1000.0);
+    let alice_pubkey = solana_sdk::pubkey::new_rand();
+    let file = NamedTempFile::new().unwrap();
+    let input_csv = file.path().to_str().unwrap().to_string();
+    let mut wtr = csv::WriterBuilder::new().from_writer(file);
+    wtr.write_record(&["recipient", "amount", "lockup_date"])
+        .unwrap();
+    wtr.write_record(&[
+        alice_pubkey.to_string(),
+        lamports_to_sol(expected_amount).to_string(),
+        "".to_string(),
+    ])
+    .unwrap();
+    wtr.flush().unwrap();
+
+    let dir = tempdir().unwrap();
+    let transaction_db = dir
+        .path()
+        .join("transactions.db")
+        .to_str()
+        .unwrap()
+        .to_string();
+
+    let output_file = NamedTempFile::new().unwrap();
+    let output_path = output_file.path().to_str().unwrap().to_string();
+
+    let sender_stake_args = SenderStakeArgs {
         stake_account_address,
         stake_authority: Box::new(stake_authority),
         withdraw_authority: Box::new(withdraw_authority),
         lockup_authority: None,
+    };
+    let stake_args = StakeArgs {
         unlocked_sol: sol_to_lamports(1.0),
+        lockup_authority: None,
+        sender_stake_args: Some(sender_stake_args),
     };
     let args = DistributeTokensArgs {
         fee_payer: Box::new(fee_payer),
@@ -1061,6 +1228,16 @@ mod tests {
         test_process_distribute_tokens_with_client(&client, alice, Some(sol_to_lamports(1.5)));
     }
 
+    #[test]
+    fn test_create_stake_allocations() {
+        let alice = Keypair::new();
+        let test_validator = TestValidator::with_no_fees(alice.pubkey(), None);
+        let url = test_validator.rpc_url();
+
+        let client = RpcClient::new_with_commitment(url, CommitmentConfig::processed());
+        test_process_create_stake_with_client(&client, alice);
+    }
+
     #[test]
     fn test_process_stake_allocations() {
         let alice = Keypair::new();
@@ -1293,7 +1470,7 @@ mod tests {
     const SET_LOCKUP_INDEX: usize = 4;
 
     #[test]
-    fn test_set_stake_lockup() {
+    fn test_set_split_stake_lockup() {
         let lockup_date_str = "2021-01-07T00:00:00Z";
         let allocation = Allocation {
             recipient: Pubkey::default().to_string(),
@@ -1303,12 +1480,17 @@ mod tests {
         let stake_account_address = solana_sdk::pubkey::new_rand();
         let new_stake_account_address = solana_sdk::pubkey::new_rand();
         let lockup_authority = Keypair::new();
-        let stake_args = StakeArgs {
+        let lockup_authority_address = lockup_authority.pubkey();
+        let sender_stake_args = SenderStakeArgs {
             stake_account_address,
             stake_authority: Box::new(Keypair::new()),
             withdraw_authority: Box::new(Keypair::new()),
             lockup_authority: Some(Box::new(lockup_authority)),
+        };
+        let stake_args = StakeArgs {
+            lockup_authority: Some(lockup_authority_address),
             unlocked_sol: sol_to_lamports(1.0),
+            sender_stake_args: Some(sender_stake_args),
         };
         let args = DistributeTokensArgs {
             fee_payer: Box::new(Keypair::new()),
@@ -1558,12 +1740,17 @@ mod tests {
             .send_and_confirm_transaction_with_spinner(&transaction)
             .unwrap();
 
-        StakeArgs {
+        let sender_stake_args = SenderStakeArgs {
             stake_account_address,
             stake_authority: Box::new(stake_authority),
             withdraw_authority: Box::new(withdraw_authority),
             lockup_authority: None,
+        };
+
+        StakeArgs {
+            lockup_authority: None,
             unlocked_sol,
+            sender_stake_args: Some(sender_stake_args),
         }
     }