stake-pool: Add user transfer authority on withdraw (#1640)

joncinque · web-flow · commit 11e207cc856c · 2021-04-27T12:53:46.000Z
The stake pool expects pool tokens to be delegated to the withdraw
authority before performing a withdrawal. If a user delegates too many
tokens to the withdraw authority, anyone else can take the rest of their
tokens by doing their own withdrawal.

Delegate pool tokens to an ephemeral keypair and sign with that.
diff --git a/stake-pool/cli/src/main.rs b/stake-pool/cli/src/main.rs
@@ -877,15 +877,20 @@ fn command_withdraw(
 
     // Construct transaction to withdraw from withdraw_accounts account list
     let mut instructions: Vec<Instruction> = vec![];
-    let mut signers = vec![config.fee_payer.as_ref(), config.token_owner.as_ref()];
+    let user_transfer_authority = Keypair::new(); // ephemeral keypair just to do the transfer
+    let mut signers = vec![
+        config.fee_payer.as_ref(),
+        config.token_owner.as_ref(),
+        &user_transfer_authority,
+    ];
     let stake_receiver_account = Keypair::new(); // Will be added to signers if creating new account
 
     instructions.push(
         // Approve spending token
         spl_token::instruction::approve(
             &spl_token::id(),
             &withdraw_from,
-            &pool_withdraw_authority,
+            &user_transfer_authority.pubkey(),
             &config.token_owner.pubkey(),
             &[],
             pool_amount,
@@ -948,6 +953,7 @@ fn command_withdraw(
             &withdraw_account.address,
             &stake_receiver.unwrap(), // Cannot be none at this point
             &config.staker.pubkey(),
+            &user_transfer_authority.pubkey(),
             &withdraw_from,
             &stake_pool.pool_mint,
             &spl_token::id(),
diff --git a/stake-pool/program/src/instruction.rs b/stake-pool/program/src/instruction.rs
@@ -225,12 +225,13 @@ pub enum StakePoolInstruction {
     ///   3. `[w]` Validator or reserve stake account to split
     ///   4. `[w]` Unitialized stake account to receive withdrawal
     ///   5. `[]` User account to set as a new withdraw authority
-    ///   6. `[w]` User account with pool tokens to burn from
-    ///   7. `[w]` Pool token mint account
-    ///   8. '[]' Sysvar clock account (required)
-    ///   9. `[]` Pool token program id
-    ///   10. `[]` Stake program id,
-    ///   userdata: amount to withdraw
+    ///   6. `[s]` User transfer authority, for pool token account
+    ///   7. `[w]` User account with pool tokens to burn from
+    ///   8. `[w]` Pool token mint account
+    ///   9. `[]` Sysvar clock account (required)
+    ///  10. `[]` Pool token program id
+    ///  11. `[]` Stake program id,
+    ///  userdata: amount of pool tokens to withdraw
     Withdraw(u64),
 
     ///  (Manager only) Update manager
@@ -645,8 +646,9 @@ pub fn withdraw(
     stake_pool_withdraw: &Pubkey,
     stake_to_split: &Pubkey,
     stake_to_receive: &Pubkey,
-    user_withdrawer: &Pubkey,
-    burn_from: &Pubkey,
+    user_stake_authority: &Pubkey,
+    user_transfer_authority: &Pubkey,
+    user_pool_token_account: &Pubkey,
     pool_mint: &Pubkey,
     token_program_id: &Pubkey,
     amount: u64,
@@ -657,8 +659,9 @@ pub fn withdraw(
         AccountMeta::new_readonly(*stake_pool_withdraw, false),
         AccountMeta::new(*stake_to_split, false),
         AccountMeta::new(*stake_to_receive, false),
-        AccountMeta::new_readonly(*user_withdrawer, false),
-        AccountMeta::new(*burn_from, false),
+        AccountMeta::new_readonly(*user_stake_authority, false),
+        AccountMeta::new_readonly(*user_transfer_authority, true),
+        AccountMeta::new(*user_pool_token_account, false),
         AccountMeta::new(*pool_mint, false),
         AccountMeta::new_readonly(sysvar::clock::id(), false),
         AccountMeta::new_readonly(*token_program_id, false),
diff --git a/stake-pool/program/src/processor.rs b/stake-pool/program/src/processor.rs
@@ -336,19 +336,12 @@ impl Processor {
     /// Issue a spl_token `Burn` instruction.
     #[allow(clippy::too_many_arguments)]
     fn token_burn<'a>(
-        stake_pool: &Pubkey,
         token_program: AccountInfo<'a>,
         burn_account: AccountInfo<'a>,
         mint: AccountInfo<'a>,
         authority: AccountInfo<'a>,
-        authority_type: &[u8],
-        bump_seed: u8,
         amount: u64,
     ) -> Result<(), ProgramError> {
-        let me_bytes = stake_pool.to_bytes();
-        let authority_signature_seeds = [&me_bytes[..32], authority_type, &[bump_seed]];
-        let signers = &[&authority_signature_seeds[..]];
-
         let ix = spl_token::instruction::burn(
             token_program.key,
             burn_account.key,
@@ -358,11 +351,7 @@ impl Processor {
             amount,
         )?;
 
-        invoke_signed(
-            &ix,
-            &[burn_account, mint, authority, token_program],
-            signers,
-        )
+        invoke(&ix, &[burn_account, mint, authority, token_program])
     }
 
     /// Issue a spl_token `MintTo` instruction.
@@ -1640,7 +1629,8 @@ impl Processor {
         let withdraw_authority_info = next_account_info(account_info_iter)?;
         let stake_split_from = next_account_info(account_info_iter)?;
         let stake_split_to = next_account_info(account_info_iter)?;
-        let user_stake_authority = next_account_info(account_info_iter)?;
+        let user_stake_authority_info = next_account_info(account_info_iter)?;
+        let user_transfer_authority_info = next_account_info(account_info_iter)?;
         let burn_from_info = next_account_info(account_info_iter)?;
         let pool_mint_info = next_account_info(account_info_iter)?;
         let clock_info = next_account_info(account_info_iter)?;
@@ -1742,13 +1732,10 @@ impl Processor {
         };
 
         Self::token_burn(
-            stake_pool_info.key,
             token_program_info.clone(),
             burn_from_info.clone(),
             pool_mint_info.clone(),
-            withdraw_authority_info.clone(),
-            AUTHORITY_WITHDRAW,
-            stake_pool.withdraw_bump_seed,
+            user_transfer_authority_info.clone(),
             pool_tokens,
         )?;
 
@@ -1768,7 +1755,7 @@ impl Processor {
             withdraw_authority_info.clone(),
             AUTHORITY_WITHDRAW,
             stake_pool.withdraw_bump_seed,
-            user_stake_authority.key,
+            user_stake_authority_info.key,
             clock_info.clone(),
             stake_program_info.clone(),
         )?;
diff --git a/stake-pool/program/tests/helpers/mod.rs b/stake-pool/program/tests/helpers/mod.rs
@@ -182,7 +182,7 @@ pub async fn delegate_tokens(
     delegate: &Pubkey,
     amount: u64,
 ) {
-    let mut transaction = Transaction::new_with_payer(
+    let transaction = Transaction::new_signed_with_payer(
         &[spl_token::instruction::approve(
             &spl_token::id(),
             &account,
@@ -193,8 +193,9 @@ pub async fn delegate_tokens(
         )
         .unwrap()],
         Some(&payer.pubkey()),
+        &[payer, manager],
+        *recent_blockhash,
     );
-    transaction.sign(&[payer, manager], *recent_blockhash);
     banks_client.process_transaction(transaction).await.unwrap();
 }
 
@@ -656,6 +657,7 @@ impl StakePoolAccounts {
         payer: &Keypair,
         recent_blockhash: &Hash,
         stake_recipient: &Pubkey,
+        user_transfer_authority: &Keypair,
         pool_account: &Pubkey,
         validator_stake_account: &Pubkey,
         recipient_new_authority: &Pubkey,
@@ -670,14 +672,15 @@ impl StakePoolAccounts {
                 validator_stake_account,
                 stake_recipient,
                 recipient_new_authority,
+                &user_transfer_authority.pubkey(),
                 pool_account,
                 &self.pool_mint.pubkey(),
                 &spl_token::id(),
                 amount,
             )
             .unwrap()],
             Some(&payer.pubkey()),
-            &[payer],
+            &[payer, user_transfer_authority],
             *recent_blockhash,
         );
         banks_client.process_transaction(transaction).await.err()
diff --git a/stake-pool/program/tests/vsa_remove.rs b/stake-pool/program/tests/vsa_remove.rs
@@ -485,13 +485,25 @@ async fn success_with_deactivating_transient_stake() {
     )
     .await;
 
+    let user_transfer_authority = Keypair::new();
     let new_authority = Pubkey::new_unique();
+    delegate_tokens(
+        &mut banks_client,
+        &payer,
+        &recent_blockhash,
+        &deposit_info.pool_account.pubkey(),
+        &deposit_info.authority,
+        &user_transfer_authority.pubkey(),
+        1,
+    )
+    .await;
     let error = stake_pool_accounts
         .withdraw_stake(
             &mut banks_client,
             &payer,
             &recent_blockhash,
             &user_stake_recipient.pubkey(),
+            &user_transfer_authority,
             &deposit_info.pool_account.pubkey(),
             &validator_stake.stake_account,
             &new_authority,
diff --git a/stake-pool/program/tests/withdraw.rs b/stake-pool/program/tests/withdraw.rs
