diff --git a/programs/budget/src/budget_processor.rs b/programs/budget/src/budget_processor.rs index 5e68cc4687..e5dfb8d44b 100644 --- a/programs/budget/src/budget_processor.rs +++ b/programs/budget/src/budget_processor.rs @@ -36,7 +36,9 @@ fn apply_signature( if &payment.to == key { budget_state.pending_budget = None; contract_keyed_account.try_account_ref_mut()?.lamports -= payment.lamports; - witness_keyed_account.try_account_ref_mut()?.lamports += payment.lamports; + witness_keyed_account + .try_account_ref_mut()? + .checked_add_lamports(payment.lamports)?; return Ok(()); } } diff --git a/runtime/src/accounts_db.rs b/runtime/src/accounts_db.rs index 28c88423aa..0c4bc4901c 100644 --- a/runtime/src/accounts_db.rs +++ b/runtime/src/accounts_db.rs @@ -6194,7 +6194,7 @@ pub mod tests { if let Some((mut account, _)) = accounts.load_without_fixed_root(&ancestors, &pubkeys[idx]) { - account.lamports += 1; + account.checked_add_lamports(1).unwrap(); accounts.store_uncached(slot, &[(&pubkeys[idx], &account)]); if account.lamports == 0 { let ancestors = vec![(slot, 0)].into_iter().collect(); diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs index aa09dbf3a0..e6d7255958 100644 --- a/runtime/src/bank.rs +++ b/runtime/src/bank.rs @@ -9870,8 +9870,12 @@ pub(crate) mod tests { dup_account.lamports -= lamports; to_account.lamports += lamports; } - keyed_accounts[0].try_account_ref_mut()?.lamports -= lamports; - keyed_accounts[1].try_account_ref_mut()?.lamports += lamports; + keyed_accounts[0] + .try_account_ref_mut()? + .checked_sub_lamports(lamports)?; + keyed_accounts[1] + .try_account_ref_mut()? + .checked_add_lamports(lamports)?; Ok(()) } diff --git a/sdk/program/src/lamports.rs b/sdk/program/src/lamports.rs new file mode 100644 index 0000000000..36ecb0ae0d --- /dev/null +++ b/sdk/program/src/lamports.rs @@ -0,0 +1,22 @@ +use crate::instruction::InstructionError; +use thiserror::Error; + +#[derive(Debug, Error)] +pub enum LamportsError { + /// arithmetic underflowed + #[error("Arithmetic underflowed")] + ArithmeticUnderflow, + + /// arithmetic overflowed + #[error("Arithmetic overflowed")] + ArithmeticOverflow, +} + +impl From for InstructionError { + fn from(error: LamportsError) -> Self { + match error { + LamportsError::ArithmeticOverflow => InstructionError::ArithmeticOverflow, + LamportsError::ArithmeticUnderflow => InstructionError::ArithmeticOverflow, + } + } +} diff --git a/sdk/program/src/lib.rs b/sdk/program/src/lib.rs index 86c4eab2d8..fcac9c4faa 100644 --- a/sdk/program/src/lib.rs +++ b/sdk/program/src/lib.rs @@ -20,6 +20,7 @@ pub mod fee_calculator; pub mod hash; pub mod incinerator; pub mod instruction; +pub mod lamports; pub mod loader_instruction; pub mod loader_upgradeable_instruction; pub mod log; diff --git a/sdk/src/account.rs b/sdk/src/account.rs index fc673977c2..55c5b04f42 100644 --- a/sdk/src/account.rs +++ b/sdk/src/account.rs @@ -1,5 +1,6 @@ use crate::{ clock::{Epoch, INITIAL_RENT_EPOCH}, + lamports::LamportsError, pubkey::Pubkey, }; use solana_program::{account_info::AccountInfo, sysvar::Sysvar}; @@ -79,6 +80,22 @@ impl From for AccountSharedData { pub trait WritableAccount: ReadableAccount { fn set_lamports(&mut self, lamports: u64); + fn checked_add_lamports(&mut self, lamports: u64) -> Result<(), LamportsError> { + self.set_lamports( + self.lamports() + .checked_add(lamports) + .ok_or(LamportsError::ArithmeticOverflow)?, + ); + Ok(()) + } + fn checked_sub_lamports(&mut self, lamports: u64) -> Result<(), LamportsError> { + self.set_lamports( + self.lamports() + .checked_sub(lamports) + .ok_or(LamportsError::ArithmeticUnderflow)?, + ); + Ok(()) + } fn data_as_mut_slice(&mut self) -> &mut [u8]; fn set_owner(&mut self, owner: Pubkey); fn copy_into_owner_from_slice(&mut self, source: &[u8]); @@ -709,6 +726,53 @@ pub mod tests { ); } + #[test] + fn test_account_add_sub_lamports() { + let key = Pubkey::new_unique(); + let (mut account1, mut account2) = make_two_accounts(&key); + assert!(accounts_equal(&account1, &account2)); + account1.checked_add_lamports(1).unwrap(); + account2.checked_add_lamports(1).unwrap(); + assert!(accounts_equal(&account1, &account2)); + assert_eq!(account1.lamports(), 2); + account1.checked_sub_lamports(2).unwrap(); + account2.checked_sub_lamports(2).unwrap(); + assert!(accounts_equal(&account1, &account2)); + assert_eq!(account1.lamports(), 0); + } + + #[test] + #[should_panic(expected = "Overflow")] + fn test_account_checked_add_lamports_overflow() { + let key = Pubkey::new_unique(); + let (mut account1, _account2) = make_two_accounts(&key); + account1.checked_add_lamports(u64::MAX).unwrap(); + } + + #[test] + #[should_panic(expected = "Underflow")] + fn test_account_checked_sub_lamports_underflow() { + let key = Pubkey::new_unique(); + let (mut account1, _account2) = make_two_accounts(&key); + account1.checked_sub_lamports(u64::MAX).unwrap(); + } + + #[test] + #[should_panic(expected = "Overflow")] + fn test_account_checked_add_lamports_overflow2() { + let key = Pubkey::new_unique(); + let (_account1, mut account2) = make_two_accounts(&key); + account2.checked_add_lamports(u64::MAX).unwrap(); + } + + #[test] + #[should_panic(expected = "Underflow")] + fn test_account_checked_sub_lamports_underflow2() { + let key = Pubkey::new_unique(); + let (_account1, mut account2) = make_two_accounts(&key); + account2.checked_sub_lamports(u64::MAX).unwrap(); + } + #[test] #[allow(clippy::redundant_clone)] fn test_account_shared_data_all_fields() { @@ -726,15 +790,15 @@ pub mod tests { for pass in 0..4 { if field_index == 0 { if pass == 0 { - account1.lamports += 1; + account1.checked_add_lamports(1).unwrap(); } else if pass == 1 { - account_expected.lamports += 1; + account_expected.checked_add_lamports(1).unwrap(); account2.set_lamports(account2.lamports + 1); } else if pass == 2 { account1.set_lamports(account1.lamports + 1); } else if pass == 3 { - account_expected.lamports += 1; - account2.lamports += 1; + account_expected.checked_add_lamports(1).unwrap(); + account2.checked_add_lamports(1).unwrap(); } } else if field_index == 1 { if pass == 0 {