Add lockups via solana-tokens (#11782)

* Allow stake distributions to update lockups

* Reorg

* Add lockup test

* Fix clippy warning

Cargo.lock

@@ -4437,6 +4437,7 @@ dependencies = [
 name = "solana-tokens"
 version = "1.4.0"
 dependencies = [
+ "bincode",
  "chrono",
  "clap",
  "console",

tokens/Cargo.toml

@@ -32,6 +32,7 @@ tokio = "0.2"
 url = "2.1"
 
 [dev-dependencies]
+bincode = "1.3.1"
 solana-banks-server = { path = "../banks-server", version = "1.4.0" }
 solana-core = { path = "../core", version = "1.4.0" }
 solana-runtime = { path = "../runtime", version = "1.4.0" }

tokens/src/arg_parser.rs

@@ -158,6 +158,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")
@@ -310,11 +318,23 @@ fn parse_distribute_stake_args(
         &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,
+    };
+
     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),

tokens/src/args.rs

@@ -16,6 +16,7 @@ pub struct StakeArgs {
     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 BalancesArgs {

tokens/src/commands.rs

@@ -1,5 +1,6 @@
 use crate::args::{BalancesArgs, DistributeTokensArgs, StakeArgs, TransactionLogArgs};
 use crate::db::{self, TransactionInfo};
+use chrono::prelude::*;
 use console::style;
 use csv::{ReaderBuilder, Trim};
 use indexmap::IndexMap;
@@ -9,6 +10,7 @@ use serde::{Deserialize, Serialize};
 use solana_banks_client::{BanksClient, BanksClientExt};
 use solana_sdk::{
     commitment_config::CommitmentLevel,
+    instruction::Instruction,
     message::Message,
     native_token::{lamports_to_sol, sol_to_lamports},
     signature::{unique_signers, Signature, Signer},
@@ -17,7 +19,7 @@ use solana_sdk::{
     transport::{self, TransportError},
 };
 use solana_stake_program::{
-    stake_instruction,
+    stake_instruction::{self, LockupArgs},
     stake_state::{Authorized, Lockup, StakeAuthorize},
 };
 use std::{
@@ -37,6 +39,7 @@ struct Bid {
 struct Allocation {
     recipient: String,
     amount: f64,
+    lockup_date: String,
 }
 
 #[derive(thiserror::Error, Debug)]
@@ -49,6 +52,8 @@ pub enum Error {
     PickleDbError(#[from] pickledb::error::Error),
     #[error("Transport error")]
     TransportError(#[from] TransportError),
+    #[error("Missing lockup authority")]
+    MissingLockupAuthority,
 }
 
 fn merge_allocations(allocations: &[Allocation]) -> Vec<Allocation> {
@@ -59,12 +64,19 @@ fn merge_allocations(allocations: &[Allocation]) -> Vec<Allocation> {
             .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<Allocation>,
     transaction_infos: &[TransactionInfo],
@@ -72,7 +84,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 {
@@ -91,6 +103,7 @@ fn create_allocation(bid: &Bid, dollars_per_sol: f64) -> Allocation {
     Allocation {
         recipient: bid.primary_address.clone(),
         amount: bid.accepted_amount_dollars / dollars_per_sol,
+        lockup_date: "".to_string(),
     }
 }
 
@@ -111,7 +124,80 @@ async fn transfer<S: Signer>(
     ))
 }
 
-async fn distribute_tokens(
+fn distribution_instructions(
+    allocation: &Allocation,
+    new_stake_account_address: &Pubkey,
+    args: &DistributeTokensArgs,
+    lockup_date: Option<DateTime<Utc>>,
+) -> Vec<Instruction> {
+    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
+}
+
+async fn distribute_allocations(
     client: &mut BanksClient,
     db: &mut PickleDb,
     allocations: &[Allocation],
@@ -126,56 +212,25 @@ async 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);
 
-        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::<DateTime<Utc>>().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));
         let result: transport::Result<(Transaction, u64)> = {
@@ -198,6 +253,7 @@ async fn distribute_tokens(
                     Some(&new_stake_account_address),
                     false,
                     last_valid_slot,
+                    lockup_date,
                 )?;
             }
             Err(e) => {
@@ -235,7 +291,7 @@ fn new_spinner_progress_bar() -> ProgressBar {
     progress_bar
 }
 
-pub async fn process_distribute_tokens(
+pub async fn process_allocations(
     client: &mut BanksClient,
     args: &DistributeTokensArgs,
 ) -> Result<Option<usize>, Error> {
@@ -313,7 +369,7 @@ pub async fn process_distribute_tokens(
         );
     }
 
-    distribute_tokens(client, &mut db, &allocations, args).await?;
+    distribute_allocations(client, &mut db, &allocations, args).await?;
 
     let opt_confirmations = finalize_transactions(client, &mut db, args.dry_run).await?;
     Ok(opt_confirmations)
@@ -469,6 +525,7 @@ pub async fn test_process_distribute_tokens_with_client(
     let allocation = Allocation {
         recipient: alice_pubkey.to_string(),
         amount: 1000.0,
+        lockup_date: "".to_string(),
     };
     let allocations_file = NamedTempFile::new().unwrap();
     let input_csv = allocations_file.path().to_str().unwrap().to_string();
@@ -494,7 +551,7 @@ pub async fn test_process_distribute_tokens_with_client(
         dollars_per_sol: None,
         stake_args: None,
     };
-    let confirmations = process_distribute_tokens(client, &args).await.unwrap();
+    let confirmations = process_allocations(client, &args).await.unwrap();
     assert_eq!(confirmations, None);
 
     let transaction_infos =
@@ -513,7 +570,7 @@ pub async fn test_process_distribute_tokens_with_client(
     );
 
     // Now, run it again, and check there's no double-spend.
-    process_distribute_tokens(client, &args).await.unwrap();
+    process_allocations(client, &args).await.unwrap();
     let transaction_infos =
         db::read_transaction_infos(&db::open_db(&transaction_db, true).unwrap());
     assert_eq!(transaction_infos.len(), 1);
@@ -578,6 +635,7 @@ pub async fn test_process_distribute_stake_with_client(
     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();
@@ -597,6 +655,7 @@ pub async fn test_process_distribute_stake_with_client(
         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 {
@@ -609,7 +668,7 @@ pub async fn test_process_distribute_stake_with_client(
         sender_keypair: Box::new(sender_keypair),
         dollars_per_sol: None,
     };
-    let confirmations = process_distribute_tokens(client, &args).await.unwrap();
+    let confirmations = process_allocations(client, &args).await.unwrap();
     assert_eq!(confirmations, None);
 
     let transaction_infos =
@@ -633,7 +692,7 @@ pub async fn test_process_distribute_stake_with_client(
     );
 
     // Now, run it again, and check there's no double-spend.
-    process_distribute_tokens(client, &args).await.unwrap();
+    process_allocations(client, &args).await.unwrap();
     let transaction_infos =
         db::read_transaction_infos(&db::open_db(&transaction_db, true).unwrap());
     assert_eq!(transaction_infos.len(), 1);
@@ -661,11 +720,12 @@ mod tests {
     use solana_banks_server::banks_server::start_local_server;
     use solana_runtime::{bank::Bank, bank_forks::BankForks};
     use solana_sdk::genesis_config::create_genesis_config;
+    use solana_stake_program::stake_instruction::StakeInstruction;
     use std::sync::{Arc, RwLock};
     use tokio::runtime::Runtime;
 
     #[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_forks = Arc::new(RwLock::new(BankForks::new(Bank::new(&genesis_config))));
         Runtime::new().unwrap().block_on(async {
@@ -676,7 +736,7 @@ mod tests {
     }
 
     #[test]
-    fn test_process_distribute_stake() {
+    fn test_process_stake_allocations() {
         let (genesis_config, sender_keypair) = create_genesis_config(sol_to_lamports(9_000_000.0));
         let bank_forks = Arc::new(RwLock::new(BankForks::new(Bank::new(&genesis_config))));
         Runtime::new().unwrap().block_on(async {
@@ -692,6 +752,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();
@@ -718,6 +779,7 @@ mod tests {
         let allocation = Allocation {
             recipient: bid.primary_address,
             amount: 84.0,
+            lockup_date: "".to_string(),
         };
         assert_eq!(
             read_allocations(&input_csv, true, Some(0.5)),
@@ -733,10 +795,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 {
@@ -751,4 +815,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(),
+            stake_args: Some(stake_args),
+            from_bids: false,
+            sender_keypair: Box::new(Keypair::new()),
+            dollars_per_sol: 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");
+        }
+    }
 }

tokens/src/db.rs

@@ -13,6 +13,7 @@ pub struct TransactionInfo {
     pub finalized_date: Option<DateTime<Utc>>,
     pub transaction: Transaction,
     pub last_valid_slot: Slot,
+    pub lockup_date: Option<DateTime<Utc>>,
 }
 
 #[derive(Serialize, Deserialize, Debug, Default, PartialEq)]
@@ -37,6 +38,7 @@ impl Default for TransactionInfo {
             finalized_date: None,
             transaction,
             last_valid_slot: 0,
+            lockup_date: None,
         }
     }
 }
@@ -106,6 +108,7 @@ pub fn set_transaction_info(
     new_stake_account_address: Option<&Pubkey>,
     finalized: bool,
     last_valid_slot: Slot,
+    lockup_date: Option<DateTime<Utc>>,
 ) -> Result<(), Error> {
     let finalized_date = if finalized { Some(Utc::now()) } else { None };
     let transaction_info = TransactionInfo {
@@ -115,6 +118,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)?;

tokens/src/main.rs

@@ -27,10 +27,7 @@ fn main() -> Result<(), Box<dyn Error>> {
 
     match command_args.command {
         Command::DistributeTokens(args) => {
-            runtime.block_on(commands::process_distribute_tokens(
-                &mut banks_client,
-                &args,
-            ))?;
+            runtime.block_on(commands::process_allocations(&mut banks_client, &args))?;
         }
         Command::Balances(args) => {
             runtime.block_on(commands::process_balances(&mut banks_client, &args))?;