Fix/ephemeral balance undelegation 2.0

src/instruction_builder/close_ephemeral_balance.rs

@@ -1,5 +1,5 @@
 use solana_program::instruction::Instruction;
-use solana_program::{instruction::AccountMeta, pubkey::Pubkey};
+use solana_program::{instruction::AccountMeta, pubkey::Pubkey, system_program};
 
 use crate::discriminator::DlpDiscriminator;
 use crate::pda::ephemeral_balance_pda_from_payer;
@@ -13,6 +13,7 @@ pub fn close_ephemeral_balance(payer: Pubkey, index: u8) -> Instruction {
         accounts: vec![
             AccountMeta::new(payer, true),
             AccountMeta::new(ephemeral_balance_pda, false),
+            AccountMeta::new_readonly(system_program::id(), false),
         ],
         data: [
             DlpDiscriminator::CloseEphemeralBalance.to_vec(),

src/instruction_builder/delegate_ephemeral_balance.rs

@@ -21,7 +21,7 @@ pub fn delegate_ephemeral_balance(
     let delegated_account = ephemeral_balance_pda_from_payer(&pubkey, args.index);
     let delegate_buffer_pda = delegate_buffer_pda_from_delegated_account_and_owner_program(
         &delegated_account,
-        &crate::id(),
+        &system_program::id(),
     );
     let delegation_record_pda = delegation_record_pda_from_delegated_account(&delegated_account);
     let delegation_metadata_pda =

src/processor/close_ephemeral_balance.rs

@@ -1,15 +1,20 @@
 use crate::ephemeral_balance_seeds_from_payer;
-use crate::processor::utils::loaders::{load_initialized_pda, load_signer};
-use crate::processor::utils::pda::close_pda;
+use crate::processor::utils::loaders::{load_pda, load_signer};
+use solana_program::msg;
+use solana_program::program::invoke_signed;
 use solana_program::program_error::ProgramError;
-use solana_program::{account_info::AccountInfo, entrypoint::ProgramResult, pubkey::Pubkey};
+use solana_program::system_instruction::transfer;
+use solana_program::{
+    account_info::AccountInfo, entrypoint::ProgramResult, pubkey::Pubkey, system_program,
+};
 
 /// Process the closing of an ephemeral balance account
 ///
 /// Accounts:
 ///
 /// 0: `[signer]` payer to pay for the transaction and receive the refund
 /// 1: `[writable]` ephemeral balance account we are closing
+/// 2: `[]` the system program
 ///
 /// Requirements:
 ///
@@ -27,21 +32,45 @@ pub fn process_close_ephemeral_balance(
     let index = *data.first().ok_or(ProgramError::InvalidInstructionData)?;
 
     // Load Accounts
-    let [payer, ephemeral_balance_account] = accounts else {
+    let [payer, ephemeral_balance_account, system_program] = accounts else {
         return Err(ProgramError::NotEnoughAccountKeys);
     };
 
     load_signer(payer, "payer")?;
 
-    load_initialized_pda(
+    let ephemeral_balance_seeds: &[&[u8]] = ephemeral_balance_seeds_from_payer!(payer.key, index);
+    let ephemeral_balance_bump = load_pda(
         ephemeral_balance_account,
-        ephemeral_balance_seeds_from_payer!(payer.key, index),
+        ephemeral_balance_seeds,
         &crate::id(),
         true,
         "ephemeral balance",
     )?;
+    if ephemeral_balance_account.owner != &system_program::id() {
+        msg!(
+            "ephemeral balance expected to be owned by system program. got: {}",
+            ephemeral_balance_account.owner
+        );
+        return Err(ProgramError::InvalidAccountOwner);
+    }
 
-    close_pda(ephemeral_balance_account, payer)?;
+    let amount = ephemeral_balance_account.lamports();
+    if amount == 0 {
+        return Ok(());
+    }
+
+    let ephemeral_balance_bump_slice: &[u8] = &[ephemeral_balance_bump];
+    let ephemeral_balance_signer_seeds =
+        [ephemeral_balance_seeds, &[ephemeral_balance_bump_slice]].concat();
+    invoke_signed(
+        &transfer(ephemeral_balance_account.key, payer.key, amount),
+        &[
+            ephemeral_balance_account.clone(),
+            payer.clone(),
+            system_program.clone(),
+        ],
+        &[&ephemeral_balance_signer_seeds],
+    )?;
 
     Ok(())
 }

src/processor/delegate.rs

@@ -63,11 +63,21 @@ pub fn process_delegate(
     load_owned_pda(delegated_account, &crate::id(), "delegated account")?;
     load_program(system_program, system_program::id(), "system program")?;
 
+    msg!("Delegating: {}", delegated_account.key);
+
     // Validate seeds if the delegate account is not on curve, i.e. is a PDA
+    // If the owner is the system program, we check if the account is derived from the delegation program,
+    // allowing delegation of escrow accounts
     if !is_on_curve(delegated_account.key) {
         let seeds_to_validate: Vec<&[u8]> = args.seeds.iter().map(|v| v.as_slice()).collect();
+        let program_id = if owner_program.key.eq(&system_program::id()) {
+            crate::id()
+        } else {
+            *owner_program.key
+        };
         let (derived_pda, _) =
-            Pubkey::find_program_address(seeds_to_validate.as_ref(), owner_program.key);
+            Pubkey::find_program_address(seeds_to_validate.as_ref(), &program_id);
+
         if derived_pda.ne(delegated_account.key) {
             msg!(
                 "Expected delegated PDA to be {}, but got {}",

src/processor/delegate_ephemeral_balance.rs

@@ -76,7 +76,7 @@ pub fn process_delegate_ephemeral_balance(
     let ix = crate::instruction_builder::delegate(
         *payer.key,
         *ephemeral_balance_account.key,
-        Some(crate::id()),
+        Some(system_program::id()),
         args.delegate_args,
     );
 

src/processor/top_up_ephemeral_balance.rs

@@ -56,7 +56,7 @@ pub fn process_top_up_ephemeral_balance(
         create_pda(
             ephemeral_balance_account,
             &system_program::id(),
-            8,
+            0,
             ephemeral_balance_seeds_from_payer!(pubkey.key, args.index),
             bump_ephemeral_balance,
             system_program,

src/state/utils/discriminator.rs

@@ -10,7 +10,7 @@ pub enum AccountDiscriminator {
 }
 
 impl AccountDiscriminator {
-    pub fn to_bytes(&self) -> [u8; 8] {
+    pub const fn to_bytes(&self) -> [u8; 8] {
         let num = (*self) as u64;
         num.to_le_bytes()
     }

src/state/utils/to_bytes.rs

@@ -23,10 +23,10 @@ macro_rules! impl_to_bytes_with_discriminator_borsh {
         impl $struct_name {
             pub fn to_bytes_with_discriminator<W: std::io::Write>(
                 &self,
-                data: &mut W,
+                writer: &mut W,
             ) -> Result<(), ::solana_program::program_error::ProgramError> {
-                data.write_all(&Self::discriminator().to_bytes())?;
-                self.serialize(data)?;
+                writer.write_all(&Self::discriminator().to_bytes())?;
+                self.serialize(writer)?;
                 Ok(())
             }
         }