Add custodian option to withdraw-stake command (bp #9662) (#9675)

* Merge stake::withdraw instructions (#9617)

* Add custodian option to withdraw-stake command (#9662)

automerge

Co-authored-by: Greg Fitzgerald <greg@solana.com>

cli/src/cli.rs

@@ -339,6 +339,7 @@ pub enum CliCommand {
         destination_account_pubkey: Pubkey,
         lamports: u64,
         withdraw_authority: SignerIndex,
+        custodian: Option<SignerIndex>,
         sign_only: bool,
         blockhash_query: BlockhashQuery,
         nonce_account: Option<Pubkey>,
@@ -1968,6 +1969,7 @@ pub fn process_command(config: &CliConfig) -> ProcessResult {
             destination_account_pubkey,
             lamports,
             withdraw_authority,
+            custodian,
             sign_only,
             blockhash_query,
             ref nonce_account,
@@ -1980,6 +1982,7 @@ pub fn process_command(config: &CliConfig) -> ProcessResult {
             &destination_account_pubkey,
             *lamports,
             *withdraw_authority,
+            *custodian,
             *sign_only,
             blockhash_query,
             nonce_account.as_ref(),
@@ -3378,6 +3381,7 @@ mod tests {
             destination_account_pubkey: to_pubkey,
             lamports: 100,
             withdraw_authority: 0,
+            custodian: None,
             sign_only: false,
             blockhash_query: BlockhashQuery::All(blockhash_query::Source::Cluster),
             nonce_account: None,

cli/src/stake.rs

@@ -313,6 +313,14 @@ impl StakeSubCommands for App<'_, '_> {
                 .arg(nonce_arg())
                 .arg(nonce_authority_arg())
                 .arg(fee_payer_arg())
+                .arg(
+                    Arg::with_name("custodian")
+                        .long("custodian")
+                        .takes_value(true)
+                        .value_name("KEYPAIR")
+                        .validator(is_valid_signer)
+                        .help("Authority to override account lockup")
+                )
         )
         .subcommand(
             SubCommand::with_name("stake-set-lockup")
@@ -676,11 +684,15 @@ pub fn parse_stake_withdraw_stake(
     let (nonce_authority, nonce_authority_pubkey) =
         signer_of(matches, NONCE_AUTHORITY_ARG.name, wallet_manager)?;
     let (fee_payer, fee_payer_pubkey) = signer_of(matches, FEE_PAYER_ARG.name, wallet_manager)?;
+    let (custodian, custodian_pubkey) = signer_of(matches, "custodian", wallet_manager)?;
 
     let mut bulk_signers = vec![withdraw_authority, fee_payer];
     if nonce_account.is_some() {
         bulk_signers.push(nonce_authority);
     }
+    if custodian.is_some() {
+        bulk_signers.push(custodian);
+    }
     let signer_info =
         generate_unique_signers(bulk_signers, matches, default_signer_path, wallet_manager)?;
 
@@ -695,6 +707,7 @@ pub fn parse_stake_withdraw_stake(
             nonce_account,
             nonce_authority: signer_info.index_of(nonce_authority_pubkey).unwrap(),
             fee_payer: signer_info.index_of(fee_payer_pubkey).unwrap(),
+            custodian: custodian_pubkey.and_then(|_| signer_info.index_of(custodian_pubkey)),
         },
         signers: signer_info.signers,
     })
@@ -1013,6 +1026,7 @@ pub fn process_withdraw_stake(
     destination_account_pubkey: &Pubkey,
     lamports: u64,
     withdraw_authority: SignerIndex,
+    custodian: Option<SignerIndex>,
     sign_only: bool,
     blockhash_query: &BlockhashQuery,
     nonce_account: Option<&Pubkey>,
@@ -1022,12 +1036,14 @@ pub fn process_withdraw_stake(
     let (recent_blockhash, fee_calculator) =
         blockhash_query.get_blockhash_and_fee_calculator(rpc_client)?;
     let withdraw_authority = config.signers[withdraw_authority];
+    let custodian = custodian.map(|index| config.signers[index]);
 
     let ixs = vec![stake_instruction::withdraw(
         stake_account_pubkey,
         &withdraw_authority.pubkey(),
         destination_account_pubkey,
         lamports,
+        custodian.map(|signer| signer.pubkey()).as_ref(),
     )];
 
     let fee_payer = config.signers[fee_payer];
@@ -1496,6 +1512,9 @@ mod tests {
         let (stake_authority_keypair_file, mut tmp_file) = make_tmp_file();
         let stake_authority_keypair = Keypair::new();
         write_keypair(&stake_authority_keypair, tmp_file.as_file_mut()).unwrap();
+        let (custodian_keypair_file, mut tmp_file) = make_tmp_file();
+        let custodian_keypair = Keypair::new();
+        write_keypair(&custodian_keypair, tmp_file.as_file_mut()).unwrap();
 
         // stake-authorize subcommand
         let stake_account_string = stake_account_pubkey.to_string();
@@ -2434,6 +2453,7 @@ mod tests {
                     destination_account_pubkey: stake_account_pubkey,
                     lamports: 42_000_000_000,
                     withdraw_authority: 0,
+                    custodian: None,
                     sign_only: false,
                     blockhash_query: BlockhashQuery::All(blockhash_query::Source::Cluster),
                     nonce_account: None,
@@ -2463,6 +2483,7 @@ mod tests {
                     destination_account_pubkey: stake_account_pubkey,
                     lamports: 42_000_000_000,
                     withdraw_authority: 1,
+                    custodian: None,
                     sign_only: false,
                     blockhash_query: BlockhashQuery::All(blockhash_query::Source::Cluster),
                     nonce_account: None,
@@ -2478,6 +2499,39 @@ mod tests {
             }
         );
 
+        // Test WithdrawStake Subcommand w/ custodian
+        let test_withdraw_stake = test_commands.clone().get_matches_from(vec![
+            "test",
+            "withdraw-stake",
+            &stake_account_string,
+            &stake_account_string,
+            "42",
+            "--custodian",
+            &custodian_keypair_file,
+        ]);
+
+        assert_eq!(
+            parse_command(&test_withdraw_stake, &default_keypair_file, &mut None).unwrap(),
+            CliCommandInfo {
+                command: CliCommand::WithdrawStake {
+                    stake_account_pubkey,
+                    destination_account_pubkey: stake_account_pubkey,
+                    lamports: 42_000_000_000,
+                    withdraw_authority: 0,
+                    custodian: Some(1),
+                    sign_only: false,
+                    blockhash_query: BlockhashQuery::All(blockhash_query::Source::Cluster),
+                    nonce_account: None,
+                    nonce_authority: 0,
+                    fee_payer: 0,
+                },
+                signers: vec![
+                    read_keypair_file(&default_keypair_file).unwrap().into(),
+                    read_keypair_file(&custodian_keypair_file).unwrap().into()
+                ],
+            }
+        );
+
         // Test WithdrawStake Subcommand w/ authority and offline nonce
         let test_withdraw_stake = test_commands.clone().get_matches_from(vec![
             "test",
@@ -2507,6 +2561,7 @@ mod tests {
                     destination_account_pubkey: stake_account_pubkey,
                     lamports: 42_000_000_000,
                     withdraw_authority: 0,
+                    custodian: None,
                     sign_only: false,
                     blockhash_query: BlockhashQuery::FeeCalculator(
                         blockhash_query::Source::NonceAccount(nonce_account),

cli/tests/stake.rs

@@ -1451,6 +1451,7 @@ fn test_offline_nonced_create_stake_account_and_withdraw() {
         destination_account_pubkey: recipient_pubkey,
         lamports: 42,
         withdraw_authority: 0,
+        custodian: None,
         sign_only: true,
         blockhash_query: BlockhashQuery::None(nonce_hash),
         nonce_account: Some(nonce_pubkey),
@@ -1466,6 +1467,7 @@ fn test_offline_nonced_create_stake_account_and_withdraw() {
         destination_account_pubkey: recipient_pubkey,
         lamports: 42,
         withdraw_authority: 0,
+        custodian: None,
         sign_only: false,
         blockhash_query: BlockhashQuery::FeeCalculator(
             blockhash_query::Source::NonceAccount(nonce_pubkey),

programs/stake/src/stake_instruction.rs

@@ -328,8 +328,9 @@ pub fn withdraw(
     withdrawer_pubkey: &Pubkey,
     to_pubkey: &Pubkey,
     lamports: u64,
+    custodian_pubkey: Option<&Pubkey>,
 ) -> Instruction {
-    let account_metas = vec![
+    let mut account_metas = vec![
         AccountMeta::new(*stake_pubkey, false),
         AccountMeta::new(*to_pubkey, false),
         AccountMeta::new_readonly(sysvar::clock::id(), false),
@@ -337,24 +338,9 @@ pub fn withdraw(
     ]
     .with_signer(withdrawer_pubkey);
 
-    Instruction::new(id(), &StakeInstruction::Withdraw(lamports), account_metas)
+    if let Some(custodian_pubkey) = custodian_pubkey {
-}
+        account_metas = account_metas.with_signer(custodian_pubkey)
-
+    }
-pub fn withdraw_early(
-    stake_pubkey: &Pubkey,
-    withdrawer_pubkey: &Pubkey,
-    to_pubkey: &Pubkey,
-    lamports: u64,
-    custodian_pubkey: &Pubkey,
-) -> Instruction {
-    let account_metas = vec![
-        AccountMeta::new(*stake_pubkey, false),
-        AccountMeta::new(*to_pubkey, false),
-        AccountMeta::new_readonly(sysvar::clock::id(), false),
-        AccountMeta::new_readonly(sysvar::stake_history::id(), false),
-    ]
-    .with_signer(withdrawer_pubkey)
-    .with_signer(custodian_pubkey);
 
     Instruction::new(id(), &StakeInstruction::Withdraw(lamports), account_metas)
 }
@@ -535,7 +521,8 @@ mod tests {
                 &Pubkey::default(),
                 &Pubkey::default(),
                 &Pubkey::new_rand(),
-                100
+                100,
+                None,
             )),
             Err(InstructionError::InvalidAccountData),
         );

runtime/tests/stake.rs

@@ -224,6 +224,7 @@ fn test_stake_account_lifetime() {
             &stake_pubkey,
             &Pubkey::new_rand(),
             1,
+            None,
         )],
         Some(&mint_pubkey),
     );
@@ -318,6 +319,7 @@ fn test_stake_account_lifetime() {
             &stake_pubkey,
             &Pubkey::new_rand(),
             lamports / 2 - split_staked + 1,
+            None,
         )],
         Some(&mint_pubkey),
     );
@@ -339,6 +341,7 @@ fn test_stake_account_lifetime() {
             &stake_pubkey,
             &Pubkey::new_rand(),
             lamports / 2,
+            None,
         )],
         Some(&mint_pubkey),
     );
@@ -354,6 +357,7 @@ fn test_stake_account_lifetime() {
             &stake_pubkey,
             &Pubkey::new_rand(),
             lamports / 2 - split_staked,
+            None,
         )],
         Some(&mint_pubkey),
     );
@@ -378,6 +382,7 @@ fn test_stake_account_lifetime() {
             &stake_pubkey,
             &Pubkey::new_rand(),
             split_staked,
+            None,
         )],
         Some(&mint_pubkey),
     );

stake-monitor/src/lib.rs

@@ -450,6 +450,7 @@ mod test {
                         &stake3_keypair.pubkey(),
                         &payer.pubkey(),
                         one_sol,
+                        None,
                     )],
                     Some(&payer.pubkey()),
                 ),