diff --git a/pinocchio/program/src/processor/unwrap_lamports.rs b/pinocchio/program/src/processor/unwrap_lamports.rs index 30e0f5d2..163e268c 100644 --- a/pinocchio/program/src/processor/unwrap_lamports.rs +++ b/pinocchio/program/src/processor/unwrap_lamports.rs @@ -60,30 +60,30 @@ pub fn process_unwrap_lamports(accounts: &[AccountInfo], instruction_data: &[u8] (source_account.amount(), 0) }; - // Comparing whether the AccountInfo's "point" to the same account or - // not - this is a faster comparison since it just checks the internal - // raw pointer. - let self_transfer = source_account_info == destination_account_info; - - if unlikely(self_transfer || amount == 0) { + if unlikely(amount == 0) { // Validates the token account owner since we are not writing // to the account. check_account_owner(source_account_info) } else { source_account.set_amount(remaining_amount); - // SAFETY: single mutable borrow to `source_account_info` lamports. - let source_lamports = unsafe { source_account_info.borrow_mut_lamports_unchecked() }; - // Note: The amount of a source token account is already validated and the - // `lamports` on the account is always greater than `amount`. - *source_lamports -= amount; + // Comparing whether the AccountInfo's "point" to the same account or + // not - this is a faster comparison since it just checks the internal + // raw pointer. + if source_account_info != destination_account_info { + // SAFETY: single mutable borrow to `source_account_info` lamports. + let source_lamports = unsafe { source_account_info.borrow_mut_lamports_unchecked() }; + // Note: The amount of a source token account is already validated and the + // `lamports` on the account is always greater than `amount`. + *source_lamports -= amount; - // SAFETY: single mutable borrow to `destination_account_info` lamports; the - // account is already validated to be different from `source_account_info`. - let destination_lamports = - unsafe { destination_account_info.borrow_mut_lamports_unchecked() }; - // Note: The total lamports supply is bound to `u64::MAX`. - *destination_lamports += amount; + // SAFETY: single mutable borrow to `destination_account_info` lamports; the + // account is already validated to be different from `source_account_info`. + let destination_lamports = + unsafe { destination_account_info.borrow_mut_lamports_unchecked() }; + // Note: The total lamports supply is bound to `u64::MAX`. + *destination_lamports += amount; + } Ok(()) } diff --git a/pinocchio/program/tests/unwrap_lamports.rs b/pinocchio/program/tests/unwrap_lamports.rs index 2acc8ff5..af4ab8bf 100644 --- a/pinocchio/program/tests/unwrap_lamports.rs +++ b/pinocchio/program/tests/unwrap_lamports.rs @@ -409,7 +409,8 @@ fn unwrap_lamports_with_self_transfer() { .unwrap(); // It should succeed to unwrap lamports with the same source and destination - // accounts. + // accounts. The amount should be deducted from the source account but the + // lamports should remain the same. let result = mollusk().process_and_validate_instruction( &instruction, @@ -432,7 +433,7 @@ fn unwrap_lamports_with_self_transfer() { let account = account.unwrap(); let token_account = spl_token_interface::state::Account::unpack(&account.data).unwrap(); - assert_eq!(token_account.amount, 2_000_000_000); + assert_eq!(token_account.amount, 1_000_000_000); } #[test]