From 257b5601094d5e7170946a8de9f6a9532d917ed5 Mon Sep 17 00:00:00 2001
From: acheron <98934430+acheroncrypto@users.noreply.github.com>
Date: Sun, 14 Apr 2024 23:04:51 +0200
Subject: [PATCH] lang: Return overflow error from `Lamports` trait operations
 (#2907)

---
 CHANGELOG.md                            |  1 +
 lang/src/lib.rs                         | 11 ++++++++--
 tests/misc/programs/lamports/src/lib.rs | 29 +++++++++++++++++++++----
 tests/misc/tests/lamports/lamports.ts   | 15 +++++--------
 4 files changed, 41 insertions(+), 15 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ffd472466..a924f308b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -44,6 +44,7 @@ The minor version will be incremented upon a breaking change and the patch versi
 - idl: Store deployment addresses for other clusters ([#2892](https://github.com/coral-xyz/anchor/pull/2892)).
 - lang: Add `Event` utility type to get events from bytes ([#2897](https://github.com/coral-xyz/anchor/pull/2897)).
 - lang, spl: Add support for [token extensions](https://solana.com/solutions/token-extensions) ([#2789](https://github.com/coral-xyz/anchor/pull/2789)).
+- lang: Return overflow error from `Lamports` trait operations ([#2907](https://github.com/coral-xyz/anchor/pull/2907)).
 
 ### Fixes
 
diff --git a/lang/src/lib.rs b/lang/src/lib.rs
index bb3766cc9..98a26176d 100644
--- a/lang/src/lib.rs
+++ b/lang/src/lib.rs
@@ -28,6 +28,7 @@ extern crate self as anchor_lang;
 use bytemuck::{Pod, Zeroable};
 use solana_program::account_info::AccountInfo;
 use solana_program::instruction::AccountMeta;
+use solana_program::program_error::ProgramError;
 use solana_program::pubkey::Pubkey;
 use std::{collections::BTreeSet, fmt::Debug, io::Write};
 
@@ -197,7 +198,10 @@ pub trait Lamports<'info>: AsRef<AccountInfo<'info>> {
     ///
     /// See [`Lamports::sub_lamports`] for subtracting lamports.
     fn add_lamports(&self, amount: u64) -> Result<&Self> {
-        **self.as_ref().try_borrow_mut_lamports()? += amount;
+        **self.as_ref().try_borrow_mut_lamports()? = self
+            .get_lamports()
+            .checked_add(amount)
+            .ok_or(ProgramError::ArithmeticOverflow)?;
         Ok(self)
     }
 
@@ -215,7 +219,10 @@ pub trait Lamports<'info>: AsRef<AccountInfo<'info>> {
     ///
     /// See [`Lamports::add_lamports`] for adding lamports.
     fn sub_lamports(&self, amount: u64) -> Result<&Self> {
-        **self.as_ref().try_borrow_mut_lamports()? -= amount;
+        **self.as_ref().try_borrow_mut_lamports()? = self
+            .get_lamports()
+            .checked_sub(amount)
+            .ok_or(ProgramError::ArithmeticOverflow)?;
         Ok(self)
     }
 }
diff --git a/tests/misc/programs/lamports/src/lib.rs b/tests/misc/programs/lamports/src/lib.rs
index 84ac072d2..e8ea9732f 100644
--- a/tests/misc/programs/lamports/src/lib.rs
+++ b/tests/misc/programs/lamports/src/lib.rs
@@ -6,7 +6,7 @@ declare_id!("Lamports11111111111111111111111111111111111");
 pub mod lamports {
     use super::*;
 
-    pub fn test_lamports_trait(ctx: Context<TestLamportsTrait>, amount: u64) -> Result<()> {
+    pub fn transfer(ctx: Context<Transfer>, amount: u64) -> Result<()> {
         let pda = &ctx.accounts.pda;
         let signer = &ctx.accounts.signer;
 
@@ -52,13 +52,29 @@ pub mod lamports {
 
         Ok(())
     }
+
+    // Return overflow error in the case of overflow (instead of panicking)
+    pub fn overflow(ctx: Context<Overflow>) -> Result<()> {
+        let pda = &ctx.accounts.pda;
+
+        match pda.add_lamports(u64::MAX) {
+            Err(e) => assert_eq!(e, ProgramError::ArithmeticOverflow.into()),
+            _ => unreachable!(),
+        }
+
+        match pda.sub_lamports(u64::MAX) {
+            Err(e) => assert_eq!(e, ProgramError::ArithmeticOverflow.into()),
+            _ => unreachable!(),
+        }
+
+        Ok(())
+    }
 }
 
 #[derive(Accounts)]
-pub struct TestLamportsTrait<'info> {
+pub struct Transfer<'info> {
     #[account(mut)]
     pub signer: Signer<'info>,
-
     #[account(
         init,
         payer = signer,
@@ -67,9 +83,14 @@ pub struct TestLamportsTrait<'info> {
         bump
     )]
     pub pda: Account<'info, LamportsPda>,
-
     pub system_program: Program<'info, System>,
 }
 
+#[derive(Accounts)]
+pub struct Overflow<'info> {
+    #[account(seeds = [b"lamports"], bump)]
+    pub pda: Account<'info, LamportsPda>,
+}
+
 #[account]
 pub struct LamportsPda {}
diff --git a/tests/misc/tests/lamports/lamports.ts b/tests/misc/tests/lamports/lamports.ts
index e5c54dadb..c2e52f59f 100644
--- a/tests/misc/tests/lamports/lamports.ts
+++ b/tests/misc/tests/lamports/lamports.ts
@@ -8,16 +8,13 @@ describe("lamports", () => {
 
   const program = anchor.workspace.Lamports as anchor.Program<Lamports>;
 
-  it("Can use the Lamports trait", async () => {
-    const signer = program.provider.publicKey!;
-    const [pda] = anchor.web3.PublicKey.findProgramAddressSync(
-      [Buffer.from("lamports")],
-      program.programId
-    );
-
+  it("Can transfer from/to PDA", async () => {
     await program.methods
-      .testLamportsTrait(new anchor.BN(anchor.web3.LAMPORTS_PER_SOL))
-      .accounts({ signer, pda })
+      .transfer(new anchor.BN(anchor.web3.LAMPORTS_PER_SOL))
       .rpc();
   });
+
+  it("Returns an error on overflow", async () => {
+    await program.methods.overflow().rpc();
+  });
 });