stake-pool: Make it impossible to leech value (#2218)

Author: joncinque

Cargo.lock

@@ -20,6 +20,7 @@ num_enum = "0.5.3"
 serde = "1.0.127"
 serde_derive = "1.0.103"
 solana-program = "1.7.7"
+spl-math = { version = "0.1", path = "../../libraries/math", features = [ "no-entrypoint" ] }
 spl-token = { version = "3.2", path = "../../token/program", features = [ "no-entrypoint" ] }
 thiserror = "1.0"
 bincode = "1.3.1"

stake-pool/program/Cargo.toml

@@ -102,6 +102,12 @@ pub enum StakePoolError {
     /// Proposed fee increase exceeds stipulated ratio
     #[error("FeeIncreaseTooHigh")]
     FeeIncreaseTooHigh,
+    /// Not enough pool tokens provided to withdraw stake with one lamport
+    #[error("WithdrawalTooSmall")]
+    WithdrawalTooSmall,
+    /// Not enough lamports provided for deposit to result in one pool token
+    #[error("DepositTooSmall")]
+    DepositTooSmall,
 }
 impl From<StakePoolError> for ProgramError {
     fn from(e: StakePoolError) -> Self {

stake-pool/program/src/error.rs

@@ -1858,6 +1858,10 @@ impl Processor {
             .calc_pool_tokens_for_deposit(all_deposit_lamports)
             .ok_or(StakePoolError::CalculationFailure)?;
 
+        if new_pool_tokens == 0 {
+            return Err(StakePoolError::DepositTooSmall.into());
+        }
+
         Self::token_mint_to(
             stake_pool_info.key,
             token_program_info.clone(),
@@ -1978,6 +1982,10 @@ impl Processor {
             .calc_lamports_withdraw_amount(pool_tokens_burnt)
             .ok_or(StakePoolError::CalculationFailure)?;
 
+        if withdraw_lamports == 0 {
+            return Err(StakePoolError::WithdrawalTooSmall.into());
+        }
+
         let has_active_stake = validator_list
             .find::<ValidatorStakeInfo>(
                 &0u64.to_le_bytes(),
@@ -2426,6 +2434,8 @@ impl PrintProgramError for StakePoolError {
             StakePoolError::IncorrectWithdrawVoteAddress => msg!("Error: The provided withdraw stake account is not the preferred deposit vote account"),
             StakePoolError::InvalidMintFreezeAuthority => msg!("Error: The mint has an invalid freeze authority"),
             StakePoolError::FeeIncreaseTooHigh => msg!("Error: The fee cannot increase by a factor exceeding the stipulated ratio"),
+            StakePoolError::WithdrawalTooSmall => msg!("Error: Not enough pool tokens provided to withdraw 1-lamport stake"),
+            StakePoolError::DepositTooSmall => msg!("Error: Not enough lamports provided for deposit to result in one pool token"),
         }
     }
 }

stake-pool/program/src/processor.rs

@@ -14,6 +14,7 @@ use {
         program_pack::{Pack, Sealed},
         pubkey::{Pubkey, PUBKEY_BYTES},
     },
+    spl_math::checked_ceil_div::CheckedCeilDiv,
     std::convert::TryFrom,
 };
 
@@ -135,12 +136,10 @@ impl StakePool {
 
     /// calculate lamports amount on withdrawal
     pub fn calc_lamports_withdraw_amount(&self, pool_tokens: u64) -> Option<u64> {
-        u64::try_from(
-            (pool_tokens as u128)
-                .checked_mul(self.total_stake_lamports as u128)?
-                .checked_div(self.pool_token_supply as u128)?,
-        )
-        .ok()
+        let (quotient, _) = (pool_tokens as u128)
+            .checked_mul(self.total_stake_lamports as u128)?
+            .checked_ceil_div(self.pool_token_supply as u128)?;
+        u64::try_from(quotient).ok()
     }
 
     /// calculate pool tokens to be deducted as withdrawal fees
@@ -159,12 +158,16 @@ impl StakePool {
         let total_stake_lamports =
             (self.total_stake_lamports as u128).checked_add(reward_lamports as u128)?;
         let fee_lamports = self.fee.apply(reward_lamports)?;
-        u64::try_from(
-            (self.pool_token_supply as u128)
-                .checked_mul(fee_lamports)?
-                .checked_div(total_stake_lamports.checked_sub(fee_lamports)?)?,
-        )
-        .ok()
+        if total_stake_lamports == fee_lamports || self.pool_token_supply == 0 {
+            Some(reward_lamports)
+        } else {
+            u64::try_from(
+                (self.pool_token_supply as u128)
+                    .checked_mul(fee_lamports)?
+                    .checked_div(total_stake_lamports.checked_sub(fee_lamports)?)?,
+            )
+            .ok()
+        }
     }
 
     /// Checks that the withdraw or deposit authority is valid
@@ -777,7 +780,22 @@ mod test {
         let fee_lamports = stake_pool
             .calc_lamports_withdraw_amount(pool_token_fee)
             .unwrap();
-        assert_eq!(fee_lamports, LAMPORTS_PER_SOL - 1); // lose 1 lamport of precision
+        assert_eq!(fee_lamports, LAMPORTS_PER_SOL);
+    }
+
+    #[test]
+    fn divide_by_zero_fee() {
+        let stake_pool = StakePool {
+            total_stake_lamports: 0,
+            fee: Fee {
+                numerator: 1,
+                denominator: 10,
+            },
+            ..StakePool::default()
+        };
+        let rewards = 10;
+        let fee = stake_pool.calc_fee_amount(rewards).unwrap();
+        assert_eq!(fee, rewards);
     }
 
     proptest! {
@@ -814,4 +832,33 @@ mod test {
                 max_fee_lamports, fee_lamports, epsilon);
         }
     }
+
+    prop_compose! {
+        fn total_tokens_and_deposit()(total_lamports in 1..u64::MAX)(
+            total_lamports in Just(total_lamports),
+            pool_token_supply in 1..=total_lamports,
+            deposit_lamports in 1..total_lamports,
+        ) -> (u64, u64, u64) {
+            (total_lamports - deposit_lamports, pool_token_supply.saturating_sub(deposit_lamports).max(1), deposit_lamports)
+        }
+    }
+
+    proptest! {
+        #[test]
+        fn deposit_and_withdraw(
+            (total_stake_lamports, pool_token_supply, deposit_stake) in total_tokens_and_deposit()
+        ) {
+            let mut stake_pool = StakePool {
+                total_stake_lamports,
+                pool_token_supply,
+                ..StakePool::default()
+            };
+            let deposit_result = stake_pool.calc_pool_tokens_for_deposit(deposit_stake).unwrap();
+            prop_assume!(deposit_result > 0);
+            stake_pool.total_stake_lamports += deposit_stake;
+            stake_pool.pool_token_supply += deposit_result;
+            let withdraw_result = stake_pool.calc_lamports_withdraw_amount(deposit_result).unwrap();
+            assert!(withdraw_result <= deposit_stake);
+        }
+    }
 }

stake-pool/program/src/state.rs

@@ -448,10 +448,10 @@ async fn fail_with_unknown_validator() {
         recent_blockhash,
         stake_pool_accounts,
         _,
-        _,
+        deposit_info,
         user_transfer_authority,
         user_stake_recipient,
-        _,
+        tokens_to_withdraw,
     ) = setup().await;
 
     let validator_stake_account =
@@ -475,80 +475,18 @@ async fn fail_with_unknown_validator() {
         )
         .await;
 
-    let user_pool_account = Keypair::new();
-    let user = Keypair::new();
-    create_token_account(
-        &mut banks_client,
-        &payer,
-        &recent_blockhash,
-        &user_pool_account,
-        &stake_pool_accounts.pool_mint.pubkey(),
-        &user.pubkey(),
-    )
-    .await
-    .unwrap();
-
-    let user = Keypair::new();
-    // make stake account
-    let user_stake = Keypair::new();
-    let lockup = stake_program::Lockup::default();
-    let authorized = stake_program::Authorized {
-        staker: stake_pool_accounts.deposit_authority,
-        withdrawer: stake_pool_accounts.deposit_authority,
-    };
-    create_independent_stake_account(
-        &mut banks_client,
-        &payer,
-        &recent_blockhash,
-        &user_stake,
-        &authorized,
-        &lockup,
-        TEST_STAKE_AMOUNT,
-    )
-    .await;
-    // make pool token account
-    let user_pool_account = Keypair::new();
-    create_token_account(
-        &mut banks_client,
-        &payer,
-        &recent_blockhash,
-        &user_pool_account,
-        &stake_pool_accounts.pool_mint.pubkey(),
-        &user.pubkey(),
-    )
-    .await
-    .unwrap();
-
-    let user_pool_account = user_pool_account.pubkey();
-    let pool_tokens = get_token_balance(&mut banks_client, &user_pool_account).await;
-
-    let tokens_to_burn = pool_tokens / 4;
-
-    // Delegate tokens for burning
-    delegate_tokens(
-        &mut banks_client,
-        &payer,
-        &recent_blockhash,
-        &user_pool_account,
-        &user,
-        &user_transfer_authority.pubkey(),
-        tokens_to_burn,
-    )
-    .await;
-
     let new_authority = Pubkey::new_unique();
-
     let transaction_error = stake_pool_accounts
         .withdraw_stake(
             &mut banks_client,
             &payer,
             &recent_blockhash,
             &user_stake_recipient.pubkey(),
             &user_transfer_authority,
-            &user_pool_account,
+            &deposit_info.pool_account.pubkey(),
             &validator_stake_account.stake_account,
             &new_authority,
-            tokens_to_burn,
+            tokens_to_withdraw,
         )
         .await
         .unwrap()