From ae82b5ebfd7066044b834cd7584a1a4d570ddad5 Mon Sep 17 00:00:00 2001
From: Trent Nelson <trent@solana.com>
Date: Fri, 12 Feb 2021 17:35:44 -0700
Subject: [PATCH] stake: add lamports overflow test for withdraw

---
 programs/stake/src/stake_state.rs | 54 +++++++++++++++++++++++++++++++
 1 file changed, 54 insertions(+)

diff --git a/programs/stake/src/stake_state.rs b/programs/stake/src/stake_state.rs
index 2cd0a17999..757da1f4bc 100644
--- a/programs/stake/src/stake_state.rs
+++ b/programs/stake/src/stake_state.rs
@@ -3239,6 +3239,60 @@ mod tests {
         );
         assert_eq!(stake_account.borrow().lamports, 0);
         assert_eq!(stake_keyed_account.state(), Ok(StakeState::Uninitialized));
+
+        // overflow
+        let rent = Rent::default();
+        let rent_exempt_reserve = rent.minimum_balance(std::mem::size_of::<StakeState>());
+        let authority_pubkey = Pubkey::new_unique();
+        let stake_pubkey = Pubkey::new_unique();
+        let stake_account = Account::new_ref_data_with_space(
+            1_000_000_000,
+            &StakeState::Initialized(Meta {
+                rent_exempt_reserve,
+                authorized: Authorized {
+                    staker: authority_pubkey,
+                    withdrawer: authority_pubkey,
+                },
+                lockup: Lockup::default(),
+            }),
+            std::mem::size_of::<StakeState>(),
+            &id(),
+        )
+        .expect("stake_account");
+
+        let stake2_pubkey = Pubkey::new_unique();
+        let stake2_account = Account::new_ref_data_with_space(
+            1_000_000_000,
+            &StakeState::Initialized(Meta {
+                rent_exempt_reserve,
+                authorized: Authorized {
+                    staker: authority_pubkey,
+                    withdrawer: authority_pubkey,
+                },
+                lockup: Lockup::default(),
+            }),
+            std::mem::size_of::<StakeState>(),
+            &id(),
+        )
+        .expect("stake2_account");
+
+        let stake_keyed_account = KeyedAccount::new(&stake_pubkey, true, &stake_account);
+        let stake2_keyed_account = KeyedAccount::new(&stake2_pubkey, false, &stake2_account);
+        let authority_account = Account::new_ref(42, 0, &system_program::id());
+        let authority_keyed_account =
+            KeyedAccount::new(&authority_pubkey, true, &authority_account);
+
+        assert_eq!(
+            stake_keyed_account.withdraw(
+                u64::MAX - 10,
+                &stake2_keyed_account,
+                &clock,
+                &StakeHistory::default(),
+                &authority_keyed_account,
+                None,
+            ),
+            Err(InstructionError::InsufficientFunds),
+        );
     }
 
     #[test]