added delegate_mining check (#38)

added `delegate_mining` check

programs/rewards/src/error.rs

@@ -80,6 +80,16 @@ pub enum MplxRewardsError {
     /// Account is already initialized
     #[error("Account is already initialized")]
     AlreadyInitialized,
+
+    /// 19
+    /// Incorrect mining address.
+    #[error("Invalid mining")]
+    InvalidMining,
+
+    /// 20
+    /// Failed to derive PDA.
+    #[error("Failed to derive PDA")]
+    DerivationError,
 }
 
 impl PrintProgramError for MplxRewardsError {

programs/rewards/src/instruction.rs

@@ -65,6 +65,7 @@ pub enum RewardsInstruction {
         lockup_period: LockupPeriod,
         /// Specifies the owner of the Mining Account
         mining_owner: Pubkey,
+        delegate: Pubkey,
     },
 
     /// Withdraws amount of supply to the mining account
@@ -77,6 +78,7 @@ pub enum RewardsInstruction {
         amount: u64,
         /// Specifies the owner of the Mining Account
         mining_owner: Pubkey,
+        delegate: Pubkey,
     },
 
     /// Claims amount of rewards
@@ -113,6 +115,8 @@ pub enum RewardsInstruction {
         additional_amount: u64,
         /// The wallet who owns the mining account
         mining_owner: Pubkey,
+        /// Wallet addres of delegate
+        delegate: Pubkey,
     },
 
     /// Distributes tokens among mining owners
@@ -138,6 +142,7 @@ pub enum RewardsInstruction {
     ChangeDelegate {
         /// Amount of staked tokens
         staked_amount: u64,
+        new_delegate: Pubkey,
     },
 }
 
@@ -240,6 +245,7 @@ pub fn deposit_mining(
     amount: u64,
     lockup_period: LockupPeriod,
     mining_owner: &Pubkey,
+    delegate: &Pubkey,
 ) -> Instruction {
     let accounts = vec![
         AccountMeta::new(*reward_pool, false),
@@ -254,12 +260,14 @@ pub fn deposit_mining(
             amount,
             lockup_period,
             mining_owner: *mining_owner,
+            delegate: *delegate,
         },
         accounts,
     )
 }
 
 /// Creates 'WithdrawMining' instruction.
+#[allow(clippy::too_many_arguments)]
 pub fn withdraw_mining(
     program_id: &Pubkey,
     reward_pool: &Pubkey,
@@ -268,6 +276,7 @@ pub fn withdraw_mining(
     delegate_mining: &Pubkey,
     amount: u64,
     mining_owner: &Pubkey,
+    delegate: &Pubkey,
 ) -> Instruction {
     let accounts = vec![
         AccountMeta::new(*reward_pool, false),
@@ -281,6 +290,7 @@ pub fn withdraw_mining(
         &RewardsInstruction::WithdrawMining {
             amount,
             mining_owner: *mining_owner,
+            delegate: *delegate,
         },
         accounts,
     )
@@ -326,6 +336,7 @@ pub fn extend_stake(
     base_amount: u64,
     additional_amount: u64,
     mining_owner: &Pubkey,
+    delegate: &Pubkey,
 ) -> Instruction {
     let accounts = vec![
         AccountMeta::new(*reward_pool, false),
@@ -343,6 +354,7 @@ pub fn extend_stake(
             base_amount,
             additional_amount,
             mining_owner: *mining_owner,
+            delegate: *delegate,
         },
         accounts,
     )
@@ -398,6 +410,7 @@ pub fn change_delegate(
     mining_owner: &Pubkey,
     old_delegate_mining: &Pubkey,
     new_delegate_mining: &Pubkey,
+    new_delegate: &Pubkey,
     staked_amount: u64,
 ) -> Instruction {
     let accounts = vec![
@@ -411,7 +424,10 @@ pub fn change_delegate(
 
     Instruction::new_with_borsh(
         *program_id,
-        &RewardsInstruction::ChangeDelegate { staked_amount },
+        &RewardsInstruction::ChangeDelegate {
+            staked_amount,
+            new_delegate: *new_delegate,
+        },
         accounts,
     )
 }

programs/rewards/src/instructions/change_delegate.rs

@@ -1,14 +1,15 @@
 use crate::{
     asserts::assert_and_get_pool_and_mining,
     error::MplxRewardsError,
-    utils::{get_delegate_mining, AccountLoader},
+    utils::{get_delegate_mining, verify_delegate_mining_address, AccountLoader},
 };
 use solana_program::{account_info::AccountInfo, entrypoint::ProgramResult, pubkey::Pubkey};
 
 pub fn process_change_delegate<'a>(
     program_id: &Pubkey,
     accounts: &'a [AccountInfo<'a>],
     staked_amount: u64,
+    new_delegate: &Pubkey,
 ) -> ProgramResult {
     let account_info_iter = &mut accounts.iter().enumerate();
 
@@ -37,6 +38,15 @@ pub fn process_change_delegate<'a>(
     )?;
 
     let new_delegate_mining = get_delegate_mining(new_delegate_mining, mining)?;
+    if let Some(new_delegate_mining) = new_delegate_mining {
+        verify_delegate_mining_address(
+            program_id,
+            new_delegate_mining,
+            new_delegate,
+            reward_pool.key,
+        )?
+    }
+
     let old_delegate_mining = get_delegate_mining(old_delegate_mining, mining)?;
 
     wrapped_reward_pool.change_delegate(

programs/rewards/src/instructions/deposit_mining.rs

@@ -1,6 +1,6 @@
 use crate::{
     asserts::assert_and_get_pool_and_mining,
-    utils::{get_delegate_mining, AccountLoader, LockupPeriod},
+    utils::{get_delegate_mining, verify_delegate_mining_address, AccountLoader, LockupPeriod},
 };
 use solana_program::{account_info::AccountInfo, entrypoint::ProgramResult, pubkey::Pubkey};
 
@@ -10,6 +10,7 @@ pub fn process_deposit_mining<'a>(
     amount: u64,
     lockup_period: LockupPeriod,
     mining_owner: &Pubkey,
+    delegate: &Pubkey,
 ) -> ProgramResult {
     let account_info_iter = &mut accounts.iter().enumerate();
 
@@ -32,6 +33,9 @@ pub fn process_deposit_mining<'a>(
     )?;
 
     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)?
+    }
 
     wrapped_reward_pool.deposit(&mut wrapped_mining, amount, lockup_period, delegate_mining)?;
 

programs/rewards/src/instructions/extend_stake.rs

@@ -1,6 +1,6 @@
 use crate::{
     asserts::assert_and_get_pool_and_mining,
-    utils::{get_delegate_mining, AccountLoader, LockupPeriod},
+    utils::{get_delegate_mining, verify_delegate_mining_address, AccountLoader, LockupPeriod},
 };
 use solana_program::{account_info::AccountInfo, entrypoint::ProgramResult, pubkey::Pubkey};
 
@@ -14,6 +14,7 @@ pub fn process_extend_stake<'a>(
     base_amount: u64,
     additional_amount: u64,
     mining_owner: &Pubkey,
+    delegate: &Pubkey,
 ) -> ProgramResult {
     let account_info_iter = &mut accounts.iter().enumerate();
 
@@ -37,6 +38,10 @@ pub fn process_extend_stake<'a>(
 
     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)?
+    }
+
     wrapped_reward_pool.extend(
         &mut wrapped_mining,
         old_lockup_period,

programs/rewards/src/instructions/mod.rs

@@ -55,16 +55,25 @@ pub fn process_instruction<'a>(
             amount,
             lockup_period,
             mining_owner,
+            delegate,
         } => {
             msg!("RewardsInstruction: DepositMining");
-            process_deposit_mining(program_id, accounts, amount, lockup_period, &mining_owner)
+            process_deposit_mining(
+                program_id,
+                accounts,
+                amount,
+                lockup_period,
+                &mining_owner,
+                &delegate,
+            )
         }
         RewardsInstruction::WithdrawMining {
             amount,
             mining_owner,
+            delegate,
         } => {
             msg!("RewardsInstruction: WithdrawMining");
-            process_withdraw_mining(program_id, accounts, amount, &mining_owner)
+            process_withdraw_mining(program_id, accounts, amount, &mining_owner, &delegate)
         }
         RewardsInstruction::Claim => {
             msg!("RewardsInstruction: Claim");
@@ -77,6 +86,7 @@ pub fn process_instruction<'a>(
             base_amount,
             additional_amount,
             mining_owner,
+            delegate,
         } => {
             msg!("RewardsInstruction: ExtendStake");
             process_extend_stake(
@@ -88,6 +98,7 @@ pub fn process_instruction<'a>(
                 base_amount,
                 additional_amount,
                 &mining_owner,
+                &delegate,
             )
         }
         RewardsInstruction::DistributeRewards => {
@@ -98,9 +109,12 @@ pub fn process_instruction<'a>(
             msg!("RewardsInstruction: CloseAccount");
             process_close_mining(program_id, accounts)
         }
-        RewardsInstruction::ChangeDelegate { staked_amount } => {
+        RewardsInstruction::ChangeDelegate {
+            staked_amount,
+            new_delegate,
+        } => {
             msg!("RewardsInstruction: ChangeDelegate");
-            process_change_delegate(program_id, accounts, staked_amount)
+            process_change_delegate(program_id, accounts, staked_amount, &new_delegate)
         }
     }
 }

programs/rewards/src/instructions/withdraw_mining.rs

@@ -3,13 +3,15 @@ use crate::{
     utils::{get_delegate_mining, AccountLoader},
 };
 
+use crate::utils::verify_delegate_mining_address;
 use solana_program::{account_info::AccountInfo, entrypoint::ProgramResult, pubkey::Pubkey};
 
 pub fn process_withdraw_mining<'a>(
     program_id: &Pubkey,
     accounts: &'a [AccountInfo<'a>],
     amount: u64,
     mining_owner: &Pubkey,
+    delegate: &Pubkey,
 ) -> ProgramResult {
     let account_info_iter = &mut accounts.iter().enumerate();
 
@@ -32,6 +34,9 @@ pub fn process_withdraw_mining<'a>(
     )?;
 
     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)?
+    }
 
     wrapped_reward_pool.withdraw(&mut wrapped_mining, amount, delegate_mining)?;
 

programs/rewards/src/utils.rs

@@ -1,7 +1,7 @@
 //! Arbitrary auxilliary functions
 use std::iter::Enumerate;
 
-use crate::error::MplxRewardsError;
+use crate::{error::MplxRewardsError, state::WrappedImmutableMining};
 use borsh::{BorshDeserialize, BorshSerialize};
 use solana_program::{
     account_info::AccountInfo,
@@ -11,7 +11,7 @@ use solana_program::{
     program::{invoke, invoke_signed},
     program_error::ProgramError,
     program_pack::Pack,
-    pubkey::Pubkey,
+    pubkey::{Pubkey, PubkeyError},
     rent::Rent,
     system_instruction,
     sysvar::Sysvar,
@@ -118,6 +118,29 @@ pub fn get_delegate_mining<'a, 'b>(
     }
 }
 
+pub fn verify_delegate_mining_address(
+    program_id: &Pubkey,
+    delegate_mining: &AccountInfo<'_>,
+    delegate: &Pubkey,
+    reward_pool_key: &Pubkey,
+) -> Result<(), ProgramError> {
+    if *delegate_mining.key
+        != create_mining_address(
+            program_id,
+            delegate,
+            reward_pool_key,
+            WrappedImmutableMining::from_bytes(&delegate_mining.data.borrow())?
+                .mining
+                .bump,
+        )
+        .map_err(|_| MplxRewardsError::DerivationError)?
+    {
+        return Err(MplxRewardsError::InvalidMining.into());
+    }
+
+    Ok(())
+}
+
 /// Helper for parsing accounts with arbitrary input conditions
 pub struct AccountLoader {}
 
@@ -321,3 +344,20 @@ impl SafeArithmeticOperations for u128 {
             .ok_or(MplxRewardsError::MathOverflow)
     }
 }
+
+pub fn create_mining_address(
+    program_id: &Pubkey,
+    mining_owner: &Pubkey,
+    reward_pool: &Pubkey,
+    bump: u8,
+) -> Result<Pubkey, PubkeyError> {
+    Pubkey::create_program_address(
+        &[
+            "mining".as_bytes(),
+            &mining_owner.to_bytes(),
+            &reward_pool.to_bytes(),
+            &[bump],
+        ],
+        program_id,
+    )
+}