From cc8f6a8d1215a983419f151e8c7158175b34dc54 Mon Sep 17 00:00:00 2001 From: taco-paco Date: Tue, 6 May 2025 13:43:16 +0900 Subject: [PATCH 1/3] fix: fixed undelegation of ephemeral balance. Now undelegated to system program. Added new version of close_ephemeral_balance to support it. kept old one for backward-compatability --- src/discriminator.rs | 9 +- .../close_ephemeral_balance.rs | 8 +- src/instruction_builder/mod.rs | 2 + .../undelegate_ephemeral_balance.rs | 18 +++ src/lib.rs | 6 + src/processor/close_ephemeral_balance_v1.rs | 72 +++++++++++ src/processor/mod.rs | 4 + src/processor/top_up_ephemeral_balance.rs | 2 +- src/processor/undelegate_ephemeral_balance.rs | 74 +++++++++++ tests/test_top_up.rs | 120 ++++++++++++------ 10 files changed, 274 insertions(+), 41 deletions(-) create mode 100644 src/instruction_builder/undelegate_ephemeral_balance.rs create mode 100644 src/processor/close_ephemeral_balance_v1.rs create mode 100644 src/processor/undelegate_ephemeral_balance.rs diff --git a/src/discriminator.rs b/src/discriminator.rs index e896fa2..197a1b9 100644 --- a/src/discriminator.rs +++ b/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 { - 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), } } diff --git a/src/instruction_builder/close_ephemeral_balance.rs b/src/instruction_builder/close_ephemeral_balance.rs index dece09b..153fdfa 100644 --- a/src/instruction_builder/close_ephemeral_balance.rs +++ b/src/instruction_builder/close_ephemeral_balance.rs @@ -1,11 +1,12 @@ 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 { @@ -13,9 +14,10 @@ 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(), + DlpDiscriminator::CloseEphemeralBalanceV1.to_vec(), vec![index], ] .concat(), diff --git a/src/instruction_builder/mod.rs b/src/instruction_builder/mod.rs index e0cae79..d174d6d 100644 --- a/src/instruction_builder/mod.rs +++ b/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::*; diff --git a/src/instruction_builder/undelegate_ephemeral_balance.rs b/src/instruction_builder/undelegate_ephemeral_balance.rs new file mode 100644 index 0000000..fec70f1 --- /dev/null +++ b/src/instruction_builder/undelegate_ephemeral_balance.rs @@ -0,0 +1,18 @@ +use crate::discriminator::DlpDiscriminator; +use crate::instruction_builder::undelegate; +use solana_program::instruction::Instruction; +use solana_program::pubkey::Pubkey; + +/// Builds an undelegate instruction. +/// See [crate::processor::process_undelegate] for docs. +#[allow(clippy::too_many_arguments)] +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 +} diff --git a/src/lib.rs b/src/lib.rs index 4b1eca8..a8fc29a 100644 --- a/src/lib.rs +++ b/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(()) } diff --git a/src/processor/close_ephemeral_balance_v1.rs b/src/processor/close_ephemeral_balance_v1.rs new file mode 100644 index 0000000..3e371c3 --- /dev/null +++ b/src/processor/close_ephemeral_balance_v1.rs @@ -0,0 +1,72 @@ +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(); + 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(()) +} diff --git a/src/processor/mod.rs b/src/processor/mod.rs index e9b25db..43eda89 100644 --- a/src/processor/mod.rs +++ b/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::*; diff --git a/src/processor/top_up_ephemeral_balance.rs b/src/processor/top_up_ephemeral_balance.rs index cca1e38..ccd0fc6 100644 --- a/src/processor/top_up_ephemeral_balance.rs +++ b/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, diff --git a/src/processor/undelegate_ephemeral_balance.rs b/src/processor/undelegate_ephemeral_balance.rs new file mode 100644 index 0000000..7179b77 --- /dev/null +++ b/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(()) +} diff --git a/tests/test_top_up.rs b/tests/test_top_up.rs index e880bd6..e29525a 100644 --- a/tests/test_top_up.rs +++ b/tests/test_top_up.rs @@ -11,7 +11,7 @@ use solana_program::rent::Rent; use solana_program::{hash::Hash, native_token::LAMPORTS_PER_SOL, system_program}; use solana_program_test::{processor, BanksClient, ProgramTest}; use solana_sdk::{ - account::Account, + account::{Account, ReadableAccount}, signature::{Keypair, Signer}, transaction::Transaction, }; @@ -125,6 +125,47 @@ async fn test_top_up_ephemeral_balance_and_delegate_for_pubkey() { assert!(res.is_ok()); } +#[tokio::test] +async fn test_undelegate() { + // Setup + let (banks, _, payer_alt, blockhash) = setup_program_test_env().await; + let validator = Keypair::from_bytes(&TEST_AUTHORITY).unwrap(); + + let ephemeral_balance_pda = ephemeral_balance_pda_from_payer(&payer_alt.pubkey(), 0); + let ephemeral_balance_owner = banks + .get_account(ephemeral_balance_pda) + .await + .unwrap() + .unwrap() + .owner; + + assert_eq!(ephemeral_balance_owner, dlp::id()); + + // Undelegate ephemeral balance Ix + let ix = dlp::instruction_builder::undelegate_ephemeral_balance( + validator.pubkey(), + ephemeral_balance_pda, + validator.pubkey(), + ); + + let tx = Transaction::new_signed_with_payer( + &[ix], + Some(&validator.pubkey()), + &[&validator], + blockhash, + ); + let res = banks.process_transaction(tx).await; + assert!(res.is_ok()); + + // Assert that the ephemeral balance account is closed + let ephemeral_balance_account = banks.get_account(ephemeral_balance_pda).await.unwrap(); + assert!(ephemeral_balance_account.is_some()); + + let actual_owner = *ephemeral_balance_account.unwrap().owner(); + let expected_owner = system_program::id(); + assert_eq!(actual_owner, expected_owner); +} + #[tokio::test] async fn test_undelegate_and_close() { // Setup @@ -149,10 +190,9 @@ async fn test_undelegate_and_close() { .lamports; // Undelegate ephemeral balance Ix - let ix = dlp::instruction_builder::undelegate( + let ix = dlp::instruction_builder::undelegate_ephemeral_balance( validator.pubkey(), ephemeral_balance_pda, - dlp::id(), validator.pubkey(), ); @@ -201,87 +241,95 @@ async fn setup_program_test_env() -> (BanksClient, Keypair, Keypair, Hash) { }, ); - let ephemeral_balance_pda = ephemeral_balance_pda_from_payer(&payer_alt.pubkey(), 0); + setup_ephemeral_balance(&mut program_test, &validator, &payer_alt).await; - // Setup the delegated account PDA + // Setup the validator keypair program_test.add_account( - ephemeral_balance_pda, + validator.pubkey(), Account { lamports: LAMPORTS_PER_SOL, data: vec![], - owner: dlp::id(), + owner: system_program::id(), executable: false, rent_epoch: 0, }, ); - // Setup the delegated record PDA - let delegation_record_data = - create_delegation_record_data(validator.pubkey(), dlp::id(), Some(LAMPORTS_PER_SOL)); + // Setup the protocol fees vault program_test.add_account( - delegation_record_pda_from_delegated_account(&ephemeral_balance_pda), + fees_vault_pda(), Account { - lamports: Rent::default().minimum_balance(delegation_record_data.len()), - data: delegation_record_data, + lamports: Rent::default().minimum_balance(0), + data: vec![], owner: dlp::id(), executable: false, rent_epoch: 0, }, ); - // Setup the delegated account metadata PDA - let delegation_metadata_data = create_delegation_metadata_data( - validator.pubkey(), - ephemeral_balance_seeds_from_payer!(payer_alt.pubkey(), 0), - true, - ); + // Setup the validator fees vault program_test.add_account( - delegation_metadata_pda_from_delegated_account(&ephemeral_balance_pda), + validator_fees_vault_pda_from_validator(&validator.pubkey()), Account { - lamports: Rent::default().minimum_balance(delegation_metadata_data.len()), - data: delegation_metadata_data, + lamports: LAMPORTS_PER_SOL, + data: vec![], owner: dlp::id(), executable: false, rent_epoch: 0, }, ); - // Setup the validator keypair + let (banks, payer, blockhash) = program_test.start().await; + (banks, payer, payer_alt, blockhash) +} + +async fn setup_ephemeral_balance( + program_test: &mut ProgramTest, + validator: &Keypair, + payer: &Keypair, +) { + let ephemeral_balance_pda = ephemeral_balance_pda_from_payer(&payer.pubkey(), 0); + + // Setup the delegated account PDA program_test.add_account( - validator.pubkey(), + ephemeral_balance_pda, Account { lamports: LAMPORTS_PER_SOL, data: vec![], - owner: system_program::id(), + owner: dlp::id(), executable: false, rent_epoch: 0, }, ); - // Setup the protocol fees vault + // Setup the delegated record PDA + let delegation_record_data = + create_delegation_record_data(validator.pubkey(), dlp::id(), Some(LAMPORTS_PER_SOL)); program_test.add_account( - fees_vault_pda(), + delegation_record_pda_from_delegated_account(&ephemeral_balance_pda), Account { - lamports: Rent::default().minimum_balance(0), - data: vec![], + lamports: Rent::default().minimum_balance(delegation_record_data.len()), + data: delegation_record_data, owner: dlp::id(), executable: false, rent_epoch: 0, }, ); - // Setup the validator fees vault + // Setup the delegated account metadata PDA + let delegation_metadata_data = create_delegation_metadata_data( + validator.pubkey(), + ephemeral_balance_seeds_from_payer!(payer.pubkey(), 0), + true, + ); program_test.add_account( - validator_fees_vault_pda_from_validator(&validator.pubkey()), + delegation_metadata_pda_from_delegated_account(&ephemeral_balance_pda), Account { - lamports: LAMPORTS_PER_SOL, - data: vec![], + lamports: Rent::default().minimum_balance(delegation_metadata_data.len()), + data: delegation_metadata_data, owner: dlp::id(), executable: false, rent_epoch: 0, }, ); - - let (banks, payer, blockhash) = program_test.start().await; - (banks, payer, payer_alt, blockhash) } From f7e71b41cc7f4a1593847a11de58705d912f070e Mon Sep 17 00:00:00 2001 From: taco-paco Date: Tue, 6 May 2025 13:54:46 +0900 Subject: [PATCH 2/3] refactor: skip transfer if amount is 0 --- src/processor/close_ephemeral_balance_v1.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/processor/close_ephemeral_balance_v1.rs b/src/processor/close_ephemeral_balance_v1.rs index 3e371c3..b2b1162 100644 --- a/src/processor/close_ephemeral_balance_v1.rs +++ b/src/processor/close_ephemeral_balance_v1.rs @@ -55,6 +55,10 @@ pub fn process_close_ephemeral_balance_v1( } 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(); From d781c42b4325ea67450d4dfec3ca8161e91f84e0 Mon Sep 17 00:00:00 2001 From: taco-paco Date: Tue, 6 May 2025 14:06:02 +0900 Subject: [PATCH 3/3] refactor: addressed greptile comments --- src/instruction_builder/undelegate_ephemeral_balance.rs | 5 ++--- tests/test_top_up.rs | 2 +- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/src/instruction_builder/undelegate_ephemeral_balance.rs b/src/instruction_builder/undelegate_ephemeral_balance.rs index fec70f1..c924fee 100644 --- a/src/instruction_builder/undelegate_ephemeral_balance.rs +++ b/src/instruction_builder/undelegate_ephemeral_balance.rs @@ -3,9 +3,8 @@ use crate::instruction_builder::undelegate; use solana_program::instruction::Instruction; use solana_program::pubkey::Pubkey; -/// Builds an undelegate instruction. -/// See [crate::processor::process_undelegate] for docs. -#[allow(clippy::too_many_arguments)] +/// 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, diff --git a/tests/test_top_up.rs b/tests/test_top_up.rs index e29525a..6712e0a 100644 --- a/tests/test_top_up.rs +++ b/tests/test_top_up.rs @@ -157,7 +157,7 @@ async fn test_undelegate() { let res = banks.process_transaction(tx).await; assert!(res.is_ok()); - // Assert that the ephemeral balance account is closed + // Assert that the ephemeral balance account still exists but is now owned by the system program let ephemeral_balance_account = banks.get_account(ephemeral_balance_pda).await.unwrap(); assert!(ephemeral_balance_account.is_some());