Allow stake lockup fields to be updated independently (#8568)

* Make Lockup fields optional for SetLockup instruction

* Use LockupArgs in cli

* Include lockup timestamp in stake-account print

cli/src/cli.rs

@@ -40,7 +40,10 @@ use solana_sdk::{
     system_instruction::{self, create_address_with_seed, SystemError, MAX_ADDRESS_SEED_LEN},
     transaction::{Transaction, TransactionError},
 };
-use solana_stake_program::stake_state::{Lockup, StakeAuthorize};
+use solana_stake_program::{
+    stake_instruction::LockupArgs,
+    stake_state::{Lockup, StakeAuthorize},
+};
 use solana_storage_program::storage_instruction::StorageAccountType;
 use solana_vote_program::vote_state::VoteAuthorize;
 use std::{
@@ -302,7 +305,7 @@ pub enum CliCommand {
     },
     StakeSetLockup {
         stake_account_pubkey: Pubkey,
-        lockup: Lockup,
+        lockup: LockupArgs,
         custodian: SignerIndex,
         sign_only: bool,
         blockhash_query: BlockhashQuery,

cli/src/stake.rs

@@ -7,7 +7,8 @@ use crate::{
     nonce::{check_nonce_account, nonce_arg, NONCE_ARG, NONCE_AUTHORITY_ARG},
     offline::*,
 };
-use clap::{App, Arg, ArgMatches, SubCommand};
+use chrono::{DateTime, NaiveDateTime, SecondsFormat, Utc};
+use clap::{App, Arg, ArgGroup, ArgMatches, SubCommand};
 use console::style;
 use solana_clap_utils::{input_parsers::*, input_validators::*, offline::*, ArgConstant};
 use solana_client::rpc_client::RpcClient;
@@ -24,7 +25,7 @@ use solana_sdk::{
     transaction::Transaction,
 };
 use solana_stake_program::{
-    stake_instruction::{self, StakeError},
+    stake_instruction::{self, LockupArgs, StakeError},
     stake_state::{Authorized, Lockup, Meta, StakeAuthorize, StakeState},
 };
 use solana_vote_program::vote_state::VoteState;
@@ -364,6 +365,9 @@ impl StakeSubCommands for App<'_, '_> {
                         .validator(is_pubkey_or_keypair)
                         .help("Identity of the new lockup custodian (can withdraw before lockup expires)")
                 )
+                .group(ArgGroup::with_name("lockup_details")
+                    .args(&["lockup_epoch", "lockup_date", "new_custodian"])
+                    .required(true))
                 .arg(
                     Arg::with_name("custodian")
                         .long("custodian")
@@ -672,9 +676,9 @@ pub fn parse_stake_set_lockup(
     wallet_manager: Option<&Arc<RemoteWalletManager>>,
 ) -> Result<CliCommandInfo, CliError> {
     let stake_account_pubkey = pubkey_of(matches, "stake_account_pubkey").unwrap();
-    let epoch = value_of(matches, "lockup_epoch").unwrap_or(0);
+    let epoch = value_of(matches, "lockup_epoch");
-    let unix_timestamp = unix_timestamp_from_rfc3339_datetime(matches, "lockup_date").unwrap_or(0);
+    let unix_timestamp = unix_timestamp_from_rfc3339_datetime(matches, "lockup_date");
-    let new_custodian = pubkey_of(matches, "new_custodian").unwrap_or_default();
+    let new_custodian = pubkey_of(matches, "new_custodian");
 
     let sign_only = matches.is_present(SIGN_ONLY_ARG.name);
     let blockhash_query = BlockhashQuery::new_from_matches(matches);
@@ -695,7 +699,7 @@ pub fn parse_stake_set_lockup(
     Ok(CliCommandInfo {
         command: CliCommand::StakeSetLockup {
             stake_account_pubkey,
-            lockup: Lockup {
+            lockup: LockupArgs {
                 custodian: new_custodian,
                 epoch,
                 unix_timestamp,
@@ -1155,7 +1159,7 @@ pub fn process_stake_set_lockup(
     rpc_client: &RpcClient,
     config: &CliConfig,
     stake_account_pubkey: &Pubkey,
-    lockup: &mut Lockup,
+    lockup: &mut LockupArgs,
     custodian: SignerIndex,
     sign_only: bool,
     blockhash_query: &BlockhashQuery,
@@ -1166,10 +1170,7 @@ pub fn process_stake_set_lockup(
     let (recent_blockhash, fee_calculator) =
         blockhash_query.get_blockhash_fee_calculator(rpc_client)?;
     let custodian = config.signers[custodian];
-    // If new custodian is not explicitly set, default to current custodian
+
-    if lockup.custodian == Pubkey::default() {
-        lockup.custodian = custodian.pubkey();
-    }
     let ixs = vec![stake_instruction::set_lockup(
         stake_account_pubkey,
         lockup,
@@ -1215,6 +1216,12 @@ pub fn print_stake_state(stake_lamports: u64, stake_state: &StakeState, use_lamp
         println!("Authorized Withdrawer: {}", authorized.withdrawer);
     }
     fn show_lockup(lockup: &Lockup) {
+        println!(
+            "Lockup Timestamp: {} (UnixTimestamp: {})",
+            DateTime::<Utc>::from_utc(NaiveDateTime::from_timestamp(lockup.unix_timestamp, 0), Utc)
+                .to_rfc3339_opts(SecondsFormat::Secs, true),
+            lockup.unix_timestamp
+        );
         println!("Lockup Epoch: {}", lockup.epoch);
         println!("Lockup Custodian: {}", lockup.custodian);
     }

cli/tests/stake.rs

@@ -14,7 +14,10 @@ use solana_sdk::{
     signature::{keypair_from_seed, Keypair, Signer},
     system_instruction::create_address_with_seed,
 };
-use solana_stake_program::stake_state::{Lockup, StakeAuthorize, StakeState};
+use solana_stake_program::{
+    stake_instruction::LockupArgs,
+    stake_state::{Lockup, StakeAuthorize, StakeState},
+};
 use std::{fs::remove_dir_all, sync::mpsc::channel, thread::sleep, time::Duration};
 
 fn check_balance(expected_balance: u64, client: &RpcClient, pubkey: &Pubkey) {
@@ -1128,10 +1131,10 @@ fn test_stake_set_lockup() {
     );
 
     // Online set lockup
-    let mut lockup = Lockup {
+    let lockup = LockupArgs {
-        unix_timestamp: 1581534570,
+        unix_timestamp: Some(1581534570),
-        epoch: 200,
+        epoch: Some(200),
-        custodian: Pubkey::default(),
+        custodian: None,
     };
     config.signers.pop();
     config.command = CliCommand::StakeSetLockup {
@@ -1151,17 +1154,21 @@ fn test_stake_set_lockup() {
         StakeState::Initialized(meta) => meta.lockup,
         _ => panic!("Unexpected stake state!"),
     };
-    lockup.custodian = config.signers[0].pubkey(); // Default new_custodian is config.signers[0]
+    assert_eq!(
-    assert_eq!(current_lockup, lockup);
+        current_lockup.unix_timestamp,
+        lockup.unix_timestamp.unwrap()
+    );
+    assert_eq!(current_lockup.epoch, lockup.epoch.unwrap());
+    assert_eq!(current_lockup.custodian, config.signers[0].pubkey());
 
     // Set custodian to another pubkey
     let online_custodian = Keypair::new();
     let online_custodian_pubkey = online_custodian.pubkey();
 
-    let lockup = Lockup {
+    let lockup = LockupArgs {
-        unix_timestamp: 1581534571,
+        unix_timestamp: Some(1581534571),
-        epoch: 201,
+        epoch: Some(201),
-        custodian: online_custodian_pubkey,
+        custodian: Some(online_custodian_pubkey),
     };
     config.command = CliCommand::StakeSetLockup {
         stake_account_pubkey,
@@ -1175,10 +1182,10 @@ fn test_stake_set_lockup() {
     };
     process_command(&config).unwrap();
 
-    let mut lockup = Lockup {
+    let lockup = LockupArgs {
-        unix_timestamp: 1581534572,
+        unix_timestamp: Some(1581534572),
-        epoch: 202,
+        epoch: Some(202),
-        custodian: Pubkey::default(),
+        custodian: None,
     };
     config.signers = vec![&default_signer, &online_custodian];
     config.command = CliCommand::StakeSetLockup {
@@ -1198,14 +1205,18 @@ fn test_stake_set_lockup() {
         StakeState::Initialized(meta) => meta.lockup,
         _ => panic!("Unexpected stake state!"),
     };
-    lockup.custodian = online_custodian_pubkey; // Default new_custodian is designated custodian
+    assert_eq!(
-    assert_eq!(current_lockup, lockup);
+        current_lockup.unix_timestamp,
+        lockup.unix_timestamp.unwrap()
+    );
+    assert_eq!(current_lockup.epoch, lockup.epoch.unwrap());
+    assert_eq!(current_lockup.custodian, online_custodian_pubkey);
 
     // Set custodian to offline pubkey
-    let lockup = Lockup {
+    let lockup = LockupArgs {
-        unix_timestamp: 1581534573,
+        unix_timestamp: Some(1581534573),
-        epoch: 203,
+        epoch: Some(203),
-        custodian: offline_pubkey,
+        custodian: Some(offline_pubkey),
     };
     config.command = CliCommand::StakeSetLockup {
         stake_account_pubkey,
@@ -1244,10 +1255,10 @@ fn test_stake_set_lockup() {
     };
 
     // Nonced offline set lockup
-    let lockup = Lockup {
+    let lockup = LockupArgs {
-        unix_timestamp: 1581534576,
+        unix_timestamp: Some(1581534576),
-        epoch: 222,
+        epoch: Some(222),
-        custodian: offline_pubkey,
+        custodian: None,
     };
     config_offline.command = CliCommand::StakeSetLockup {
         stake_account_pubkey,
@@ -1280,7 +1291,12 @@ fn test_stake_set_lockup() {
         StakeState::Initialized(meta) => meta.lockup,
         _ => panic!("Unexpected stake state!"),
     };
-    assert_eq!(current_lockup, lockup);
+    assert_eq!(
+        current_lockup.unix_timestamp,
+        lockup.unix_timestamp.unwrap()
+    );
+    assert_eq!(current_lockup.epoch, lockup.epoch.unwrap());
+    assert_eq!(current_lockup.custodian, offline_pubkey);
 
     server.close().unwrap();
     remove_dir_all(ledger_path).unwrap();

programs/stake/src/stake_instruction.rs

@@ -7,6 +7,7 @@ use num_derive::{FromPrimitive, ToPrimitive};
 use serde_derive::{Deserialize, Serialize};
 use solana_sdk::{
     account::{get_signers, KeyedAccount},
+    clock::{Epoch, UnixTimestamp},
     instruction::{AccountMeta, Instruction, InstructionError, WithSigner},
     program_utils::{limited_deserialize, next_keyed_account, DecodeError},
     pubkey::Pubkey,
@@ -121,7 +122,14 @@ pub enum StakeInstruction {
     /// Expects 1 Account:
     ///    0 - initialized StakeAccount
     ///
-    SetLockup(Lockup),
+    SetLockup(LockupArgs),
+}
+
+#[derive(Default, Debug, Serialize, Deserialize, PartialEq, Clone, Copy)]
+pub struct LockupArgs {
+    pub unix_timestamp: Option<UnixTimestamp>,
+    pub epoch: Option<Epoch>,
+    pub custodian: Option<Pubkey>,
 }
 
 fn initialize(stake_pubkey: &Pubkey, authorized: &Authorized, lockup: &Lockup) -> Instruction {
@@ -361,7 +369,7 @@ pub fn deactivate_stake(stake_pubkey: &Pubkey, authorized_pubkey: &Pubkey) -> In
 
 pub fn set_lockup(
     stake_pubkey: &Pubkey,
-    lockup: &Lockup,
+    lockup: &LockupArgs,
     custodian_pubkey: &Pubkey,
 ) -> Instruction {
     let account_metas = vec![AccountMeta::new(*stake_pubkey, false)].with_signer(custodian_pubkey);
@@ -540,7 +548,7 @@ mod tests {
         assert_eq!(
             process_instruction(&set_lockup(
                 &Pubkey::default(),
-                &Lockup::default(),
+                &LockupArgs::default(),
                 &Pubkey::default()
             )),
             Err(InstructionError::InvalidAccountData),

programs/stake/src/stake_state.rs

@@ -3,7 +3,11 @@
 //! * keep track of rewards
 //! * own mining pools
 
-use crate::{config::Config, id, stake_instruction::StakeError};
+use crate::{
+    config::Config,
+    id,
+    stake_instruction::{LockupArgs, StakeError},
+};
 use serde_derive::{Deserialize, Serialize};
 use solana_sdk::{
     account::{Account, KeyedAccount},
@@ -120,13 +124,21 @@ pub struct Meta {
 impl Meta {
     pub fn set_lockup(
         &mut self,
-        lockup: &Lockup,
+        lockup: &LockupArgs,
         signers: &HashSet<Pubkey>,
     ) -> Result<(), InstructionError> {
         if !signers.contains(&self.lockup.custodian) {
             return Err(InstructionError::MissingRequiredSignature);
         }
-        self.lockup = *lockup;
+        if let Some(unix_timestamp) = lockup.unix_timestamp {
+            self.lockup.unix_timestamp = unix_timestamp;
+        }
+        if let Some(epoch) = lockup.epoch {
+            self.lockup.epoch = epoch;
+        }
+        if let Some(custodian) = lockup.custodian {
+            self.lockup.custodian = custodian;
+        }
         Ok(())
     }
 
@@ -524,7 +536,7 @@ pub trait StakeAccount {
     fn deactivate(&self, clock: &Clock, signers: &HashSet<Pubkey>) -> Result<(), InstructionError>;
     fn set_lockup(
         &self,
-        lockup: &Lockup,
+        lockup: &LockupArgs,
         signers: &HashSet<Pubkey>,
     ) -> Result<(), InstructionError>;
     fn split(
@@ -635,7 +647,7 @@ impl<'a> StakeAccount for KeyedAccount<'a> {
     }
     fn set_lockup(
         &self,
-        lockup: &Lockup,
+        lockup: &LockupArgs,
         signers: &HashSet<Pubkey>,
     ) -> Result<(), InstructionError> {
         match self.state()? {
@@ -1594,7 +1606,7 @@ mod tests {
         // wrong state, should fail
         let stake_keyed_account = KeyedAccount::new(&stake_pubkey, false, &stake_account);
         assert_eq!(
-            stake_keyed_account.set_lockup(&Lockup::default(), &HashSet::default(),),
+            stake_keyed_account.set_lockup(&LockupArgs::default(), &HashSet::default(),),
             Err(InstructionError::InvalidAccountData)
         );
 
@@ -1613,16 +1625,16 @@ mod tests {
             .unwrap();
 
         assert_eq!(
-            stake_keyed_account.set_lockup(&Lockup::default(), &HashSet::default(),),
+            stake_keyed_account.set_lockup(&LockupArgs::default(), &HashSet::default(),),
             Err(InstructionError::MissingRequiredSignature)
         );
 
         assert_eq!(
             stake_keyed_account.set_lockup(
-                &Lockup {
+                &LockupArgs {
-                    unix_timestamp: 1,
+                    unix_timestamp: Some(1),
-                    epoch: 1,
+                    epoch: Some(1),
-                    custodian,
+                    custodian: Some(custodian),
                 },
                 &vec![custodian].into_iter().collect()
             ),
@@ -1654,10 +1666,10 @@ mod tests {
 
         assert_eq!(
             stake_keyed_account.set_lockup(
-                &Lockup {
+                &LockupArgs {
-                    unix_timestamp: 1,
+                    unix_timestamp: Some(1),
-                    epoch: 1,
+                    epoch: Some(1),
-                    custodian,
+                    custodian: Some(custodian),
                 },
                 &HashSet::default(),
             ),
@@ -1665,15 +1677,129 @@ mod tests {
         );
         assert_eq!(
             stake_keyed_account.set_lockup(
+                &LockupArgs {
+                    unix_timestamp: Some(1),
+                    epoch: Some(1),
+                    custodian: Some(custodian),
+                },
+                &vec![custodian].into_iter().collect()
+            ),
+            Ok(())
+        );
+    }
+
+    #[test]
+    fn test_optional_lockup() {
+        let stake_pubkey = Pubkey::new_rand();
+        let stake_lamports = 42;
+        let stake_account = Account::new_ref_data_with_space(
+            stake_lamports,
+            &StakeState::Uninitialized,
+            std::mem::size_of::<StakeState>(),
+            &id(),
+        )
+        .expect("stake_account");
+        let stake_keyed_account = KeyedAccount::new(&stake_pubkey, false, &stake_account);
+
+        let custodian = Pubkey::new_rand();
+        stake_keyed_account
+            .initialize(
+                &Authorized::auto(&stake_pubkey),
                 &Lockup {
                     unix_timestamp: 1,
                     epoch: 1,
                     custodian,
                 },
+                &Rent::free(),
+            )
+            .unwrap();
+
+        assert_eq!(
+            stake_keyed_account.set_lockup(
+                &LockupArgs {
+                    unix_timestamp: None,
+                    epoch: None,
+                    custodian: None,
+                },
                 &vec![custodian].into_iter().collect()
             ),
             Ok(())
         );
+
+        assert_eq!(
+            stake_keyed_account.set_lockup(
+                &LockupArgs {
+                    unix_timestamp: Some(2),
+                    epoch: None,
+                    custodian: None,
+                },
+                &vec![custodian].into_iter().collect()
+            ),
+            Ok(())
+        );
+
+        if let StakeState::Initialized(Meta { lockup, .. }) =
+            StakeState::from(&stake_keyed_account.account.borrow()).unwrap()
+        {
+            assert_eq!(lockup.unix_timestamp, 2);
+            assert_eq!(lockup.epoch, 1);
+            assert_eq!(lockup.custodian, custodian);
+        } else {
+            assert!(false);
+        }
+
+        assert_eq!(
+            stake_keyed_account.set_lockup(
+                &LockupArgs {
+                    unix_timestamp: None,
+                    epoch: Some(3),
+                    custodian: None,
+                },
+                &vec![custodian].into_iter().collect()
+            ),
+            Ok(())
+        );
+
+        if let StakeState::Initialized(Meta { lockup, .. }) =
+            StakeState::from(&stake_keyed_account.account.borrow()).unwrap()
+        {
+            assert_eq!(lockup.unix_timestamp, 2);
+            assert_eq!(lockup.epoch, 3);
+            assert_eq!(lockup.custodian, custodian);
+        } else {
+            assert!(false);
+        }
+
+        let new_custodian = Pubkey::new_rand();
+        assert_eq!(
+            stake_keyed_account.set_lockup(
+                &LockupArgs {
+                    unix_timestamp: None,
+                    epoch: None,
+                    custodian: Some(new_custodian),
+                },
+                &vec![custodian].into_iter().collect()
+            ),
+            Ok(())
+        );
+
+        if let StakeState::Initialized(Meta { lockup, .. }) =
+            StakeState::from(&stake_keyed_account.account.borrow()).unwrap()
+        {
+            assert_eq!(lockup.unix_timestamp, 2);
+            assert_eq!(lockup.epoch, 3);
+            assert_eq!(lockup.custodian, new_custodian);
+        } else {
+            assert!(false);
+        }
+
+        assert_eq!(
+            stake_keyed_account.set_lockup(
+                &LockupArgs::default(),
+                &vec![custodian].into_iter().collect()
+            ),
+            Err(InstructionError::MissingRequiredSignature)
+        );
     }
 
     #[test]