Skip to content

Commit

Permalink
Force users to pass in ATA in remaining_accounts in a strict order to…
Browse files Browse the repository at this point in the history
… reduce computations complexity
  • Loading branch information
kstepanovdev committed Oct 17, 2024
1 parent da0769f commit 55a22cf
Showing 3 changed files with 167 additions and 11 deletions.
2 changes: 1 addition & 1 deletion program-states/src/error.rs
Original file line number Diff line number Diff line change
@@ -129,6 +129,6 @@ pub enum MplStakingError {
#[msg("Cannot deserialize an account")]
DeserializationError,
// 6042 / 0x179a
#[msg("Invalid associated token accounts, they should match the number mint configs of a registrar")]
#[msg("Invalid associated token accounts, they should match the number mint configs of a registrar and must be passed in the strict order")]
InvalidAssoctiatedTokenAccounts,
}
32 changes: 22 additions & 10 deletions programs/voter-stake-registry/src/instructions/close_voter.rs
Original file line number Diff line number Diff line change
@@ -10,9 +10,15 @@ use mplx_staking_states::{
use spl_associated_token_account::get_associated_token_address;
use std::ops::DerefMut;

// Remaining accounts must be all the token token accounts owned by voter, he wants to close,
// they should be writable so that they can be closed and sol required for rent
// can then be sent back to the sol_destination
/// Remaining accounts must be all the token token accounts owned by voter,
/// they should be writable so that they can be closed and sol required for rent
/// can then be sent back to the sol_destination
///
/// Remaining account must be passed in the order of the mint configs in the registrar
/// that aren't default Pubkey addresses. E.g.
/// Registrar { voting_mint: [Pubkey::default, mint2, mint3] }
/// then remaining accounts must be:
/// [mint2 ATA, mint3 ATA]
#[derive(Accounts)]
pub struct CloseVoter<'info> {
pub registrar: AccountLoader<'info, Registrar>,
@@ -109,12 +115,18 @@ pub fn close_voter<'info>(ctx: Context<'_, '_, '_, 'info, CloseVoter<'info>>) ->
get_associated_token_address(&ctx.accounts.voter.key(), &voting_mint_config.mint)
});

for calculated_ata_to_close in calculated_atas_to_close {
let ata_info_to_close = ctx
.remaining_accounts
.iter()
.find(|ata| *ata.key == calculated_ata_to_close)
.ok_or(MplStakingError::InvalidAssoctiatedTokenAccounts)?;
for (index, calculated_ata_to_close) in calculated_atas_to_close.enumerate() {
require!(
ctx.remaining_accounts.len() > index,
MplStakingError::InvalidAssoctiatedTokenAccounts
);

let ata_info_to_close = ctx.remaining_accounts[index].to_account_info();
require_keys_eq!(
*ata_info_to_close.key,
calculated_ata_to_close,
MplStakingError::InvalidAssoctiatedTokenAccounts
);

if ata_info_to_close.data_is_empty()
&& ata_info_to_close.owner == &system_program::ID
@@ -123,7 +135,7 @@ pub fn close_voter<'info>(ctx: Context<'_, '_, '_, 'info, CloseVoter<'info>>) ->
continue;
}

let ata = Account::<TokenAccount>::try_from(ata_info_to_close)
let ata = Account::<TokenAccount>::try_from(&ata_info_to_close)
.map_err(|_| MplStakingError::DeserializationError)?;
require_keys_eq!(
ata.owner,
144 changes: 144 additions & 0 deletions programs/voter-stake-registry/tests/close_voter.rs
Original file line number Diff line number Diff line change
@@ -138,6 +138,7 @@ async fn two_the_same_voting_mints_fail() -> Result<(), TransportError> {
.await?;

let mints = &[context.mints[0].clone(), context.mints[0].clone()];

context
.addin
.close_voter(
@@ -596,3 +597,146 @@ async fn success() -> Result<(), TransportError> {

Ok(())
}

#[tokio::test]
async fn wrong_order_of_passed_in_ata() -> Result<(), TransportError> {
let context = TestContext::new().await;

let payer = &context.users[0].key;
let realm_authority = Keypair::new();
let realm = context
.governance
.create_realm(
"testrealm",
realm_authority.pubkey(),
&context.mints[0],
payer,
&context.addin.program_id,
)
.await;

let deposit_authority = &context.users[1].key;
let token_owner_record = realm
.create_token_owner_record(deposit_authority.pubkey(), payer)
.await;

let fill_authority = Keypair::from_bytes(&context.users[3].key.to_bytes()).unwrap();
let distribution_authority = Keypair::new();
let (registrar, rewards_pool) = context
.addin
.create_registrar(
&realm,
&realm_authority,
payer,
&fill_authority.pubkey(),
&distribution_authority.pubkey(),
&context.rewards.program_id,
)
.await;

let mngo_voting_mint = context
.addin
.configure_voting_mint(
&registrar,
&realm_authority,
payer,
0,
&context.mints[0],
None,
Some(&[context.mints[1].pubkey.unwrap()]),
)
.await;

context
.addin
.configure_voting_mint(
&registrar,
&realm_authority,
payer,
1,
&context.mints[1],
None,
Some(&[context.mints[0].pubkey.unwrap()]),
)
.await;

// TODO: ??? voter_authority == deposit_authority ???
let voter_authority = deposit_authority;
let (deposit_mining, _) = find_deposit_mining_addr(
&context.rewards.program_id,
&voter_authority.pubkey(),
&rewards_pool,
);

let voter = context
.addin
.create_voter(
&registrar,
&token_owner_record,
voter_authority,
payer,
&rewards_pool,
&deposit_mining,
&context.rewards.program_id,
)
.await;

let balance_initial = voter.deposit_amount(&context.solana, 0).await;
assert_eq!(balance_initial, 0);

context
.addin
.create_deposit_entry(
&registrar,
&voter,
&voter,
&mngo_voting_mint,
0,
LockupKind::None,
LockupPeriod::None,
)
.await?;

// test deposit and withdraw
let reference_account = context.users[1].token_accounts[0];
context
.addin
.deposit(
&registrar,
&voter,
&mngo_voting_mint,
deposit_authority,
reference_account,
0,
10000,
)
.await?;

context
.addin
.withdraw(
&registrar,
&voter,
&mngo_voting_mint,
deposit_authority,
reference_account,
0,
10000,
)
.await?;

let mints = &[context.mints[1].clone(), context.mints[0].clone()];
context
.addin
.close_voter(
&registrar,
&voter,
&mints[..],
voter_authority,
&context.rewards.program_id,
)
.await
.assert_on_chain_err(MplStakingError::InvalidAssoctiatedTokenAccounts);

Ok(())
}

0 comments on commit 55a22cf

Please sign in to comment.