From 39acc86718c62a34ad352267afb906fbf427427d Mon Sep 17 00:00:00 2001
From: Onyewuchi Emeka <marrnuel123@.com>
Date: Mon, 7 Oct 2024 01:15:41 +0100
Subject: [PATCH] Refactor SOL transfer program for improved security and
 clarity

- Added custom error handling (CustomError) for insufficient funds
- Implemented balance checks before transferring SOL
---
 .../anchor/programs/transfer-sol/src/lib.rs      | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/basics/transfer-sol/anchor/programs/transfer-sol/src/lib.rs b/basics/transfer-sol/anchor/programs/transfer-sol/src/lib.rs
index 26d15e578..794ab0e70 100644
--- a/basics/transfer-sol/anchor/programs/transfer-sol/src/lib.rs
+++ b/basics/transfer-sol/anchor/programs/transfer-sol/src/lib.rs
@@ -9,7 +9,13 @@ declare_id!("4fQVnLWKKKYxtxgGn7Haw8v2g2Hzbu8K61JvWKvqAi7W");
 pub mod transfer_sol {
     use super::*;
 
+    /// Transfer SOL using Cross-Program Invocation (CPI)
     pub fn transfer_sol_with_cpi(ctx: Context<TransferSolWithCpi>, amount: u64) -> Result<()> {
+        // Check for sufficient balance before transferring
+        let payer_balance = **ctx.accounts.payer.try_borrow_lamports()?;
+        require!(payer_balance >= amount, CustomError::InsufficientFunds);
+
+        // CPI-based transfer using the system program's transfer function
         system_program::transfer(
             CpiContext::new(
                 ctx.accounts.system_program.to_account_info(),
@@ -29,6 +35,10 @@ pub mod transfer_sol {
         ctx: Context<TransferSolWithProgram>,
         amount: u64,
     ) -> Result<()> {
+        // Check if the payer has sufficient funds for the transfer
+        let payer_balance = **ctx.accounts.payer.try_borrow_lamports()?;
+        require!(payer_balance >= amount, CustomError::InsufficientFunds);
+
         **ctx.accounts.payer.try_borrow_mut_lamports()? -= amount;
         **ctx.accounts.recipient.try_borrow_mut_lamports()? += amount;
         Ok(())
@@ -55,3 +65,9 @@ pub struct TransferSolWithProgram<'info> {
     #[account(mut)]
     recipient: SystemAccount<'info>,
 }
+
+#[error_code]
+pub enum CustomError {
+    #[msg("Insufficient funds for the transfer.")]
+    InsufficientFunds,
+}