From 109786a172579acd1b9b862630459c23ce38832a Mon Sep 17 00:00:00 2001
From: rwwwx <vlad.matiukhin@gmail.com>
Date: Mon, 26 Aug 2024 15:40:12 +0300
Subject: [PATCH] added `verify_delegate_mining_address` for change_delegate

---
 programs/rewards/src/instruction.rs                  |  7 ++++++-
 programs/rewards/src/instructions/change_delegate.rs | 12 +++++++++++-
 programs/rewards/src/instructions/mod.rs             | 12 ++++++++++--
 programs/rewards/tests/rewards/change_delegate.rs    |  2 ++
 programs/rewards/tests/rewards/utils.rs              |  2 ++
 5 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/programs/rewards/src/instruction.rs b/programs/rewards/src/instruction.rs
index bfaa7772..a1c6954d 100644
--- a/programs/rewards/src/instruction.rs
+++ b/programs/rewards/src/instruction.rs
@@ -142,6 +142,7 @@ pub enum RewardsInstruction {
     ChangeDelegate {
         /// Amount of staked tokens
         staked_amount: u64,
+        new_delegate: Pubkey,
     },
 }
 
@@ -409,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![
@@ -422,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,
     )
 }
diff --git a/programs/rewards/src/instructions/change_delegate.rs b/programs/rewards/src/instructions/change_delegate.rs
index b88584fb..680c150d 100644
--- a/programs/rewards/src/instructions/change_delegate.rs
+++ b/programs/rewards/src/instructions/change_delegate.rs
@@ -1,7 +1,7 @@
 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};
 
@@ -9,6 +9,7 @@ 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(
diff --git a/programs/rewards/src/instructions/mod.rs b/programs/rewards/src/instructions/mod.rs
index ee9877ec..3d16717d 100644
--- a/programs/rewards/src/instructions/mod.rs
+++ b/programs/rewards/src/instructions/mod.rs
@@ -109,9 +109,17 @@ 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,
+            )
         }
     }
 }
diff --git a/programs/rewards/tests/rewards/change_delegate.rs b/programs/rewards/tests/rewards/change_delegate.rs
index 203356e7..2a6bfdcc 100644
--- a/programs/rewards/tests/rewards/change_delegate.rs
+++ b/programs/rewards/tests/rewards/change_delegate.rs
@@ -70,6 +70,7 @@ async fn change_delegate_to_the_same() {
             &user_a,
             &user_mining_a,
             &user_mining_a,
+            &user_a.pubkey(),
             6_000_000,
         )
         .await
@@ -121,6 +122,7 @@ async fn change_delegate_then_claim() {
             &user_a,
             &delegate_mining,
             &user_mining_a,
+            &delegate.pubkey(),
             1_000_000,
         )
         .await
diff --git a/programs/rewards/tests/rewards/utils.rs b/programs/rewards/tests/rewards/utils.rs
index 127215cf..57c1ddd7 100644
--- a/programs/rewards/tests/rewards/utils.rs
+++ b/programs/rewards/tests/rewards/utils.rs
@@ -126,6 +126,7 @@ impl TestRewards {
         mining_owner: &Keypair,
         new_delegate_mining: &Pubkey,
         old_delegate_mining: &Pubkey,
+        new_delegate: &Pubkey,
         amount: u64,
     ) -> BanksClientResult<()> {
         let tx = Transaction::new_signed_with_payer(
@@ -137,6 +138,7 @@ impl TestRewards {
                 &mining_owner.pubkey(),
                 old_delegate_mining,
                 new_delegate_mining,
+                new_delegate,
                 amount,
             )],
             Some(&context.payer.pubkey()),