Skip to content

Commit aa1cfd9

Browse files
authored
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>
1 parent dc24958 commit aa1cfd9

21 files changed

Lines changed: 341 additions & 42 deletions

idl/rewards_program.json

Lines changed: 28 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,24 @@
9494
"kind": "accountNode",
9595
"name": "directDistribution"
9696
},
97+
{
98+
"data": {
99+
"fields": [
100+
{
101+
"kind": "structFieldTypeNode",
102+
"name": "bump",
103+
"type": {
104+
"endian": "le",
105+
"format": "u8",
106+
"kind": "numberTypeNode"
107+
}
108+
}
109+
],
110+
"kind": "structTypeNode"
111+
},
112+
"kind": "accountNode",
113+
"name": "directDistributionClosed"
114+
},
97115
{
98116
"data": {
99117
"fields": [
@@ -1942,6 +1960,12 @@
19421960
"kind": "errorNode",
19431961
"message": "Nothing to revoke — user has zero balance",
19441962
"name": "pointsNothingToRevoke"
1963+
},
1964+
{
1965+
"code": 42,
1966+
"kind": "errorNode",
1967+
"message": "Distribution has been permanently closed",
1968+
"name": "distributionPermanentlyClosed"
19451969
}
19461970
],
19471971
"instructions": [
@@ -1976,7 +2000,7 @@
19762000
},
19772001
{
19782002
"docs": [
1979-
"PDA: [b\"direct_distribution\", mint, authority, seeds] (created)"
2003+
"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."
19802004
],
19812005
"isSigner": false,
19822006
"isWritable": true,
@@ -2395,7 +2419,7 @@
23952419
},
23962420
{
23972421
"docs": [
2398-
"PDA: DirectDistribution account (closed)"
2422+
"PDA: DirectDistribution account. Flipped to a compact DirectDistributionClosed marker in place; freed rent is refunded to authority."
23992423
],
24002424
"isSigner": false,
24012425
"isWritable": true,
@@ -3314,10 +3338,10 @@
33143338
},
33153339
{
33163340
"docs": [
3317-
"PDA: [b\"merkle_claim\", distribution, claimant] (read-only, may not exist)"
3341+
"PDA: [b\"merkle_claim\", distribution, claimant] (writable, may not exist)"
33183342
],
33193343
"isSigner": false,
3320-
"isWritable": false,
3344+
"isWritable": true,
33213345
"kind": "instructionAccountNode",
33223346
"name": "claimAccount"
33233347
},

program/src/errors.rs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -172,6 +172,10 @@ pub enum RewardsProgramError {
172172
/// (41) Nothing to revoke — user has zero balance
173173
#[error("Nothing to revoke — user has zero balance")]
174174
PointsNothingToRevoke,
175+
176+
/// (42) Distribution has been permanently closed
177+
#[error("Distribution has been permanently closed")]
178+
DistributionPermanentlyClosed,
175179
}
176180

177181
impl From<RewardsProgramError> for ProgramError {

program/src/instructions/definition.rs

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ pub enum RewardsProgramInstruction {
1414
#[codama(account(
1515
name = "distribution",
1616
writable,
17-
docs = "PDA: [b\"direct_distribution\", mint, authority, seeds] (created)"
17+
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."
1818
))]
1919
#[codama(account(name = "mint", docs = "SPL token mint"))]
2020
#[codama(account(
@@ -110,7 +110,11 @@ pub enum RewardsProgramInstruction {
110110
writable,
111111
docs = "Distribution authority; receives rent + remaining distribution vault tokens"
112112
))]
113-
#[codama(account(name = "distribution", writable, docs = "PDA: DirectDistribution account (closed)"))]
113+
#[codama(account(
114+
name = "distribution",
115+
writable,
116+
docs = "PDA: DirectDistribution account. Flipped to a compact DirectDistributionClosed marker in place; freed rent is refunded to authority."
117+
))]
114118
#[codama(account(name = "mint", docs = "SPL token mint"))]
115119
#[codama(account(
116120
name = "distribution_vault",
@@ -324,7 +328,8 @@ pub enum RewardsProgramInstruction {
324328
#[codama(account(name = "distribution", writable, docs = "PDA: MerkleDistribution account"))]
325329
#[codama(account(
326330
name = "claim_account",
327-
docs = "PDA: [b\"merkle_claim\", distribution, claimant] (read-only, may not exist)"
331+
writable,
332+
docs = "PDA: [b\"merkle_claim\", distribution, claimant] (writable, may not exist)"
328333
))]
329334
#[codama(account(
330335
name = "revocation_marker",

program/src/instructions/direct/close_distribution/processor.rs

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ use crate::{
66
events::DistributionClosedEvent,
77
state::DirectDistribution,
88
traits::{Distribution, DistributionSigner, EventSerialize, InstructionData},
9-
utils::{close_pda_account, emit_event, get_current_timestamp, get_mint_decimals, get_token_account_balance},
9+
utils::{emit_event, get_current_timestamp, get_mint_decimals, get_token_account_balance},
1010
ID,
1111
};
1212

@@ -24,6 +24,7 @@ pub fn process_close_direct_distribution(
2424
let distribution = DirectDistribution::from_account(&distribution_data, ix.accounts.distribution, &ID)?;
2525
distribution.validate_authority(ix.accounts.authority.address())?;
2626
distribution.validate_mint(ix.accounts.mint.address())?;
27+
drop(distribution_data);
2728

2829
if distribution.clawback_ts != 0 {
2930
let current_ts = get_current_timestamp()?;
@@ -60,9 +61,8 @@ pub fn process_close_direct_distribution(
6061
.invoke_signed(signers)
6162
})?;
6263

63-
drop(distribution_data);
64-
65-
close_pda_account(ix.accounts.distribution, ix.accounts.authority)?;
64+
// Flip the distribution PDA to its permanently-closed state and refund the freed rent.
65+
DirectDistribution::close_in_place(ix.accounts.distribution, ix.accounts.authority)?;
6666

6767
let event = DistributionClosedEvent::new(*ix.accounts.distribution.address(), remaining_amount);
6868
emit_event(&ID, ix.accounts.event_authority, ix.accounts.program, &event.to_bytes())?;

program/src/instructions/direct/close_recipient/accounts.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,8 +33,8 @@ impl<'a> TryFrom<&'a [AccountView]> for CloseDirectRecipientAccounts<'a> {
3333
verify_current_program(program)?;
3434
verify_event_authority(event_authority)?;
3535

36-
// 4. Validate accounts owned by current program
37-
verify_current_program_account(distribution)?;
36+
// 4. distribution ownership is validated in processor:
37+
// program-owned => active distribution, anything else => treated as closed.
3838
verify_current_program_account(recipient_account)?;
3939

4040
Ok(Self { recipient, original_payer, distribution, recipient_account, event_authority, program })

program/src/instructions/direct/close_recipient/processor.rs

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,9 +19,15 @@ pub fn process_close_direct_recipient(
1919
let ix = CloseDirectRecipient::try_from((instruction_data, accounts))?;
2020
ix.data.validate()?;
2121

22-
let distribution_data = ix.accounts.distribution.try_borrow()?;
23-
let _distribution = DirectDistribution::from_account(&distribution_data, ix.accounts.distribution, &ID)?;
24-
drop(distribution_data);
22+
// The distribution PDA always lives on: active as `DirectDistribution`,
23+
// permanently closed as a compact `DirectDistributionClosed` marker.
24+
let distribution_closed = DirectDistribution::is_closed(ix.accounts.distribution, &ID)?;
25+
26+
if !distribution_closed {
27+
let distribution_data = ix.accounts.distribution.try_borrow()?;
28+
let _distribution = DirectDistribution::from_account(&distribution_data, ix.accounts.distribution, &ID)?;
29+
drop(distribution_data);
30+
}
2531

2632
let recipient_data = ix.accounts.recipient_account.try_borrow()?;
2733
let recipient = DirectRecipient::from_account(&recipient_data, ix.accounts.recipient_account, &ID)?;
@@ -35,7 +41,7 @@ pub fn process_close_direct_recipient(
3541
return Err(ProgramError::InvalidAccountData);
3642
}
3743

38-
if recipient.claimed_amount < recipient.total_amount {
44+
if !distribution_closed && recipient.claimed_amount < recipient.total_amount {
3945
return Err(RewardsProgramError::ClaimNotFullyVested.into());
4046
}
4147

program/src/instructions/direct/create_distribution/accounts.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ impl<'a> TryFrom<&'a [AccountView]> for CreateDirectDistributionAccounts<'a> {
5454
verify_current_program(program)?;
5555
verify_event_authority(event_authority)?;
5656

57-
// 4. (no accounts owned by current program for this instruction)
57+
// 4. (distribution may be uninitialized or hold a closed marker; processor validates)
5858

5959
// 5. Validate token account ownership
6060
verify_owned_by(mint, token_program.address())?;

program/src/instructions/direct/create_distribution/processor.rs

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ use pinocchio::{account::AccountView, error::ProgramError, Address, ProgramResul
22
use pinocchio_associated_token_account::instructions::CreateIdempotent;
33

44
use crate::{
5+
errors::RewardsProgramError,
56
events::DistributionCreatedEvent,
67
state::DirectDistribution,
78
traits::{AccountSerialize, AccountSize, EventSerialize, InstructionData, PdaSeeds},
@@ -30,6 +31,12 @@ pub fn process_create_direct_distribution(
3031

3132
distribution.validate_pda(ix.accounts.distribution, &ID, ix.data.bump)?;
3233

34+
// If the distribution PDA was previously closed, it still lives as a
35+
// `DirectDistributionClosed` marker. Reject re-creation.
36+
if DirectDistribution::is_closed(ix.accounts.distribution, &ID)? {
37+
return Err(RewardsProgramError::DistributionPermanentlyClosed.into());
38+
}
39+
3340
let bump_seed = [ix.data.bump];
3441
let distribution_seeds = distribution.seeds_with_bump(&bump_seed);
3542
let distribution_seeds_array: [_; 5] = distribution_seeds.try_into().map_err(|_| ProgramError::InvalidArgument)?;

program/src/instructions/merkle/revoke_claim/accounts.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,13 +43,13 @@ impl<'a> TryFrom<&'a [AccountView]> for RevokeMerkleClaimAccounts<'a> {
4343

4444
// 2. Validate writable
4545
verify_writable(distribution, true)?;
46+
verify_writable(claim_account, true)?;
4647
verify_writable(revocation_marker, true)?;
4748
verify_writable(distribution_vault, true)?;
4849
verify_writable(claimant_token_account, true)?;
4950
verify_writable(authority_token_account, true)?;
5051

5152
// 2b. Validate read-only accounts
52-
verify_readonly(claim_account)?;
5353
verify_readonly(claimant)?;
5454
verify_readonly(mint)?;
5555

program/src/instructions/merkle/revoke_claim/processor.rs

Lines changed: 18 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,8 @@ use crate::{
66
events::RecipientRevokedEvent,
77
state::{MerkleClaim, MerkleClaimSeeds, MerkleDistribution, Revocation, RevocationSeeds},
88
traits::{
9-
AccountParse, AccountSerialize, AccountSize, Distribution, DistributionSigner, EventSerialize, InstructionData,
10-
PdaSeeds, VestingParams,
9+
AccountParse, AccountSerialize, AccountSize, ClaimTracker, Distribution, DistributionSigner, EventSerialize,
10+
InstructionData, PdaSeeds, VestingParams,
1111
},
1212
utils::{
1313
compute_leaf_hash, create_pda_account, emit_event, get_current_timestamp, get_mint_decimals,
@@ -61,14 +61,15 @@ pub fn process_revoke_merkle_claim(
6161
};
6262
claim_seeds.validate_pda_address(ix.accounts.claim_account, &ID)?;
6363

64-
let claimed_amount = if is_pda_uninitialized(ix.accounts.claim_account) {
65-
0u64
64+
let mut claim = if is_pda_uninitialized(ix.accounts.claim_account) {
65+
None
6666
} else {
6767
let claim_data = ix.accounts.claim_account.try_borrow()?;
6868
let claim = MerkleClaim::parse_from_bytes(&claim_data)?;
6969
drop(claim_data);
70-
claim.claimed_amount
70+
Some(claim)
7171
};
72+
let claimed_amount = claim.as_ref().map_or(0u64, |existing_claim| existing_claim.claimed_amount);
7273

7374
// Calculate vesting
7475
let vested_amount = VestingParams::calculate_unlocked(&ix.data, current_ts)?;
@@ -77,6 +78,7 @@ pub fn process_revoke_merkle_claim(
7778

7879
// Apply revoke mode
7980
let decimals = get_mint_decimals(ix.accounts.mint)?;
81+
let mut claim_needs_write = false;
8082

8183
let (vested_transferred, total_freed) = match ix.data.revoke_mode {
8284
RevokeMode::NonVested => {
@@ -96,6 +98,10 @@ pub fn process_revoke_merkle_claim(
9698
}
9799

98100
Distribution::add_claimed(&mut distribution, vested_unclaimed)?;
101+
if let Some(existing_claim) = claim.as_mut() {
102+
ClaimTracker::add_claimed(existing_claim, vested_unclaimed)?;
103+
claim_needs_write = vested_unclaimed > 0;
104+
}
99105

100106
(vested_unclaimed, unvested)
101107
}
@@ -125,6 +131,13 @@ pub fn process_revoke_merkle_claim(
125131
distribution.write_to_slice(&mut distribution_data)?;
126132
drop(distribution_data);
127133

134+
if claim_needs_write {
135+
let existing_claim = claim.as_ref().ok_or(RewardsProgramError::InvalidAccountData)?;
136+
let mut claim_data = ix.accounts.claim_account.try_borrow_mut()?;
137+
existing_claim.write_to_slice(&mut claim_data)?;
138+
drop(claim_data);
139+
}
140+
128141
// Create revocation PDA
129142
let revocation_bump_seed = [revocation_bump];
130143
let revocation_pda_seeds = revocation_seeds.seeds_with_bump(&revocation_bump_seed);

0 commit comments

Comments
 (0)