feat: allow reward authority to sign for first stake account declarat…

…ion (#521)

* do it

* do it

* go

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
+    );
+}

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,

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