From bfc5ca65d3caaad93e503e0305a15f10154778c4 Mon Sep 17 00:00:00 2001
From: Keyvan Khademi <keyvankhademi@gmail.com>
Date: Wed, 4 Sep 2024 08:13:14 -0700
Subject: [PATCH] feat: add update reward program authority (#519)

* feat: add update reward program authority

* fix: typo

* refactor: update reward program authority in initialize_pool.rs
---
 .../src/integrity_pool/instructions.rs        | 34 +++++++++++++
 .../tests/initialize_pool.rs                  | 49 +++++++++++++++++++
 .../programs/integrity-pool/src/context.rs    | 15 ++++++
 staking/programs/integrity-pool/src/lib.rs    |  8 +++
 4 files changed, 106 insertions(+)

diff --git a/staking/integration-tests/src/integrity_pool/instructions.rs b/staking/integration-tests/src/integrity_pool/instructions.rs
index cccf2ad6..82b16afc 100644
--- a/staking/integration-tests/src/integrity_pool/instructions.rs
+++ b/staking/integration-tests/src/integrity_pool/instructions.rs
@@ -167,6 +167,40 @@ pub fn update_y(
     svm.send_transaction(update_y_tx.clone())
 }
 
+pub fn update_reward_program_authority(
+    svm: &mut litesvm::LiteSVM,
+    payer: &Keypair,
+    reward_program_authority: &Keypair,
+    new_reward_program_authority: Pubkey,
+) -> TransactionResult {
+    let pool_config_pubkey = get_pool_config_address();
+
+    let instruction_data = integrity_pool::instruction::UpdateRewardProgramAuthority {
+        reward_program_authority: new_reward_program_authority,
+    };
+
+    let instruction_accs = integrity_pool::accounts::UpdateRewardProgramAuthority {
+        pool_config:              pool_config_pubkey,
+        reward_program_authority: reward_program_authority.pubkey(),
+        system_program:           system_program::ID,
+    };
+
+    let instruction = Instruction::new_with_bytes(
+        integrity_pool::ID,
+        &instruction_data.data(),
+        instruction_accs.to_account_metas(None),
+    );
+
+    let transaction = Transaction::new_signed_with_payer(
+        &[instruction],
+        Some(&payer.pubkey()),
+        &[payer, reward_program_authority],
+        svm.latest_blockhash(),
+    );
+
+    svm.send_transaction(transaction)
+}
+
 pub fn update_delegation_fee(
     svm: &mut litesvm::LiteSVM,
     payer: &Keypair,
diff --git a/staking/integration-tests/tests/initialize_pool.rs b/staking/integration-tests/tests/initialize_pool.rs
index 33742903..30ee8be2 100644
--- a/staking/integration-tests/tests/initialize_pool.rs
+++ b/staking/integration-tests/tests/initialize_pool.rs
@@ -4,6 +4,7 @@ use {
         integrity_pool::{
             instructions::{
                 create_pool_data_account,
+                update_reward_program_authority,
                 update_y,
             },
             pda::get_pool_config_address,
@@ -116,3 +117,51 @@ fn test_update_y() {
         0
     );
 }
+
+#[test]
+fn test_update_reward_program_authority() {
+    let SetupResult {
+        mut svm,
+        payer,
+        reward_program_authority,
+        ..
+    } = setup(SetupProps {
+        init_config:            true,
+        init_target:            true,
+        init_mint:              true,
+        init_pool_data:         true,
+        init_publishers:        true,
+        reward_amount_override: None,
+    });
+
+    let pool_config_pubkey = get_pool_config_address();
+    let pool_config: PoolConfig = fetch_account_data(&mut svm, &pool_config_pubkey);
+
+    assert!(pool_config.reward_program_authority == reward_program_authority.pubkey());
+
+    let new_reward_program_authority = Pubkey::new_unique();
+    update_reward_program_authority(
+        &mut svm,
+        &payer,
+        &reward_program_authority,
+        new_reward_program_authority,
+    )
+    .unwrap();
+
+    let pool_config: PoolConfig = fetch_account_data(&mut svm, &pool_config_pubkey);
+    assert!(pool_config.reward_program_authority == new_reward_program_authority);
+
+    // Trying to update the reward program authority without the correct authority should fail
+    let new_reward_program_authority2 = Keypair::new();
+
+    assert_anchor_program_error!(
+        update_reward_program_authority(
+            &mut svm,
+            &payer,
+            &reward_program_authority,
+            new_reward_program_authority2.pubkey(),
+        ),
+        IntegrityPoolError::InvalidRewardProgramAuthority,
+        0
+    );
+}
diff --git a/staking/programs/integrity-pool/src/context.rs b/staking/programs/integrity-pool/src/context.rs
index 64e244bb..400cd03d 100644
--- a/staking/programs/integrity-pool/src/context.rs
+++ b/staking/programs/integrity-pool/src/context.rs
@@ -53,6 +53,21 @@ pub struct UpdateY<'info> {
     pub system_program: Program<'info, System>,
 }
 
+#[derive(Accounts)]
+pub struct UpdateRewardProgramAuthority<'info> {
+    pub reward_program_authority: Signer<'info>,
+
+    #[account(
+        mut,
+        seeds = [POOL_CONFIG.as_bytes()],
+        bump,
+        has_one = reward_program_authority @ IntegrityPoolError::InvalidRewardProgramAuthority,
+    )]
+    pub pool_config: Account<'info, PoolConfig>,
+
+    pub system_program: Program<'info, System>,
+}
+
 #[derive(Accounts)]
 pub struct UpdateDelegationFee<'info> {
     pub reward_program_authority: Signer<'info>,
diff --git a/staking/programs/integrity-pool/src/lib.rs b/staking/programs/integrity-pool/src/lib.rs
index bbf5c0f5..65e60fb6 100644
--- a/staking/programs/integrity-pool/src/lib.rs
+++ b/staking/programs/integrity-pool/src/lib.rs
@@ -57,6 +57,14 @@ pub mod integrity_pool {
         Ok(())
     }
 
+    pub fn update_reward_program_authority(
+        ctx: Context<UpdateRewardProgramAuthority>,
+        reward_program_authority: Pubkey,
+    ) -> Result<()> {
+        ctx.accounts.pool_config.reward_program_authority = reward_program_authority;
+        Ok(())
+    }
+
     pub fn update_delegation_fee(
         ctx: Context<UpdateDelegationFee>,
         delegation_fee: frac64,