Make amount optional

Author: febo

p-interface/src/instruction.rs

@@ -510,7 +510,9 @@ pub enum TokenInstruction {
     ///
     /// Data expected by this instruction:
     ///
-    ///   - `u64` The amount of lamports to transfer.
+    ///   - `Option<u64>` The amount of lamports to transfer. When an amount is
+    ///     not specified, the entire balance of the source account will be
+    ///     transferred.
     UnwrapLamports = 45,
 
     /// Executes a batch of instructions. The instructions to be executed are

p-token/src/processor/unwrap_lamports.rs

@@ -1,17 +1,31 @@
 use {
     super::validate_owner,
-    crate::processor::{check_account_owner, unpack_amount},
+    crate::processor::{check_account_owner, U64_BYTES},
     pinocchio::{account_info::AccountInfo, program_error::ProgramError, ProgramResult},
     pinocchio_token_interface::{
         error::TokenError,
+        likely,
         state::{account::Account, load_mut},
     },
 };
 
 #[allow(clippy::arithmetic_side_effects)]
 pub fn process_unwrap_lamports(accounts: &[AccountInfo], instruction_data: &[u8]) -> ProgramResult {
-    // Amount being unwrapped.
-    let amount = unpack_amount(instruction_data)?;
+    // instruction data: expected u8 (1) + optional u64 (8)
+    let [has_amount, maybe_amount @ ..] = instruction_data else {
+        return Err(TokenError::InvalidInstruction.into());
+    };
+
+    let maybe_amount = if likely(*has_amount == 0) {
+        None
+    } else if maybe_amount.len() >= 8 {
+        // SAFETY: The slice is guaranteed to be at least 8 bytes long.
+        Some(u64::from_le_bytes(unsafe {
+            *(maybe_amount.as_ptr() as *const [u8; U64_BYTES])
+        }))
+    } else {
+        return Err(TokenError::InvalidInstruction.into());
+    };
 
     let [source_account_info, destination_account_info, authority_info, remaining @ ..] = accounts
     else {
@@ -31,10 +45,19 @@ pub fn process_unwrap_lamports(accounts: &[AccountInfo], instruction_data: &[u8]
     // a multisig.
     unsafe { validate_owner(&source_account.owner, authority_info, remaining)? };
 
-    let remaining_amount = source_account
-        .amount()
-        .checked_sub(amount)
-        .ok_or(TokenError::InsufficientFunds)?;
+    // If we have an amount, we need to validate whether there are enough lamports
+    // to unwrap or not; otherwise we just use the full amount.
+    let (amount, remaining_amount) = if let Some(amount) = maybe_amount {
+        (
+            amount,
+            source_account
+                .amount()
+                .checked_sub(amount)
+                .ok_or(TokenError::InsufficientFunds)?,
+        )
+    } else {
+        (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
@@ -44,7 +67,7 @@ pub fn process_unwrap_lamports(accounts: &[AccountInfo], instruction_data: &[u8]
     if self_transfer || amount == 0 {
         // Validates the token account owner since we are not writing
         // to the account.
-        check_account_owner(source_account_info)?;
+        check_account_owner(source_account_info)
     } else {
         source_account.set_amount(remaining_amount);
 
@@ -55,13 +78,12 @@ pub fn process_unwrap_lamports(accounts: &[AccountInfo], instruction_data: &[u8]
         *source_lamports -= amount;
 
         // SAFETY: single mutable borrow to `destination_account_info` lamports; the
-        // account is already validated to be different from
-        // `source_account_info`.
+        // 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(())
+        Ok(())
+    }
 }

p-token/tests/unwrap_lamports.rs

@@ -53,11 +53,23 @@ fn create_token_account(
     }
 }
 
+/// Creates a Mollusk instance with the default feature set, excluding the
+/// `bpf_account_data_direct_mapping` feature.
+fn mollusk() -> Mollusk {
+    let mut mollusk = Mollusk::default();
+    mollusk.add_program(
+        &TOKEN_PROGRAM_ID,
+        "pinocchio_token_program",
+        &bpf_loader_upgradeable::id(),
+    );
+    mollusk
+}
+
 fn unwrap_lamports_instruction(
     source: &Pubkey,
     destination: &Pubkey,
     authority: &Pubkey,
-    amount: u64,
+    amount: Option<u64>,
 ) -> Result<Instruction, ProgramError> {
     let accounts = vec![
         AccountMeta::new(*source, false),
@@ -67,7 +79,13 @@ fn unwrap_lamports_instruction(
 
     // Start with the batch discriminator
     let mut data: Vec<u8> = vec![TokenInstruction::UnwrapLamports as u8];
-    data.extend_from_slice(&amount.to_le_bytes());
+
+    if let Some(amount) = amount {
+        data.push(1);
+        data.extend_from_slice(&amount.to_le_bytes());
+    } else {
+        data.push(0);
+    }
 
     Ok(Instruction {
         program_id: spl_token::ID,
@@ -76,20 +94,63 @@ fn unwrap_lamports_instruction(
     })
 }
 
-/// Creates a Mollusk instance with the default feature set, excluding the
-/// `bpf_account_data_direct_mapping` feature.
-fn mollusk() -> Mollusk {
-    let mut mollusk = Mollusk::default();
-    mollusk.add_program(
+#[tokio::test]
+async fn unwrap_lamports() {
+    let native_mint = Pubkey::new_from_array(native_mint::ID);
+    let authority_key = Pubkey::new_unique();
+    let destination_account_key = Pubkey::new_unique();
+
+    // native account:
+    //   - amount: 2_000_000_000
+    let source_account_key = Pubkey::new_unique();
+    let source_account = create_token_account(
+        &native_mint,
+        &authority_key,
+        true,
+        2_000_000_000,
         &TOKEN_PROGRAM_ID,
-        "pinocchio_token_program",
-        &bpf_loader_upgradeable::id(),
     );
-    mollusk
+
+    let instruction = unwrap_lamports_instruction(
+        &source_account_key,
+        &destination_account_key,
+        &authority_key,
+        None,
+    )
+    .unwrap();
+
+    // It should succeed to unwrap 2_000_000_000 lamports.
+
+    let result = mollusk().process_and_validate_instruction(
+        &instruction,
+        &[
+            (source_account_key, source_account),
+            (destination_account_key, Account::default()),
+            (authority_key, Account::default()),
+        ],
+        &[
+            Check::success(),
+            Check::account(&destination_account_key)
+                .lamports(2_000_000_000)
+                .build(),
+            Check::account(&source_account_key)
+                .lamports(Rent::default().minimum_balance(size_of::<TokenAccount>()))
+                .build(),
+        ],
+    );
+
+    // And the remaining amount must be 0.
+
+    result.resulting_accounts.iter().for_each(|(key, account)| {
+        if *key == source_account_key {
+            let token_account = spl_token::state::Account::unpack(&account.data).unwrap();
+            assert_eq!(token_account.amount, 0);
+        }
+    });
 }
 
 #[tokio::test]
-async fn unwrap_lamports() {
+async fn unwrap_lamports_with_amount() {
     let native_mint = Pubkey::new_from_array(native_mint::ID);
     let authority_key = Pubkey::new_unique();
     let destination_account_key = Pubkey::new_unique();
@@ -109,7 +170,7 @@ async fn unwrap_lamports() {
         &source_account_key,
         &destination_account_key,
         &authority_key,
-        2_000_000_000,
+        Some(2_000_000_000),
     )
     .unwrap();
 
@@ -164,7 +225,7 @@ async fn fail_unwrap_lamports_with_insufficient_funds() {
         &source_account_key,
         &destination_account_key,
         &authority_key,
-        2_000_000_000,
+        Some(2_000_000_000),
     )
     .unwrap();
 
@@ -205,7 +266,7 @@ async fn unwrap_lamports_with_parial_amount() {
         &source_account_key,
         &destination_account_key,
         &authority_key,
-        1_000_000_000,
+        Some(1_000_000_000),
     )
     .unwrap();
 
@@ -263,7 +324,7 @@ async fn fail_unwrap_lamports_with_invalid_authority() {
         &source_account_key,
         &destination_account_key,
         &fake_authority_key, // <-- wrong authority
-        2_000_000_000,
+        Some(2_000_000_000),
     )
     .unwrap();
 
@@ -305,7 +366,7 @@ async fn fail_unwrap_lamports_with_non_native_account() {
         &source_account_key,
         &destination_account_key,
         &authority_key,
-        1_000_000_000,
+        Some(1_000_000_000),
     )
     .unwrap();
 
@@ -353,7 +414,7 @@ async fn unwrap_lamports_with_self_transfer() {
         &source_account_key,
         &source_account_key, // <-- destination same as source
         &authority_key,
-        1_000_000_000,
+        Some(1_000_000_000),
     )
     .unwrap();
 
@@ -407,7 +468,7 @@ async fn fail_unwrap_lamports_with_invalid_native_account() {
         &source_account_key,
         &destination_account_key,
         &authority_key,
-        1_000_000_000,
+        Some(1_000_000_000),
     )
     .unwrap();