From 6cdda5e3b50f2498e75610d36e2e157efee20b5c Mon Sep 17 00:00:00 2001
From: guibescos <59208140+guibescos@users.noreply.github.com>
Date: Mon, 9 Sep 2024 14:02:59 +0100
Subject: [PATCH] feat: allow reward authority to sign for first stake account
 declaration (#521)

* do it

* do it

* go
---
 .../tests/set_publisher_stake_account.rs      | 58 ++++++++++++++++++-
 staking/programs/integrity-pool/src/error.rs  |  2 +-
 staking/programs/integrity-pool/src/lib.rs    |  9 +--
 3 files changed, 63 insertions(+), 6 deletions(-)

diff --git a/staking/integration-tests/tests/set_publisher_stake_account.rs b/staking/integration-tests/tests/set_publisher_stake_account.rs
index 35529e70..76baefe6 100644
--- a/staking/integration-tests/tests/set_publisher_stake_account.rs
+++ b/staking/integration-tests/tests/set_publisher_stake_account.rs
@@ -71,7 +71,7 @@ fn test_set_publisher_stake_account() {
             None,
             stake_account_positions,
         ),
-        IntegrityPoolError::PublisherNeedsToSign,
+        IntegrityPoolError::PublisherOrRewardAuthorityNeedsToSign,
         0
     );
 
@@ -371,3 +371,59 @@ fn test_set_publisher_stake_account() {
         assert_eq!(pool_data.self_del_state[i], DelegationState::default());
     }
 }
+
+
+#[test]
+fn test_set_publisher_stake_account_reward_authority() {
+    let SetupResult {
+        mut svm,
+        payer,
+        pyth_token_mint,
+        publisher_keypair,
+        pool_data_pubkey,
+        reward_program_authority,
+        maybe_publisher_index,
+    } = setup(SetupProps {
+        init_config:            true,
+        init_target:            true,
+        init_mint:              true,
+        init_pool_data:         true,
+        init_publishers:        true,
+        reward_amount_override: None,
+    });
+    let publisher_index = maybe_publisher_index.unwrap();
+
+    let stake_account_positions =
+        initialize_new_stake_account(&mut svm, &payer, &pyth_token_mint, true, true);
+
+    set_publisher_stake_account(
+        &mut svm,
+        &payer,
+        &reward_program_authority,
+        publisher_keypair.pubkey(),
+        None,
+        stake_account_positions,
+    )
+    .unwrap();
+
+    let pool_data: PoolData = fetch_account_data_bytemuck(&mut svm, &pool_data_pubkey);
+
+    assert_eq!(
+        pool_data.publisher_stake_accounts[publisher_index],
+        stake_account_positions,
+    );
+
+    // can't set it once it's set
+    assert_anchor_program_error!(
+        set_publisher_stake_account(
+            &mut svm,
+            &payer,
+            &reward_program_authority,
+            publisher_keypair.pubkey(),
+            Some(stake_account_positions),
+            stake_account_positions,
+        ),
+        IntegrityPoolError::StakeAccountOwnerNeedsToSign,
+        0
+    );
+}
diff --git a/staking/programs/integrity-pool/src/error.rs b/staking/programs/integrity-pool/src/error.rs
index 6051075c..90c8a17f 100644
--- a/staking/programs/integrity-pool/src/error.rs
+++ b/staking/programs/integrity-pool/src/error.rs
@@ -3,7 +3,7 @@ use anchor_lang::error_code;
 #[error_code]
 pub enum IntegrityPoolError {
     PublisherNotFound,
-    PublisherNeedsToSign,
+    PublisherOrRewardAuthorityNeedsToSign,
     StakeAccountOwnerNeedsToSign,
     OutdatedPublisherAccounting,
     TooManyPublishers,
diff --git a/staking/programs/integrity-pool/src/lib.rs b/staking/programs/integrity-pool/src/lib.rs
index a3764d28..356b647b 100644
--- a/staking/programs/integrity-pool/src/lib.rs
+++ b/staking/programs/integrity-pool/src/lib.rs
@@ -226,6 +226,7 @@ pub mod integrity_pool {
         let signer = &ctx.accounts.signer;
         let publisher = &ctx.accounts.publisher;
         let pool_data = &mut ctx.accounts.pool_data.load_mut()?;
+        let pool_config = &ctx.accounts.pool_config;
         let new_stake_account =
             DynamicPositionArray::load(&ctx.accounts.new_stake_account_positions)?;
 
@@ -236,10 +237,10 @@ pub mod integrity_pool {
         let publisher_index = pool_data.get_publisher_index(publisher.key)?;
 
         if pool_data.publisher_stake_accounts[publisher_index] == Pubkey::default() {
-            require_eq!(
-                signer.key(),
-                publisher.key(),
-                IntegrityPoolError::PublisherNeedsToSign
+            require!(
+                signer.key() == publisher.key()
+                    || signer.key() == pool_config.reward_program_authority,
+                IntegrityPoolError::PublisherOrRewardAuthorityNeedsToSign
             );
         } else if let Some(current_stake_account_positions) =
             &ctx.accounts.current_stake_account_positions_option