diff --git a/idl/rewards_program.json b/idl/rewards_program.json index e91ed37..d4d0e4c 100644 --- a/idl/rewards_program.json +++ b/idl/rewards_program.json @@ -94,6 +94,24 @@ "kind": "accountNode", "name": "directDistribution" }, + { + "data": { + "fields": [ + { + "kind": "structFieldTypeNode", + "name": "bump", + "type": { + "endian": "le", + "format": "u8", + "kind": "numberTypeNode" + } + } + ], + "kind": "structTypeNode" + }, + "kind": "accountNode", + "name": "directDistributionClosed" + }, { "data": { "fields": [ @@ -1942,6 +1960,12 @@ "kind": "errorNode", "message": "Nothing to revoke — user has zero balance", "name": "pointsNothingToRevoke" + }, + { + "code": 42, + "kind": "errorNode", + "message": "Distribution has been permanently closed", + "name": "distributionPermanentlyClosed" } ], "instructions": [ @@ -1976,7 +2000,7 @@ }, { "docs": [ - "PDA: [b\"direct_distribution\", mint, authority, seeds] (created)" + "PDA: [b\"direct_distribution\", mint, authority, seeds]. If a previous distribution was closed at this address, the account now holds a DirectDistributionClosed marker and re-creation is rejected." ], "isSigner": false, "isWritable": true, @@ -2395,7 +2419,7 @@ }, { "docs": [ - "PDA: DirectDistribution account (closed)" + "PDA: DirectDistribution account. Flipped to a compact DirectDistributionClosed marker in place; freed rent is refunded to authority." ], "isSigner": false, "isWritable": true, @@ -3314,10 +3338,10 @@ }, { "docs": [ - "PDA: [b\"merkle_claim\", distribution, claimant] (read-only, may not exist)" + "PDA: [b\"merkle_claim\", distribution, claimant] (writable, may not exist)" ], "isSigner": false, - "isWritable": false, + "isWritable": true, "kind": "instructionAccountNode", "name": "claimAccount" }, diff --git a/program/src/errors.rs b/program/src/errors.rs index 3a0c9d9..8daed83 100644 --- a/program/src/errors.rs +++ b/program/src/errors.rs @@ -172,6 +172,10 @@ pub enum RewardsProgramError { /// (41) Nothing to revoke — user has zero balance #[error("Nothing to revoke — user has zero balance")] PointsNothingToRevoke, + + /// (42) Distribution has been permanently closed + #[error("Distribution has been permanently closed")] + DistributionPermanentlyClosed, } impl From for ProgramError { diff --git a/program/src/instructions/definition.rs b/program/src/instructions/definition.rs index 215e2fa..8ab9f07 100644 --- a/program/src/instructions/definition.rs +++ b/program/src/instructions/definition.rs @@ -14,7 +14,7 @@ pub enum RewardsProgramInstruction { #[codama(account( name = "distribution", writable, - docs = "PDA: [b\"direct_distribution\", mint, authority, seeds] (created)" + docs = "PDA: [b\"direct_distribution\", mint, authority, seeds]. If a previous distribution was closed at this address, the account now holds a DirectDistributionClosed marker and re-creation is rejected." ))] #[codama(account(name = "mint", docs = "SPL token mint"))] #[codama(account( @@ -110,7 +110,11 @@ pub enum RewardsProgramInstruction { writable, docs = "Distribution authority; receives rent + remaining distribution vault tokens" ))] - #[codama(account(name = "distribution", writable, docs = "PDA: DirectDistribution account (closed)"))] + #[codama(account( + name = "distribution", + writable, + docs = "PDA: DirectDistribution account. Flipped to a compact DirectDistributionClosed marker in place; freed rent is refunded to authority." + ))] #[codama(account(name = "mint", docs = "SPL token mint"))] #[codama(account( name = "distribution_vault", @@ -324,7 +328,8 @@ pub enum RewardsProgramInstruction { #[codama(account(name = "distribution", writable, docs = "PDA: MerkleDistribution account"))] #[codama(account( name = "claim_account", - docs = "PDA: [b\"merkle_claim\", distribution, claimant] (read-only, may not exist)" + writable, + docs = "PDA: [b\"merkle_claim\", distribution, claimant] (writable, may not exist)" ))] #[codama(account( name = "revocation_marker", diff --git a/program/src/instructions/direct/close_distribution/processor.rs b/program/src/instructions/direct/close_distribution/processor.rs index 7d61197..2429247 100644 --- a/program/src/instructions/direct/close_distribution/processor.rs +++ b/program/src/instructions/direct/close_distribution/processor.rs @@ -6,7 +6,7 @@ use crate::{ events::DistributionClosedEvent, state::DirectDistribution, traits::{Distribution, DistributionSigner, EventSerialize, InstructionData}, - utils::{close_pda_account, emit_event, get_current_timestamp, get_mint_decimals, get_token_account_balance}, + utils::{emit_event, get_current_timestamp, get_mint_decimals, get_token_account_balance}, ID, }; @@ -24,6 +24,7 @@ pub fn process_close_direct_distribution( let distribution = DirectDistribution::from_account(&distribution_data, ix.accounts.distribution, &ID)?; distribution.validate_authority(ix.accounts.authority.address())?; distribution.validate_mint(ix.accounts.mint.address())?; + drop(distribution_data); if distribution.clawback_ts != 0 { let current_ts = get_current_timestamp()?; @@ -60,9 +61,8 @@ pub fn process_close_direct_distribution( .invoke_signed(signers) })?; - drop(distribution_data); - - close_pda_account(ix.accounts.distribution, ix.accounts.authority)?; + // Flip the distribution PDA to its permanently-closed state and refund the freed rent. + DirectDistribution::close_in_place(ix.accounts.distribution, ix.accounts.authority)?; let event = DistributionClosedEvent::new(*ix.accounts.distribution.address(), remaining_amount); emit_event(&ID, ix.accounts.event_authority, ix.accounts.program, &event.to_bytes())?; diff --git a/program/src/instructions/direct/close_recipient/accounts.rs b/program/src/instructions/direct/close_recipient/accounts.rs index c40bb2e..455691f 100644 --- a/program/src/instructions/direct/close_recipient/accounts.rs +++ b/program/src/instructions/direct/close_recipient/accounts.rs @@ -33,8 +33,8 @@ impl<'a> TryFrom<&'a [AccountView]> for CloseDirectRecipientAccounts<'a> { verify_current_program(program)?; verify_event_authority(event_authority)?; - // 4. Validate accounts owned by current program - verify_current_program_account(distribution)?; + // 4. distribution ownership is validated in processor: + // program-owned => active distribution, anything else => treated as closed. verify_current_program_account(recipient_account)?; Ok(Self { recipient, original_payer, distribution, recipient_account, event_authority, program }) diff --git a/program/src/instructions/direct/close_recipient/processor.rs b/program/src/instructions/direct/close_recipient/processor.rs index e19e80c..09613ab 100644 --- a/program/src/instructions/direct/close_recipient/processor.rs +++ b/program/src/instructions/direct/close_recipient/processor.rs @@ -19,9 +19,15 @@ pub fn process_close_direct_recipient( let ix = CloseDirectRecipient::try_from((instruction_data, accounts))?; ix.data.validate()?; - let distribution_data = ix.accounts.distribution.try_borrow()?; - let _distribution = DirectDistribution::from_account(&distribution_data, ix.accounts.distribution, &ID)?; - drop(distribution_data); + // The distribution PDA always lives on: active as `DirectDistribution`, + // permanently closed as a compact `DirectDistributionClosed` marker. + let distribution_closed = DirectDistribution::is_closed(ix.accounts.distribution, &ID)?; + + if !distribution_closed { + let distribution_data = ix.accounts.distribution.try_borrow()?; + let _distribution = DirectDistribution::from_account(&distribution_data, ix.accounts.distribution, &ID)?; + drop(distribution_data); + } let recipient_data = ix.accounts.recipient_account.try_borrow()?; let recipient = DirectRecipient::from_account(&recipient_data, ix.accounts.recipient_account, &ID)?; @@ -35,7 +41,7 @@ pub fn process_close_direct_recipient( return Err(ProgramError::InvalidAccountData); } - if recipient.claimed_amount < recipient.total_amount { + if !distribution_closed && recipient.claimed_amount < recipient.total_amount { return Err(RewardsProgramError::ClaimNotFullyVested.into()); } diff --git a/program/src/instructions/direct/create_distribution/accounts.rs b/program/src/instructions/direct/create_distribution/accounts.rs index 452d653..1e68fb7 100644 --- a/program/src/instructions/direct/create_distribution/accounts.rs +++ b/program/src/instructions/direct/create_distribution/accounts.rs @@ -54,7 +54,7 @@ impl<'a> TryFrom<&'a [AccountView]> for CreateDirectDistributionAccounts<'a> { verify_current_program(program)?; verify_event_authority(event_authority)?; - // 4. (no accounts owned by current program for this instruction) + // 4. (distribution may be uninitialized or hold a closed marker; processor validates) // 5. Validate token account ownership verify_owned_by(mint, token_program.address())?; diff --git a/program/src/instructions/direct/create_distribution/processor.rs b/program/src/instructions/direct/create_distribution/processor.rs index b68ccb9..acf97ba 100644 --- a/program/src/instructions/direct/create_distribution/processor.rs +++ b/program/src/instructions/direct/create_distribution/processor.rs @@ -2,6 +2,7 @@ use pinocchio::{account::AccountView, error::ProgramError, Address, ProgramResul use pinocchio_associated_token_account::instructions::CreateIdempotent; use crate::{ + errors::RewardsProgramError, events::DistributionCreatedEvent, state::DirectDistribution, traits::{AccountSerialize, AccountSize, EventSerialize, InstructionData, PdaSeeds}, @@ -30,6 +31,12 @@ pub fn process_create_direct_distribution( distribution.validate_pda(ix.accounts.distribution, &ID, ix.data.bump)?; + // If the distribution PDA was previously closed, it still lives as a + // `DirectDistributionClosed` marker. Reject re-creation. + if DirectDistribution::is_closed(ix.accounts.distribution, &ID)? { + return Err(RewardsProgramError::DistributionPermanentlyClosed.into()); + } + let bump_seed = [ix.data.bump]; let distribution_seeds = distribution.seeds_with_bump(&bump_seed); let distribution_seeds_array: [_; 5] = distribution_seeds.try_into().map_err(|_| ProgramError::InvalidArgument)?; diff --git a/program/src/instructions/merkle/revoke_claim/accounts.rs b/program/src/instructions/merkle/revoke_claim/accounts.rs index dc56552..28f49b5 100644 --- a/program/src/instructions/merkle/revoke_claim/accounts.rs +++ b/program/src/instructions/merkle/revoke_claim/accounts.rs @@ -43,13 +43,13 @@ impl<'a> TryFrom<&'a [AccountView]> for RevokeMerkleClaimAccounts<'a> { // 2. Validate writable verify_writable(distribution, true)?; + verify_writable(claim_account, true)?; verify_writable(revocation_marker, true)?; verify_writable(distribution_vault, true)?; verify_writable(claimant_token_account, true)?; verify_writable(authority_token_account, true)?; // 2b. Validate read-only accounts - verify_readonly(claim_account)?; verify_readonly(claimant)?; verify_readonly(mint)?; diff --git a/program/src/instructions/merkle/revoke_claim/processor.rs b/program/src/instructions/merkle/revoke_claim/processor.rs index 99a5e63..eaef668 100644 --- a/program/src/instructions/merkle/revoke_claim/processor.rs +++ b/program/src/instructions/merkle/revoke_claim/processor.rs @@ -6,8 +6,8 @@ use crate::{ events::RecipientRevokedEvent, state::{MerkleClaim, MerkleClaimSeeds, MerkleDistribution, Revocation, RevocationSeeds}, traits::{ - AccountParse, AccountSerialize, AccountSize, Distribution, DistributionSigner, EventSerialize, InstructionData, - PdaSeeds, VestingParams, + AccountParse, AccountSerialize, AccountSize, ClaimTracker, Distribution, DistributionSigner, EventSerialize, + InstructionData, PdaSeeds, VestingParams, }, utils::{ compute_leaf_hash, create_pda_account, emit_event, get_current_timestamp, get_mint_decimals, @@ -61,14 +61,15 @@ pub fn process_revoke_merkle_claim( }; claim_seeds.validate_pda_address(ix.accounts.claim_account, &ID)?; - let claimed_amount = if is_pda_uninitialized(ix.accounts.claim_account) { - 0u64 + let mut claim = if is_pda_uninitialized(ix.accounts.claim_account) { + None } else { let claim_data = ix.accounts.claim_account.try_borrow()?; let claim = MerkleClaim::parse_from_bytes(&claim_data)?; drop(claim_data); - claim.claimed_amount + Some(claim) }; + let claimed_amount = claim.as_ref().map_or(0u64, |existing_claim| existing_claim.claimed_amount); // Calculate vesting let vested_amount = VestingParams::calculate_unlocked(&ix.data, current_ts)?; @@ -77,6 +78,7 @@ pub fn process_revoke_merkle_claim( // Apply revoke mode let decimals = get_mint_decimals(ix.accounts.mint)?; + let mut claim_needs_write = false; let (vested_transferred, total_freed) = match ix.data.revoke_mode { RevokeMode::NonVested => { @@ -96,6 +98,10 @@ pub fn process_revoke_merkle_claim( } Distribution::add_claimed(&mut distribution, vested_unclaimed)?; + if let Some(existing_claim) = claim.as_mut() { + ClaimTracker::add_claimed(existing_claim, vested_unclaimed)?; + claim_needs_write = vested_unclaimed > 0; + } (vested_unclaimed, unvested) } @@ -125,6 +131,13 @@ pub fn process_revoke_merkle_claim( distribution.write_to_slice(&mut distribution_data)?; drop(distribution_data); + if claim_needs_write { + let existing_claim = claim.as_ref().ok_or(RewardsProgramError::InvalidAccountData)?; + let mut claim_data = ix.accounts.claim_account.try_borrow_mut()?; + existing_claim.write_to_slice(&mut claim_data)?; + drop(claim_data); + } + // Create revocation PDA let revocation_bump_seed = [revocation_bump]; let revocation_pda_seeds = revocation_seeds.seeds_with_bump(&revocation_bump_seed); diff --git a/program/src/state/direct_distribution.rs b/program/src/state/direct_distribution.rs index d776ac6..9148ddc 100644 --- a/program/src/state/direct_distribution.rs +++ b/program/src/state/direct_distribution.rs @@ -5,14 +5,16 @@ use pinocchio::{ account::AccountView, cpi::{Seed, Signer}, error::ProgramError, - Address, + Address, ProgramResult, }; use crate::errors::RewardsProgramError; +use crate::state::DirectDistributionClosed; use crate::traits::{ AccountParse, AccountSerialize, AccountSize, AccountValidation, Discriminator, Distribution, DistributionSigner, PdaAccount, PdaSeeds, RewardsAccountDiscriminators, Versioned, }; +use crate::utils::refund_excess_rent; use crate::{assert_no_padding, require_account_len, validate_discriminator, validate_version}; /// DirectDistribution account state @@ -207,6 +209,40 @@ impl DirectDistribution { Ok(state) } + /// Returns `true` if the given account holds a `DirectDistributionClosed` marker + /// — i.e., a distribution that was previously closed at this PDA address. + /// + /// An account owned by a different program (including the system program + /// for a fresh/uninitialized PDA) is never considered closed. + #[inline(always)] + pub fn is_closed(account: &AccountView, program_id: &Address) -> Result { + if !account.owned_by(program_id) { + return Ok(false); + } + let data = account.try_borrow()?; + Ok(!data.is_empty() && data[0] == DirectDistributionClosed::DISCRIMINATOR) + } + + /// Transition an active distribution PDA to its permanently-closed state: + /// overwrite the account header with the `DirectDistributionClosed` + /// discriminator + version (preserving the bump at byte 2), shrink the + /// account to the closed marker's size, and refund the freed rent to + /// `rent_recipient`. + /// + /// The caller is responsible for transferring any remaining token balance + /// and closing the distribution vault before invoking this. + pub fn close_in_place(account: &AccountView, rent_recipient: &AccountView) -> ProgramResult { + { + let mut data = account.try_borrow_mut()?; + data[0] = DirectDistributionClosed::DISCRIMINATOR; + data[1] = DirectDistributionClosed::VERSION; + // data[2] already holds the bump (unchanged from active layout) + } + + account.resize(DirectDistributionClosed::LEN)?; + refund_excess_rent(account, rent_recipient, DirectDistributionClosed::LEN) + } + pub fn remaining_unallocated(&self, vault_balance: u64) -> Result { let outstanding = self.total_allocated.checked_sub(self.total_claimed).ok_or(RewardsProgramError::MathOverflow)?; diff --git a/program/src/state/direct_distribution_closed.rs b/program/src/state/direct_distribution_closed.rs new file mode 100644 index 0000000..eef0ca2 --- /dev/null +++ b/program/src/state/direct_distribution_closed.rs @@ -0,0 +1,89 @@ +use alloc::vec::Vec; +use codama::CodamaAccount; +use pinocchio::error::ProgramError; + +use crate::traits::{ + AccountParse, AccountSerialize, AccountSize, AccountValidation, Discriminator, RewardsAccountDiscriminators, + Versioned, +}; +use crate::{require_account_len, validate_discriminator, validate_version}; + +/// `DirectDistributionClosed` is the permanently-closed state of a direct +/// distribution PDA. The same PDA address is reused — only the discriminator +/// flips on close — so no separate tombstone account is needed. +/// +/// After close, the distribution account is resized down to `LEN` (3 bytes: +/// discriminator + version + bump), and the freed rent is refunded to the +/// authority. On subsequent `create_direct_distribution` calls, the presence +/// of this discriminator at the PDA address triggers `DistributionPermanentlyClosed`. +#[derive(Clone, Debug, PartialEq, CodamaAccount)] +#[repr(C)] +pub struct DirectDistributionClosed { + pub bump: u8, +} + +impl Discriminator for DirectDistributionClosed { + const DISCRIMINATOR: u8 = RewardsAccountDiscriminators::DirectDistributionClosed as u8; +} + +impl Versioned for DirectDistributionClosed { + const VERSION: u8 = 1; +} + +impl AccountSize for DirectDistributionClosed { + const DATA_LEN: usize = 1; // bump +} + +impl AccountParse for DirectDistributionClosed { + fn parse_from_bytes(data: &[u8]) -> Result { + require_account_len!(data, Self::LEN); + validate_discriminator!(data, Self::DISCRIMINATOR); + validate_version!(data, Self::VERSION); + + Ok(Self { bump: data[2] }) + } +} + +impl AccountSerialize for DirectDistributionClosed { + #[inline(always)] + fn to_bytes_inner(&self) -> Vec { + let mut data = Vec::with_capacity(Self::DATA_LEN); + data.push(self.bump); + data + } +} + +impl AccountValidation for DirectDistributionClosed {} + +impl DirectDistributionClosed { + #[inline(always)] + pub fn new(bump: u8) -> Self { + Self { bump } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_direct_distribution_closed_roundtrip() { + let closed = DirectDistributionClosed::new(255); + let bytes = closed.to_bytes(); + assert_eq!(bytes.len(), DirectDistributionClosed::LEN); + assert_eq!(bytes[0], DirectDistributionClosed::DISCRIMINATOR); + assert_eq!(bytes[1], DirectDistributionClosed::VERSION); + assert_eq!(bytes[2], 255); + + let decoded = DirectDistributionClosed::parse_from_bytes(&bytes).unwrap(); + assert_eq!(decoded, closed); + } + + #[test] + fn test_parse_rejects_wrong_discriminator() { + let closed = DirectDistributionClosed::new(100); + let mut bytes = closed.to_bytes(); + bytes[0] = 0xFF; // corrupt discriminator + assert!(DirectDistributionClosed::parse_from_bytes(&bytes).is_err()); + } +} diff --git a/program/src/state/mod.rs b/program/src/state/mod.rs index 0c6248b..7a37b31 100644 --- a/program/src/state/mod.rs +++ b/program/src/state/mod.rs @@ -1,4 +1,5 @@ pub mod direct_distribution; +pub mod direct_distribution_closed; pub mod direct_recipient; pub mod merkle_claim; pub mod merkle_distribution; @@ -8,6 +9,7 @@ pub mod reward_pool; pub mod user_reward_account; pub use direct_distribution::*; +pub use direct_distribution_closed::*; pub use direct_recipient::*; pub use merkle_claim::*; pub use merkle_distribution::*; diff --git a/program/src/traits/account.rs b/program/src/traits/account.rs index bcb6012..4bca1db 100644 --- a/program/src/traits/account.rs +++ b/program/src/traits/account.rs @@ -91,6 +91,7 @@ pub enum RewardsAccountDiscriminators { RewardPool = 5, UserRewardAccount = 6, PointsConfig = 7, + DirectDistributionClosed = 8, } /// Manual account deserialization (non-zero-copy) diff --git a/program/src/utils/pda_utils.rs b/program/src/utils/pda_utils.rs index 5f50617..900c5cb 100644 --- a/program/src/utils/pda_utils.rs +++ b/program/src/utils/pda_utils.rs @@ -47,6 +47,25 @@ pub fn close_pda_account(pda_account: &AccountView, recipient: &AccountView) -> Ok(()) } +/// Refund any lamports above the rent-exempt minimum for `new_size` back to `recipient`. +/// +/// Call after shrinking a program-owned PDA via `account.resize(new_size)` to return +/// the freed rent. The account must be owned by the invoking program so that direct +/// lamport manipulation is valid. +pub fn refund_excess_rent(account: &AccountView, recipient: &AccountView, new_size: usize) -> ProgramResult { + let rent = Rent::get()?; + let required = rent.try_minimum_balance(new_size).map_err(|_| RewardsProgramError::RentCalculationFailed)?; + let current = account.lamports(); + let excess = current.saturating_sub(required); + + if excess > 0 { + account.set_lamports(current.saturating_sub(excess)); + recipient.set_lamports(recipient.lamports().checked_add(excess).ok_or(RewardsProgramError::MathOverflow)?); + } + + Ok(()) +} + /// Create a PDA account for the given seeds. /// /// Supports pre-funded, system-owned PDA addresses with zero data by diff --git a/tests/integration-tests/src/fixtures/close_direct_distribution.rs b/tests/integration-tests/src/fixtures/close_direct_distribution.rs index ee4eba1..9000a4a 100644 --- a/tests/integration-tests/src/fixtures/close_direct_distribution.rs +++ b/tests/integration-tests/src/fixtures/close_direct_distribution.rs @@ -169,10 +169,6 @@ impl InstructionTestFixture for CloseDirectDistributionFixture { &[0, 1, 3, 4] } - fn system_program_index() -> Option { - None - } - fn current_program_index() -> Option { Some(7) } diff --git a/tests/integration-tests/src/fixtures/revoke_merkle_claim.rs b/tests/integration-tests/src/fixtures/revoke_merkle_claim.rs index 17a8c88..0e69e08 100644 --- a/tests/integration-tests/src/fixtures/revoke_merkle_claim.rs +++ b/tests/integration-tests/src/fixtures/revoke_merkle_claim.rs @@ -273,12 +273,13 @@ impl InstructionTestFixture for RevokeMerkleClaimFixture { /// Account indices that must be writable: /// 1: payer /// 2: distribution + /// 3: claim_account /// 4: revocation_marker /// 7: distribution_vault /// 8: claimant_token_account /// 9: authority_token_account fn required_writable() -> &'static [usize] { - &[1, 2, 4, 7, 8, 9] + &[1, 2, 3, 4, 7, 8, 9] } fn system_program_index() -> Option { diff --git a/tests/integration-tests/src/test_close_direct_distribution.rs b/tests/integration-tests/src/test_close_direct_distribution.rs index c43885e..ea8cf81 100644 --- a/tests/integration-tests/src/test_close_direct_distribution.rs +++ b/tests/integration-tests/src/test_close_direct_distribution.rs @@ -5,10 +5,31 @@ use crate::fixtures::{ CreateDirectDistributionSetup, }; use crate::utils::{ - assert_account_closed, assert_rewards_error, test_empty_data, test_missing_signer, test_not_writable, - test_wrong_current_program, RewardsError, TestContext, + assert_rewards_error, test_empty_data, test_missing_signer, test_not_writable, test_wrong_current_program, + RewardsError, TestContext, PROGRAM_ID, }; +/// `DirectDistributionClosed::DISCRIMINATOR` (matches `RewardsAccountDiscriminators::DirectDistributionClosed`). +const DIRECT_DISTRIBUTION_CLOSED_DISCRIMINATOR: u8 = 8; +/// `DirectDistributionClosed::LEN` = 1 (disc) + 1 (version) + 1 (bump). +const DIRECT_DISTRIBUTION_CLOSED_LEN: usize = 3; + +/// After close, the distribution PDA still lives on but has been flipped to a +/// `DirectDistributionClosed` marker: discriminator byte = 8, total data length = 3. +fn assert_distribution_closed_marker(ctx: &TestContext, distribution_pda: &solana_sdk::pubkey::Pubkey) { + let account = ctx.get_account(distribution_pda).expect("Distribution PDA should still exist after close"); + assert_eq!(account.owner, PROGRAM_ID, "Distribution should still be owned by program"); + assert_eq!( + account.data.len(), + DIRECT_DISTRIBUTION_CLOSED_LEN, + "Closed distribution should be resized to DirectDistributionClosed::LEN" + ); + assert_eq!( + account.data[0], DIRECT_DISTRIBUTION_CLOSED_DISCRIMINATOR, + "First byte should be DirectDistributionClosed discriminator" + ); +} + #[test] fn test_close_direct_distribution_missing_authority_signer() { let mut ctx = TestContext::new(); @@ -53,7 +74,7 @@ fn test_close_direct_distribution_success() { let test_ix = setup.build_instruction(&ctx); test_ix.send_expect_success(&mut ctx); - assert_account_closed(&ctx, &setup.distribution_pda); + assert_distribution_closed_marker(&ctx, &setup.distribution_pda); } #[test] @@ -64,7 +85,7 @@ fn test_close_direct_distribution_success_token_2022() { let test_ix = setup.build_instruction(&ctx); test_ix.send_expect_success(&mut ctx); - assert_account_closed(&ctx, &setup.distribution_pda); + assert_distribution_closed_marker(&ctx, &setup.distribution_pda); } #[test] @@ -123,7 +144,7 @@ fn test_close_direct_distribution_clawback_ts_zero_succeeds() { let setup = CloseDirectDistributionSetup::new(&mut ctx); let test_ix = setup.build_instruction(&ctx); test_ix.send_expect_success(&mut ctx); - assert_account_closed(&ctx, &setup.distribution_pda); + assert_distribution_closed_marker(&ctx, &setup.distribution_pda); } #[test] @@ -180,5 +201,5 @@ fn test_close_direct_distribution_clawback_ts_after_timestamp_succeeds() { let test_ix = close_setup.build_instruction(&ctx); test_ix.send_expect_success(&mut ctx); - assert_account_closed(&ctx, &close_setup.distribution_pda); + assert_distribution_closed_marker(&ctx, &close_setup.distribution_pda); } diff --git a/tests/integration-tests/src/test_close_direct_recipient.rs b/tests/integration-tests/src/test_close_direct_recipient.rs index d59d7ec..b65eee1 100644 --- a/tests/integration-tests/src/test_close_direct_recipient.rs +++ b/tests/integration-tests/src/test_close_direct_recipient.rs @@ -3,7 +3,8 @@ use rewards_program_client::types::VestingSchedule; use solana_sdk::{instruction::InstructionError, signature::Signer}; use crate::fixtures::{ - ClaimDirectSetup, CloseDirectRecipientFixture, CloseDirectRecipientSetup, CreateDirectDistributionSetup, + AddDirectRecipientSetup, ClaimDirectSetup, CloseDirectDistributionSetup, CloseDirectRecipientFixture, + CloseDirectRecipientSetup, CreateDirectDistributionSetup, }; use crate::utils::{ assert_account_closed, assert_instruction_error, assert_rewards_error, find_direct_recipient_pda, @@ -70,6 +71,45 @@ fn test_close_direct_recipient_success_token_2022() { assert_account_closed(&ctx, &setup.recipient_pda); } +#[test] +fn test_close_direct_recipient_after_distribution_closed() { + let mut ctx = TestContext::new(); + + let recipient_setup = AddDirectRecipientSetup::new(&mut ctx); + let add_ix = recipient_setup.build_instruction(&ctx); + add_ix.send_expect_success(&mut ctx); + + let close_distribution_setup = CloseDirectDistributionSetup { + authority: recipient_setup.authority.insecure_clone(), + distribution_pda: recipient_setup.distribution_pda, + mint: recipient_setup.mint, + distribution_vault: recipient_setup.distribution_vault, + authority_token_account: recipient_setup.authority_token_account, + token_program: recipient_setup.token_program, + }; + + let close_distribution_ix = close_distribution_setup.build_instruction(&ctx); + close_distribution_ix.send_expect_success(&mut ctx); + // After close, the distribution PDA still lives on as a DirectDistributionClosed marker (3 bytes, disc=8). + let closed_account = + ctx.get_account(&recipient_setup.distribution_pda).expect("Distribution PDA should still exist after close"); + assert_eq!(closed_account.data.len(), 3, "Closed distribution should be 3 bytes"); + assert_eq!(closed_account.data[0], 8, "First byte should be DirectDistributionClosed discriminator"); + + let close_recipient_setup = CloseDirectRecipientSetup { + recipient: recipient_setup.recipient.insecure_clone(), + original_payer: ctx.payer.insecure_clone(), + distribution_pda: recipient_setup.distribution_pda, + recipient_pda: recipient_setup.recipient_pda, + token_program: recipient_setup.token_program, + }; + + let close_recipient_ix = close_recipient_setup.build_instruction(&ctx); + close_recipient_ix.send_expect_success(&mut ctx); + + assert_account_closed(&ctx, &recipient_setup.recipient_pda); +} + #[test] fn test_close_direct_recipient_claim_not_fully_vested() { let mut ctx = TestContext::new(); diff --git a/tests/integration-tests/src/test_create_direct_distribution.rs b/tests/integration-tests/src/test_create_direct_distribution.rs index 540dde7..2e6a487 100644 --- a/tests/integration-tests/src/test_create_direct_distribution.rs +++ b/tests/integration-tests/src/test_create_direct_distribution.rs @@ -1,10 +1,10 @@ use solana_sdk::{account::Account, signer::Signer}; use solana_system_interface::program::ID as SYSTEM_PROGRAM_ID; -use crate::fixtures::{CreateDirectDistributionFixture, CreateDirectDistributionSetup}; +use crate::fixtures::{CloseDirectDistributionSetup, CreateDirectDistributionFixture, CreateDirectDistributionSetup}; use crate::utils::{ - assert_direct_distribution, test_empty_data, test_missing_signer, test_not_writable, test_truncated_data, - test_wrong_current_program, test_wrong_system_program, TestContext, + assert_direct_distribution, assert_rewards_error, test_empty_data, test_missing_signer, test_not_writable, + test_truncated_data, test_wrong_current_program, test_wrong_system_program, RewardsError, TestContext, }; #[test] @@ -112,3 +112,28 @@ fn test_create_direct_distribution_prefunded_distribution_pda() { setup.bump, ); } + +#[test] +fn test_create_direct_distribution_fails_if_closed_before() { + let mut ctx = TestContext::new(); + let setup = CreateDirectDistributionSetup::new(&mut ctx); + let create_ix = setup.build_instruction(&ctx); + create_ix.send_expect_success(&mut ctx); + + let authority_token_account = + ctx.create_ata_for_program(&setup.authority.pubkey(), &setup.mint.pubkey(), &setup.token_program); + let close_setup = CloseDirectDistributionSetup { + authority: setup.authority.insecure_clone(), + distribution_pda: setup.distribution_pda, + mint: setup.mint.pubkey(), + distribution_vault: setup.distribution_vault, + authority_token_account, + token_program: setup.token_program, + }; + close_setup.build_instruction(&ctx).send_expect_success(&mut ctx); + + ctx.advance_slot(); + let recreate_ix = setup.build_instruction(&ctx); + let error = recreate_ix.send_expect_error(&mut ctx); + assert_rewards_error(error, RewardsError::DistributionPermanentlyClosed); +} diff --git a/tests/integration-tests/src/test_revoke_merkle_claim.rs b/tests/integration-tests/src/test_revoke_merkle_claim.rs index d6fc6e4..bc50b4a 100644 --- a/tests/integration-tests/src/test_revoke_merkle_claim.rs +++ b/tests/integration-tests/src/test_revoke_merkle_claim.rs @@ -4,8 +4,8 @@ use solana_sdk::instruction::InstructionError; use crate::fixtures::{RevokeMerkleClaimFixture, RevokeMerkleClaimSetup}; use crate::utils::{ - assert_rewards_error, expected_linear_unlock, test_empty_data, test_missing_signer, test_not_writable, - test_wrong_account, test_wrong_current_program, RewardsError, TestContext, PROGRAM_ID, + assert_merkle_claim, assert_rewards_error, expected_linear_unlock, test_empty_data, test_missing_signer, + test_not_writable, test_wrong_account, test_wrong_current_program, RewardsError, TestContext, PROGRAM_ID, }; // ── Generic fixture tests ────────────────────────────────────────── @@ -28,6 +28,12 @@ fn test_revoke_merkle_distribution_not_writable() { test_not_writable::(&mut ctx, 2); } +#[test] +fn test_revoke_merkle_claim_account_not_writable() { + let mut ctx = TestContext::new(); + test_not_writable::(&mut ctx, 3); +} + #[test] fn test_revoke_merkle_revocation_marker_not_writable() { let mut ctx = TestContext::new(); @@ -306,6 +312,8 @@ fn test_revoke_merkle_non_vested_after_partial_claim() { claimed_at_quarter + vested_unclaimed, "total_claimed should reflect both claim and revoke transfer" ); + + assert_merkle_claim(&ctx, &setup.claim_pda, vested_at_midpoint, setup.claim_bump); } #[test] @@ -336,6 +344,8 @@ fn test_revoke_merkle_full_after_partial_claim() { let dist_account = ctx.get_account(&setup.distribution_pda).expect("Distribution should exist"); let dist = MerkleDistribution::from_bytes(&dist_account.data).expect("Should deserialize"); assert_eq!(dist.total_claimed, claimed_at_quarter, "total_claimed should only reflect the original claim"); + + assert_merkle_claim(&ctx, &setup.claim_pda, claimed_at_quarter, setup.claim_bump); } // ── Post-revocation behavior ──────────────────────────────────────