fix(rewards): harden direct close lifecycle and revoke accounting (#32)

* fix(rewards): harden direct close lifecycle and revoke accounting

---------

Co-authored-by: Jo D <dev-jodee@users.noreply.github.com>

Author: dev-jodee

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"
           },

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<RewardsProgramError> for ProgramError {

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",

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())?;

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 })

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());
     }
 

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())?;

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)?;

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)?;
 

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);