From d0992ea52a211939dbc1fbd7ac07d1657112e74f Mon Sep 17 00:00:00 2001
From: Jon Cinque <jon.cinque@gmail.com>
Date: Tue, 29 Nov 2022 12:52:56 +0100
Subject: [PATCH] Add more tolerance on withdrawal

---
 stake-pool/program/src/processor.rs           | 18 ++++++++++++-
 .../program/tests/withdraw_edge_cases.rs      | 25 +++++++++++--------
 2 files changed, 31 insertions(+), 12 deletions(-)

diff --git a/stake-pool/program/src/processor.rs b/stake-pool/program/src/processor.rs
index 99e580fb448..2b92ace9739 100644
--- a/stake-pool/program/src/processor.rs
+++ b/stake-pool/program/src/processor.rs
@@ -2551,7 +2551,23 @@ impl Processor {
                 }
                 StakeWithdrawSource::ValidatorRemoval => {
                     let split_from_lamports = stake_split_from.lamports();
-                    let upper_bound = split_from_lamports.saturating_add(lamports_per_pool_token);
+                    // The upper bound for reasonable tolerance is twice the lamports per
+                    // pool token because we have two sources of rounding. The first happens
+                    // when reducing the stake account to as close to the minimum as possible,
+                    // and the second happens on this withdrawal.
+                    //
+                    // For example, if the pool token is extremely valuable, it might only
+                    // be possible to reduce the stake account to a minimum of
+                    // `stake_rent + minimum_delegation + lamports_per_pool_token - 1`.
+                    //
+                    // After that, the minimum amount of pool tokens to get to this amount
+                    // may actually be worth
+                    // `stake_rent + minimum_delegation + lamports_per_pool_token * 2 - 2`.
+                    // We give an extra grace on this check of two lamports, which should be
+                    // reasonable. At worst, it just means that a withdrawer is losing out
+                    // on two lamports.
+                    let upper_bound = split_from_lamports
+                        .saturating_add(lamports_per_pool_token.saturating_mul(2));
                     if withdraw_lamports < split_from_lamports || withdraw_lamports > upper_bound {
                         msg!(
                             "Cannot withdraw a whole account worth {} lamports, \
diff --git a/stake-pool/program/tests/withdraw_edge_cases.rs b/stake-pool/program/tests/withdraw_edge_cases.rs
index afd21ee2c41..83a9d3d6c86 100644
--- a/stake-pool/program/tests/withdraw_edge_cases.rs
+++ b/stake-pool/program/tests/withdraw_edge_cases.rs
@@ -180,9 +180,10 @@ async fn fail_remove_validator() {
     );
 }
 
-#[test_case(1; "1 to 1")]
-#[test_case(5; "bigger multiple")]
-#[test_case(11; "biggest multiple")]
+#[test_case(0; "equal")]
+#[test_case(5; "big")]
+#[test_case(11; "bigger")]
+#[test_case(29; "biggest")]
 #[tokio::test]
 async fn success_remove_validator(multiple: u64) {
     let (
@@ -216,8 +217,12 @@ async fn success_remove_validator(multiple: u64) {
 
     let rent = context.banks_client.get_rent().await.unwrap();
     let stake_rent = rent.minimum_balance(std::mem::size_of::<stake::state::StakeState>());
+    let stake_pool = stake_pool_accounts
+        .get_stake_pool(&mut context.banks_client)
+        .await;
+    let lamports_per_pool_token = stake_pool.get_lamports_per_pool_token().unwrap();
 
-    // decrease all of stake except for 7 lamports, should be withdrawable
+    // decrease all of stake except for lamports_per_pool_token lamports, must be withdrawable
     let error = stake_pool_accounts
         .decrease_validator_stake(
             &mut context.banks_client,
@@ -225,7 +230,7 @@ async fn success_remove_validator(multiple: u64) {
             &context.last_blockhash,
             &validator_stake.stake_account,
             &validator_stake.transient_stake_account,
-            deposit_info.stake_lamports + stake_rent - multiple,
+            deposit_info.stake_lamports + stake_rent - lamports_per_pool_token,
             validator_stake.transient_stake_seed,
         )
         .await;
@@ -259,14 +264,12 @@ async fn success_remove_validator(multiple: u64) {
     )
     .await;
     // make sure it's actually more than the minimum
-    assert_ne!(remaining_lamports, stake_rent + stake_minimum_delegation);
+    assert!(remaining_lamports > stake_rent + stake_minimum_delegation);
 
-    let stake_pool = stake_pool_accounts
-        .get_stake_pool(&mut context.banks_client)
-        .await;
-    // round up to force one more pool token than needed
+    // round up to force one more pool token if needed
     let pool_tokens_post_fee =
-        remaining_lamports * stake_pool.pool_token_supply / stake_pool.total_lamports + 1;
+        (remaining_lamports * stake_pool.pool_token_supply + stake_pool.total_lamports - 1)
+            / stake_pool.total_lamports;
     let new_user_authority = Pubkey::new_unique();
     let pool_tokens = stake_pool_accounts.calculate_inverse_withdrawal_fee(pool_tokens_post_fee);
     let error = stake_pool_accounts