[mtg-578] Add both restrict and allow claim instructions with tests

programs/mpl-staking/src/cpi_instructions.rs

@@ -163,6 +163,124 @@ pub enum RewardsInstruction {
         staked_amount: u64,
         new_delegate: Pubkey,
     },
+
+    /// Restricts claiming rewards from the specified mining account
+    ///
+    /// Accounts:
+    /// [RS] Deposit authority
+    /// [S] Reward pool account
+    /// [W] Mining
+    RestrictTokenFlow { mining_owner: Pubkey },
+
+    /// Allows claiming rewards from the specified mining account
+    ///
+    /// Accounts:
+    /// [RS] Deposit authority
+    /// [S] Reward pool account
+    /// [W] Mining
+    AllowTokenFlow { mining_owner: Pubkey },
+
+    /// Restricts batch minting until the specified time
+    ///
+    /// Accounts:
+    /// [RS] Deposit authority
+    /// [S] Reward pool account
+    /// [W] Mining
+    RestrictBatchMinting {
+        /// Time until batch minting is restricted
+        until_ts: u64,
+        /// Owner of the mining account
+        mining_owner: Pubkey,
+    },
+}
+
+pub fn restrict_batch_minting<'a>(
+    program_id: AccountInfo<'a>,
+    deposit_authority: AccountInfo<'a>,
+    reward_pool: AccountInfo<'a>,
+    mining: AccountInfo<'a>,
+    mining_owner: &Pubkey,
+    until_ts: u64,
+    signers_seeds: &[&[u8]],
+) -> ProgramResult {
+    let accounts = vec![
+        AccountMeta::new_readonly(deposit_authority.key(), true),
+        AccountMeta::new_readonly(reward_pool.key(), false),
+        AccountMeta::new(mining.key(), false),
+    ];
+
+    let ix = Instruction::new_with_borsh(
+        program_id.key(),
+        &RewardsInstruction::RestrictBatchMinting {
+            until_ts,
+            mining_owner: *mining_owner,
+        },
+        accounts,
+    );
+
+    invoke_signed(
+        &ix,
+        &[deposit_authority, reward_pool, mining, program_id],
+        &[signers_seeds],
+    )
+}
+
+pub fn restrict_tokenflow<'a>(
+    program_id: AccountInfo<'a>,
+    deposit_authority: AccountInfo<'a>,
+    reward_pool: AccountInfo<'a>,
+    mining: AccountInfo<'a>,
+    mining_owner: &Pubkey,
+    signers_seeds: &[&[u8]],
+) -> ProgramResult {
+    let accounts = vec![
+        AccountMeta::new_readonly(deposit_authority.key(), true),
+        AccountMeta::new_readonly(reward_pool.key(), false),
+        AccountMeta::new(mining.key(), false),
+    ];
+
+    let ix = Instruction::new_with_borsh(
+        program_id.key(),
+        &RewardsInstruction::RestrictTokenFlow {
+            mining_owner: *mining_owner,
+        },
+        accounts,
+    );
+
+    invoke_signed(
+        &ix,
+        &[deposit_authority, reward_pool, mining, program_id],
+        &[signers_seeds],
+    )
+}
+
+pub fn allow_tokenflow<'a>(
+    program_id: AccountInfo<'a>,
+    deposit_authority: AccountInfo<'a>,
+    reward_pool: AccountInfo<'a>,
+    mining: AccountInfo<'a>,
+    mining_owner: &Pubkey,
+    signers_seeds: &[&[u8]],
+) -> ProgramResult {
+    let accounts = vec![
+        AccountMeta::new_readonly(deposit_authority.key(), true),
+        AccountMeta::new_readonly(reward_pool.key(), false),
+        AccountMeta::new(mining.key(), false),
+    ];
+
+    let ix = Instruction::new_with_borsh(
+        program_id.key(),
+        &RewardsInstruction::AllowTokenFlow {
+            mining_owner: *mining_owner,
+        },
+        accounts,
+    );
+
+    invoke_signed(
+        &ix,
+        &[deposit_authority, reward_pool, mining, program_id],
+        &[signers_seeds],
+    )
 }
 
 /// This function initializes pool. Some sort of a "root"

programs/mpl-staking/src/instructions/mod.rs

@@ -1,3 +1,4 @@
+use anchor_lang::prelude::SolanaSysvar;
 pub use change_delegate::*;
 pub use claim::*;
 pub use close_deposit_entry::*;
@@ -9,7 +10,8 @@ pub use create_voter::*;
 pub use deposit::*;
 pub use extend_stake::*;
 pub use log_voter_info::*;
-use solana_program::{clock::Clock, sysvar::Sysvar};
+pub use penalties::*;
+pub use solana_program::sysvar::clock::Clock;
 pub use stake::*;
 pub use unlock_tokens::*;
 pub use update_voter_weight_record::*;
@@ -26,6 +28,7 @@ mod create_voter;
 mod deposit;
 mod extend_stake;
 mod log_voter_info;
+mod penalties;
 mod stake;
 mod unlock_tokens;
 mod update_voter_weight_record;

programs/mpl-staking/src/instructions/penalties/allow_tokenflow.rs

@@ -0,0 +1,33 @@
+use super::Penalty;
+use crate::cpi_instructions;
+use anchor_lang::prelude::*;
+use mplx_staking_states::error::MplStakingError;
+
+/// Restricts claiming rewards from the specified mining account.
+pub fn allow_tokenflow(ctx: Context<Penalty>, mining_owner: Pubkey) -> Result<()> {
+    let registrar = ctx.accounts.registrar.load()?;
+
+    require_keys_eq!(
+        registrar.realm_authority,
+        ctx.accounts.realm_authority.key(),
+        MplStakingError::InvalidRealmAuthority
+    );
+
+    let signers_seeds = &[
+        &registrar.realm.key().to_bytes(),
+        b"registrar".as_ref(),
+        &registrar.realm_governing_token_mint.key().to_bytes(),
+        &[registrar.bump][..],
+    ];
+
+    cpi_instructions::allow_tokenflow(
+        ctx.accounts.rewards_program.to_account_info(),
+        ctx.accounts.registrar.to_account_info(),
+        ctx.accounts.reward_pool.to_account_info(),
+        ctx.accounts.deposit_mining.to_account_info(),
+        &mining_owner,
+        signers_seeds,
+    )?;
+
+    Ok(())
+}

programs/mpl-staking/src/instructions/penalties/mod.rs

@@ -0,0 +1,36 @@
+pub use allow_tokenflow::*;
+use anchor_lang::{
+    prelude::{AccountLoader, Signer, ToAccountInfo, UncheckedAccount},
+    Accounts,
+};
+pub use restrict_batch_minting::*;
+pub use restrict_tokenflow::*;
+
+mod allow_tokenflow;
+mod restrict_batch_minting;
+mod restrict_tokenflow;
+
+use mplx_staking_states::state::Registrar;
+
+#[derive(Accounts)]
+pub struct Penalty<'info> {
+    pub registrar: AccountLoader<'info, Registrar>,
+
+    pub realm_authority: Signer<'info>,
+
+    /// CHECK:
+    /// Ownership of the account will be checked in the rewards contract
+    /// It's the core account for the rewards contract, which will
+    /// keep track of all rewards and staking logic.
+    pub reward_pool: UncheckedAccount<'info>,
+
+    /// CHECK: mining PDA will be checked in the rewards contract
+    /// PDA(["mining", mining owner <aka voter_authority in our case>, reward_pool],
+    /// reward_program)
+    #[account(mut)]
+    pub deposit_mining: UncheckedAccount<'info>,
+
+    /// CHECK: Rewards Program account
+    #[account(executable)]
+    pub rewards_program: UncheckedAccount<'info>,
+}

programs/mpl-staking/src/instructions/penalties/restrict_batch_minting.rs

@@ -0,0 +1,38 @@
+use super::Penalty;
+use crate::cpi_instructions;
+use anchor_lang::prelude::*;
+use mplx_staking_states::error::MplStakingError;
+
+/// Restricts batch minting operation for the account until the specified timestamp.
+pub fn restrict_batch_minting(
+    ctx: Context<Penalty>,
+    until_ts: u64,
+    mining_owner: Pubkey,
+) -> Result<()> {
+    let registrar = ctx.accounts.registrar.load()?;
+
+    require_keys_eq!(
+        registrar.realm_authority,
+        ctx.accounts.realm_authority.key(),
+        MplStakingError::InvalidRealmAuthority
+    );
+
+    let signers_seeds = &[
+        &registrar.realm.key().to_bytes(),
+        b"registrar".as_ref(),
+        &registrar.realm_governing_token_mint.key().to_bytes(),
+        &[registrar.bump][..],
+    ];
+
+    cpi_instructions::restrict_batch_minting(
+        ctx.accounts.rewards_program.to_account_info(),
+        ctx.accounts.registrar.to_account_info(),
+        ctx.accounts.reward_pool.to_account_info(),
+        ctx.accounts.deposit_mining.to_account_info(),
+        &mining_owner,
+        until_ts,
+        signers_seeds,
+    )?;
+
+    Ok(())
+}

programs/mpl-staking/src/instructions/penalties/restrict_tokenflow.rs

@@ -0,0 +1,33 @@
+use super::Penalty;
+use crate::cpi_instructions;
+use anchor_lang::prelude::*;
+use mplx_staking_states::error::MplStakingError;
+
+/// Restricts claiming rewards from the specified mining account.
+pub fn restrict_tokenflow(ctx: Context<Penalty>, mining_owner: Pubkey) -> Result<()> {
+    let registrar = ctx.accounts.registrar.load()?;
+
+    require_keys_eq!(
+        registrar.realm_authority,
+        ctx.accounts.realm_authority.key(),
+        MplStakingError::InvalidRealmAuthority
+    );
+
+    let signers_seeds = &[
+        &registrar.realm.key().to_bytes(),
+        b"registrar".as_ref(),
+        &registrar.realm_governing_token_mint.key().to_bytes(),
+        &[registrar.bump][..],
+    ];
+
+    cpi_instructions::restrict_tokenflow(
+        ctx.accounts.rewards_program.to_account_info(),
+        ctx.accounts.registrar.to_account_info(),
+        ctx.accounts.reward_pool.to_account_info(),
+        ctx.accounts.deposit_mining.to_account_info(),
+        &mining_owner,
+        signers_seeds,
+    )?;
+
+    Ok(())
+}

programs/mpl-staking/src/lib.rs

@@ -180,6 +180,22 @@ pub mod mpl_staking {
     pub fn change_delegate(ctx: Context<ChangeDelegate>, deposit_entry_index: u8) -> Result<()> {
         instructions::change_delegate(ctx, deposit_entry_index)
     }
+
+    pub fn restrict_tokenflow(ctx: Context<Penalty>, mining_owner: Pubkey) -> Result<()> {
+        instructions::restrict_tokenflow(ctx, mining_owner)
+    }
+
+    pub fn allow_tokenflow(ctx: Context<Penalty>, mining_owner: Pubkey) -> Result<()> {
+        instructions::allow_tokenflow(ctx, mining_owner)
+    }
+
+    pub fn restrict_batch_minting(
+        ctx: Context<Penalty>,
+        until_ts: u64,
+        mining_owner: Pubkey,
+    ) -> Result<()> {
+        instructions::restrict_batch_minting(ctx, until_ts, mining_owner)
+    }
 }
 
 #[derive(Accounts)]