fix: fixed undelegation of ephemeral balance

src/discriminator.rs

@@ -33,6 +33,10 @@ pub enum DlpDiscriminator {
     CommitStateFromBuffer = 13,
     /// See [crate::processor::process_close_validator_fees_vault] for docs.
     CloseValidatorFeesVault = 14,
+    /// See [crate::processor::process_undelegate_ephemeral_balance] for docs.
+    UndelegateEphemeralBalance = 15,
+    /// See [crate::processor::process_close_ephemeral_balance_v1] for docs.
+    CloseEphemeralBalanceV1 = 16
 }
 
 impl DlpDiscriminator {
@@ -45,7 +49,8 @@ impl DlpDiscriminator {
 impl TryFrom<[u8; 8]> for DlpDiscriminator {
     type Error = ProgramError;
     fn try_from(bytes: [u8; 8]) -> Result<Self, Self::Error> {
-        match bytes[0] {
+        let discriminator = u64::from_le_bytes(bytes);
+        match discriminator {
             0x0 => Ok(DlpDiscriminator::Delegate),
             0x1 => Ok(DlpDiscriminator::CommitState),
             0x2 => Ok(DlpDiscriminator::Finalize),
@@ -60,6 +65,8 @@ impl TryFrom<[u8; 8]> for DlpDiscriminator {
             0xc => Ok(DlpDiscriminator::ProtocolClaimFees),
             0xd => Ok(DlpDiscriminator::CommitStateFromBuffer),
             0xe => Ok(DlpDiscriminator::CloseValidatorFeesVault),
+            0xf => Ok(DlpDiscriminator::UndelegateEphemeralBalance),
+            0x10 => Ok(DlpDiscriminator::CloseEphemeralBalanceV1),
             _ => Err(ProgramError::InvalidInstructionData),
         }
     }

src/instruction_builder/close_ephemeral_balance.rs

@@ -1,21 +1,23 @@
 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;
 
 /// Creates instruction to close an ephemeral balance account
-/// See [crate::processor::process_close_ephemeral_balance] for docs.
+/// See [crate::processor::process_close_ephemeral_balance_v1] for docs.
+/// [crate::processor::process_close_ephemeral_balance] now deprecated
 pub fn close_ephemeral_balance(payer: Pubkey, index: u8) -> Instruction {
     let ephemeral_balance_pda = ephemeral_balance_pda_from_payer(&payer, index);
     Instruction {
         program_id: crate::id(),
         accounts: vec![
             AccountMeta::new(payer, true),
             AccountMeta::new(ephemeral_balance_pda, false),
+            AccountMeta::new_readonly(system_program::id(), false),
         ],
         data: [
-            DlpDiscriminator::CloseEphemeralBalance.to_vec(),
+            DlpDiscriminator::CloseEphemeralBalanceV1.to_vec(),
             vec![index],
         ]
         .concat(),

src/instruction_builder/mod.rs

@@ -11,6 +11,7 @@ mod init_validator_fees_vault;
 mod protocol_claim_fees;
 mod top_up_ephemeral_balance;
 mod undelegate;
+mod undelegate_ephemeral_balance;
 mod validator_claim_fees;
 mod whitelist_validator_for_program;
 
@@ -26,5 +27,6 @@ pub use init_validator_fees_vault::*;
 pub use protocol_claim_fees::*;
 pub use top_up_ephemeral_balance::*;
 pub use undelegate::*;
+pub use undelegate_ephemeral_balance::*;
 pub use validator_claim_fees::*;
 pub use whitelist_validator_for_program::*;

src/instruction_builder/undelegate_ephemeral_balance.rs

@@ -0,0 +1,17 @@
+use crate::discriminator::DlpDiscriminator;
+use crate::instruction_builder::undelegate;
+use solana_program::instruction::Instruction;
+use solana_program::pubkey::Pubkey;
+
+/// Builds an undelegate instruction for ephemeral balance.
+/// See [crate::processor::process_undelegate_ephemeral_balance] for docs.
+pub fn undelegate_ephemeral_balance(
+    validator: Pubkey,
+    delegated_account: Pubkey,
+    rent_reimbursement: Pubkey,
+) -> Instruction {
+    let mut ix = undelegate(validator, delegated_account, crate::ID, rent_reimbursement);
+    ix.data = DlpDiscriminator::UndelegateEphemeralBalance.to_vec();
+
+    ix
+}

src/lib.rs

@@ -93,6 +93,12 @@ pub fn process_instruction(
         discriminator::DlpDiscriminator::CloseValidatorFeesVault => {
             processor::process_close_validator_fees_vault(program_id, accounts, data)?
         }
+        discriminator::DlpDiscriminator::UndelegateEphemeralBalance => {
+            processor::process_undelegate_ephemeral_balance(program_id, accounts, data)?
+        }
+        discriminator::DlpDiscriminator::CloseEphemeralBalanceV1 => {
+            processor::process_close_ephemeral_balance_v1(program_id, accounts, data)?
+        }
     }
     Ok(())
 }

src/processor/close_ephemeral_balance_v1.rs

@@ -0,0 +1,76 @@
+use crate::ephemeral_balance_seeds_from_payer;
+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::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:
+///
+/// - ephemeral balance account is initialized
+///
+/// Steps:
+///
+/// 1. Closes the ephemeral balance account and refunds the payer with the
+///    escrowed lamports
+pub fn process_close_ephemeral_balance_v1(
+    _program_id: &Pubkey,
+    accounts: &[AccountInfo],
+    data: &[u8],
+) -> ProgramResult {
+    let index = *data.first().ok_or(ProgramError::InvalidInstructionData)?;
+
+    // Load Accounts
+    let [payer, ephemeral_balance_account, system_program] = accounts else {
+        return Err(ProgramError::NotEnoughAccountKeys);
+    };
+
+    load_signer(payer, "payer")?;
+
+    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,
+        &crate::id(),
+        true,
+        "ephemeral balance",
+    )?;
+    if ephemeral_balance_account.owner != &system_program::id() {
+        msg!(
+            "ephemeral balance expected to be owned by system program: {}",
+            system_program::id()
+        );
+        return Err(ProgramError::InvalidAccountOwner);
+    }
+
+    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/mod.rs

@@ -1,4 +1,5 @@
 mod close_ephemeral_balance;
+mod close_ephemeral_balance_v1;
 mod close_validator_fees_vault;
 mod commit_state;
 mod commit_state_from_buffer;
@@ -10,11 +11,13 @@ mod init_validator_fees_vault;
 mod protocol_claim_fees;
 mod top_up_ephemeral_balance;
 mod undelegate;
+mod undelegate_ephemeral_balance;
 mod utils;
 mod validator_claim_fees;
 mod whitelist_validator_for_program;
 
 pub use close_ephemeral_balance::*;
+pub use close_ephemeral_balance_v1::*;
 pub use close_validator_fees_vault::*;
 pub use commit_state::*;
 pub use commit_state_from_buffer::*;
@@ -26,5 +29,6 @@ pub use init_validator_fees_vault::*;
 pub use protocol_claim_fees::*;
 pub use top_up_ephemeral_balance::*;
 pub use undelegate::*;
+pub use undelegate_ephemeral_balance::*;
 pub use validator_claim_fees::*;
 pub use whitelist_validator_for_program::*;

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/processor/undelegate_ephemeral_balance.rs

@@ -0,0 +1,74 @@
+use crate::instruction_builder::undelegate;
+use solana_program::msg;
+use solana_program::program::invoke;
+use solana_program::program_error::ProgramError;
+use solana_program::{
+    account_info::AccountInfo, entrypoint::ProgramResult, pubkey::Pubkey, system_program,
+};
+
+/// Undelegate ephemeral balance
+///
+/// Accounts:
+///
+///  0: `[signer]`   the validator account
+///  1: `[writable]` the delegated account
+///  2: `[]`         the owner program of the delegated account
+///  3: `[writable]` the undelegate buffer PDA we use to store the data temporarily
+///  4: `[]`         the commit state PDA
+///  5: `[]`         the commit record PDA
+///  6: `[writable]` the delegation record PDA
+///  7: `[writable]` the delegation metadata PDA
+///  8: `[]`         the rent reimbursement account
+///  9: `[writable]` the protocol fees vault account
+/// 10: `[writable]` the validator fees vault account
+/// 11: `[]`         the system program
+///
+/// Requirements:
+///
+/// - delegated account is owned by delegation program
+/// - delegation record is initialized
+/// - delegation metadata is initialized
+/// - protocol fees vault is initialized
+/// - validator fees vault is initialized
+/// - commit state is uninitialized
+/// - commit record is uninitialized
+/// - delegated account is NOT undelegatable
+/// - owner program account matches the owner in the delegation record
+/// - rent reimbursement account matches the rent payer in the delegation metadata
+///
+/// Steps:
+///
+/// - Undelegate using CPI into [`crate::processor::undelegate`]
+/// - Assigns ownership back to system program
+pub fn process_undelegate_ephemeral_balance(
+    _program_id: &Pubkey,
+    accounts: &[AccountInfo],
+    _data: &[u8],
+) -> ProgramResult {
+    let [validator, delegated_account, owner_program, _, _, _, _, _, rent_reimbursement, _, _, _] =
+        accounts
+    else {
+        return Err(ProgramError::NotEnoughAccountKeys);
+    };
+
+    if owner_program.key != &crate::ID {
+        msg!(
+            "Unexpected owner program. expected dlp, got: {}",
+            owner_program.key
+        );
+        return Err(ProgramError::IncorrectProgramId);
+    }
+
+    // Propagate to undelegate which also runs all necessary checks.
+    let undelegate_ix = undelegate(
+        *validator.key,
+        *delegated_account.key,
+        *owner_program.key,
+        *rent_reimbursement.key,
+    );
+    invoke(&undelegate_ix, accounts)?;
+
+    // Assign ownership back to system_program
+    delegated_account.assign(&system_program::ID);
+    Ok(())
+}