Merge pull request #39 from adm-metaex/penalties/block-the-clam-opera…

…tion

[mtg-578] Penalties/block the clam operation

programs/rewards/src/error.rs

@@ -90,6 +90,25 @@ pub enum MplxRewardsError {
     /// Failed to derive PDA.
     #[error("Failed to derive PDA")]
     DerivationError,
+
+    /// 21
+    #[error("Mining already restricted")]
+    MiningAlreadyRestricted,
+
+    /// 22
+    /// Mining is not restricted
+    #[error("Mining is not restricted")]
+    MiningNotRestricted,
+
+    /// 23
+    /// Claiming is restricted
+    #[error("Claiming is restricted")]
+    ClaimingRestricted,
+
+    /// 24
+    /// Withdrawal is restricted while claiming is restricted
+    #[error("Withdrawal is restricted while claiming is restricted")]
+    WithdrawalRestricted,
 }
 
 impl PrintProgramError for MplxRewardsError {

programs/rewards/src/instruction.rs

@@ -144,6 +144,21 @@ pub enum RewardsInstruction {
         staked_amount: u64,
         new_delegate: Pubkey,
     },
+
+    /// Prevents the mining account from rewards withdrawing
+    #[account(0, signer, name = "deposit_authority", desc = "The address of the Staking program's Registrar, which is PDA and is responsible for signing CPIs")]
+    #[account(1, name = "reward_pool", desc = "The address of the reward pool")]
+    #[account(2, writable, name = "mining", desc = "The address of the mining account which belongs to the user and stores info about user's rewards")]
+    RestrictTokenFlow {
+        mining_owner: Pubkey,
+    },
+
+    #[account(0, signer, name = "deposit_authority", desc = "The address of the Staking program's Registrar, which is PDA and is responsible for signing CPIs")]
+    #[account(1, name = "reward_pool", desc = "The address of the reward pool")]
+    #[account(2, writable, name = "mining", desc = "The address of the mining account which belongs to the user and stores info about user's rewards")]
+    AllowTokenFlow {
+        mining_owner: Pubkey,
+    },
 }
 
 /// Creates 'InitializePool' instruction.
@@ -431,3 +446,47 @@ pub fn change_delegate(
         accounts,
     )
 }
+
+pub fn restrict_tokenflow(
+    program_id: &Pubkey,
+    deposit_authority: &Pubkey,
+    reward_pool: &Pubkey,
+    mining: &Pubkey,
+    mining_owner: &Pubkey,
+) -> Instruction {
+    let accounts = vec![
+        AccountMeta::new_readonly(*deposit_authority, true),
+        AccountMeta::new_readonly(*reward_pool, false),
+        AccountMeta::new(*mining, false),
+    ];
+
+    Instruction::new_with_borsh(
+        *program_id,
+        &RewardsInstruction::RestrictTokenFlow {
+            mining_owner: *mining_owner,
+        },
+        accounts,
+    )
+}
+
+pub fn allow_tokenflow(
+    program_id: &Pubkey,
+    deposit_authority: &Pubkey,
+    reward_pool: &Pubkey,
+    mining: &Pubkey,
+    mining_owner: &Pubkey,
+) -> Instruction {
+    let accounts = vec![
+        AccountMeta::new_readonly(*deposit_authority, true),
+        AccountMeta::new_readonly(*reward_pool, false),
+        AccountMeta::new(*mining, false),
+    ];
+
+    Instruction::new_with_borsh(
+        *program_id,
+        &RewardsInstruction::AllowTokenFlow {
+            mining_owner: *mining_owner,
+        },
+        accounts,
+    )
+}

programs/rewards/src/instructions/claim.rs

@@ -1,5 +1,6 @@
 use crate::{
     asserts::{assert_account_key, assert_account_owner},
+    error::MplxRewardsError,
     state::{WrappedMining, WrappedRewardPool},
     utils::{spl_transfer, AccountLoader},
 };
@@ -42,6 +43,10 @@ pub fn process_claim<'a>(program_id: &Pubkey, accounts: &'a [AccountInfo<'a>]) -
             let mining_data = &mut mining.data.borrow_mut();
             let mut wrapped_mining = WrappedMining::from_bytes_mut(mining_data)?;
 
+            if wrapped_mining.mining.is_tokenflow_restricted() {
+                return Err(MplxRewardsError::ClaimingRestricted.into());
+            }
+
             assert_account_owner(reward_pool, program_id)?;
             assert_account_key(mining_owner, &wrapped_mining.mining.owner)?;
             assert_account_key(reward_pool, &wrapped_mining.mining.reward_pool)?;

programs/rewards/src/instructions/mod.rs

@@ -12,6 +12,7 @@ mod extend_stake;
 mod fill_vault;
 mod initialize_mining;
 mod initialize_pool;
+mod penalties;
 mod withdraw_mining;
 
 pub(crate) use change_delegate::*;
@@ -23,6 +24,7 @@ pub(crate) use extend_stake::*;
 pub(crate) use fill_vault::*;
 pub(crate) use initialize_mining::*;
 pub(crate) use initialize_pool::*;
+pub(crate) use penalties::*;
 pub(crate) use withdraw_mining::*;
 
 pub fn process_instruction<'a>(
@@ -116,5 +118,13 @@ pub fn process_instruction<'a>(
             msg!("RewardsInstruction: ChangeDelegate");
             process_change_delegate(program_id, accounts, staked_amount, &new_delegate)
         }
+        RewardsInstruction::RestrictTokenFlow { mining_owner } => {
+            msg!("RewardsInstruction: RestrictClaiming");
+            process_restrict_tokenflow(program_id, accounts, &mining_owner)
+        }
+        RewardsInstruction::AllowTokenFlow { mining_owner } => {
+            msg!("RewardsInstruction: AllowClaiming");
+            process_allow_tokenflow(program_id, accounts, &mining_owner)
+        }
     }
 }

programs/rewards/src/instructions/penalties/allow_tokenflow.rs

@@ -0,0 +1,31 @@
+use crate::{asserts::assert_and_get_pool_and_mining, utils::AccountLoader};
+use solana_program::{account_info::AccountInfo, entrypoint::ProgramResult, pubkey::Pubkey};
+
+pub fn process_allow_tokenflow<'a>(
+    program_id: &Pubkey,
+    accounts: &'a [AccountInfo<'a>],
+    mining_owner: &Pubkey,
+) -> ProgramResult {
+    let account_info_iter = &mut accounts.iter().enumerate();
+
+    let deposit_authority = AccountLoader::next_signer(account_info_iter)?;
+    let reward_pool = AccountLoader::next_with_owner(account_info_iter, program_id)?;
+    let mining = AccountLoader::next_with_owner(account_info_iter, program_id)?;
+
+    let reward_pool_data = &mut reward_pool.data.borrow_mut();
+    let mining_data = &mut mining.data.borrow_mut();
+
+    let (_, wrapped_mining) = assert_and_get_pool_and_mining(
+        program_id,
+        mining_owner,
+        mining,
+        reward_pool,
+        deposit_authority,
+        reward_pool_data,
+        mining_data,
+    )?;
+
+    wrapped_mining.mining.allow_tokenflow()?;
+
+    Ok(())
+}

programs/rewards/src/instructions/penalties/mod.rs

@@ -0,0 +1,5 @@
+pub(crate) use allow_tokenflow::*;
+pub(crate) use restrict_tokenflow::*;
+
+mod allow_tokenflow;
+mod restrict_tokenflow;

programs/rewards/src/instructions/penalties/restrict_tokenflow.rs

@@ -0,0 +1,31 @@
+use crate::{asserts::assert_and_get_pool_and_mining, utils::AccountLoader};
+use solana_program::{account_info::AccountInfo, entrypoint::ProgramResult, pubkey::Pubkey};
+
+pub fn process_restrict_tokenflow<'a>(
+    program_id: &Pubkey,
+    accounts: &'a [AccountInfo<'a>],
+    mining_owner: &Pubkey,
+) -> ProgramResult {
+    let account_info_iter = &mut accounts.iter().enumerate();
+
+    let deposit_authority = AccountLoader::next_signer(account_info_iter)?;
+    let reward_pool = AccountLoader::next_with_owner(account_info_iter, program_id)?;
+    let mining = AccountLoader::next_with_owner(account_info_iter, program_id)?;
+
+    let reward_pool_data = &mut reward_pool.data.borrow_mut();
+    let mining_data = &mut mining.data.borrow_mut();
+
+    let (_, wrapped_mining) = assert_and_get_pool_and_mining(
+        program_id,
+        mining_owner,
+        mining,
+        reward_pool,
+        deposit_authority,
+        reward_pool_data,
+        mining_data,
+    )?;
+
+    wrapped_mining.mining.restrict_tokenflow()?;
+
+    Ok(())
+}

programs/rewards/src/instructions/withdraw_mining.rs

@@ -1,5 +1,6 @@
 use crate::{
     asserts::assert_and_get_pool_and_mining,
+    error::MplxRewardsError,
     utils::{get_delegate_mining, AccountLoader},
 };
 
@@ -33,6 +34,10 @@ pub fn process_withdraw_mining<'a>(
         mining_data,
     )?;
 
+    if wrapped_mining.mining.is_tokenflow_restricted() {
+        return Err(MplxRewardsError::WithdrawalRestricted.into());
+    }
+
     let delegate_mining = get_delegate_mining(delegate_mining, mining)?;
     if let Some(delegate_mining) = delegate_mining {
         verify_delegate_mining_address(program_id, delegate_mining, delegate, reward_pool.key)?

programs/rewards/src/state/mining.rs

@@ -32,6 +32,9 @@ pub struct WrappedImmutableMining<'a> {
     pub weighted_stake_diffs: &'a MiningWeightedStakeDiffs,
 }
 
+pub const ACCOUNT_TYPE_BYTE: usize = 0;
+pub const CLAIMING_RESTRICTION_BYTE: usize = 1;
+
 impl<'a> WrappedMining<'a> {
     pub const LEN: usize = 1776;
 
@@ -99,8 +102,9 @@ pub struct Mining {
     pub bump: u8,
     /// Account type - Mining. This discriminator should exist in order to prevent
     /// shenanigans with customly modified accounts and their fields.
-    /// 1: account type
-    /// 2-7: unused
+    /// 0: account type
+    /// 1: claim is restricted
+    /// 2-6: unused
     pub data: [u8; 7],
 }
 
@@ -114,7 +118,7 @@ impl Mining {
     pub fn initialize(reward_pool: Pubkey, owner: Pubkey, bump: u8) -> Mining {
         let account_type = AccountType::Mining.into();
         let mut data = [0; 7];
-        data[0] = account_type;
+        data[ACCOUNT_TYPE_BYTE] = account_type;
         Mining {
             data,
             reward_pool,
@@ -125,7 +129,7 @@ impl Mining {
     }
 
     pub fn account_type(&self) -> AccountType {
-        AccountType::from(self.data[0])
+        AccountType::from(self.data[ACCOUNT_TYPE_BYTE])
     }
 
     /// Claim reward
@@ -193,11 +197,33 @@ impl Mining {
 
         Ok(())
     }
+
+    pub fn restrict_tokenflow(&mut self) -> ProgramResult {
+        if self.data[CLAIMING_RESTRICTION_BYTE] == 1 {
+            Err(MplxRewardsError::MiningAlreadyRestricted.into())
+        } else {
+            self.data[CLAIMING_RESTRICTION_BYTE] = 1;
+            Ok(())
+        }
+    }
+
+    pub fn allow_tokenflow(&mut self) -> ProgramResult {
+        if self.data[CLAIMING_RESTRICTION_BYTE] == 0 {
+            Err(MplxRewardsError::MiningNotRestricted.into())
+        } else {
+            self.data[CLAIMING_RESTRICTION_BYTE] = 0;
+            Ok(())
+        }
+    }
+
+    pub fn is_tokenflow_restricted(&self) -> bool {
+        self.data[CLAIMING_RESTRICTION_BYTE] == 1
+    }
 }
 
 impl IsInitialized for Mining {
     fn is_initialized(&self) -> bool {
-        self.data[0] == <u8>::from(AccountType::Mining)
+        self.data[ACCOUNT_TYPE_BYTE] == <u8>::from(AccountType::Mining)
     }
 }
 

programs/rewards/tests/rewards/tests.rs

@@ -11,4 +11,6 @@ mod precision;
 mod utils;
 mod withdraw_mining;
 
+mod tokenflow_restrictions;
+
 mod extend_stake;