feat: enforce slash custody (#522)

* enforce slash custody * add comments * clippy
pyth-network · Sep 9, 2024 · 8ec2f9d · 8ec2f9d
1 parent 6cdda5e
commit 8ec2f9d
Show file tree

Hide file tree

Showing 10 changed files with 87 additions and 12 deletions.
diff --git a/staking/cli/src/cli.rs b/staking/cli/src/cli.rs
@@ -52,6 +52,8 @@ pub enum Action {
         y:                        u64,
         #[clap(long, help = "Reward program authority parameter")]
         reward_program_authority: Pubkey,
+        #[clap(long, help = "Slash custody parameter")]
+        slash_custody:            Pubkey,
     },
     Advance {
         #[clap(long, help = "Url to the hermes to fetch publisher caps")]

diff --git a/staking/cli/src/instructions.rs b/staking/cli/src/instructions.rs
@@ -351,6 +351,7 @@ pub fn initialize_pool(
     pool_data_keypair: &Keypair,
     reward_program_authority: Pubkey,
     y: u64,
+    slash_custody: Pubkey,
 ) {
     let pool_data_space: u64 = PoolData::LEN.try_into().unwrap();
     let config_address = get_config_address();
@@ -375,10 +376,11 @@ pub fn initialize_pool(
     };
 
     let initialize_pool_accs = integrity_pool::accounts::InitializePool {
-        payer:          payer.pubkey(),
-        pool_data:      pool_data_keypair.pubkey(),
-        pool_config:    pool_config_pubkey,
+        payer: payer.pubkey(),
+        pool_data: pool_data_keypair.pubkey(),
+        pool_config: pool_config_pubkey,
         config_account: config_address,
+        slash_custody,
         system_program: system_program::ID,
     };
 

diff --git a/staking/cli/src/main.rs b/staking/cli/src/main.rs
@@ -30,13 +30,15 @@ fn main() {
             pool_data_keypair,
             reward_program_authority,
             y,
+            slash_custody,
         } => {
             initialize_pool(
                 &rpc_client,
                 &keypair,
                 &pool_data_keypair,
                 reward_program_authority,
                 y,
+                slash_custody,
             );
         }
         Action::Advance {

diff --git a/staking/integration-tests/src/integrity_pool/helper_functions.rs b/staking/integration-tests/src/integrity_pool/helper_functions.rs
@@ -7,8 +7,10 @@ use {
         airdrop_spl,
         initialize_ata,
     },
+    anchor_spl::associated_token::get_associated_token_address,
     integrity_pool::utils::types::FRAC_64_MULTIPLIER,
     solana_sdk::{
+        pubkey::Pubkey,
         signature::Keypair,
         signer::Signer,
     },
@@ -34,3 +36,12 @@ pub fn initialize_pool_reward_custody(
         reward_amount_override.unwrap_or(1_000_000 * FRAC_64_MULTIPLIER),
     );
 }
+
+/// Default slash custody in the tests is the associated token account of the reward program
+/// authority
+pub fn get_default_slash_custody(
+    reward_program_authority: &Pubkey,
+    pyth_token_mint: &Pubkey,
+) -> Pubkey {
+    get_associated_token_address(reward_program_authority, pyth_token_mint)
+}
diff --git a/staking/integration-tests/src/integrity_pool/instructions.rs b/staking/integration-tests/src/integrity_pool/instructions.rs
@@ -1,12 +1,18 @@
 use {
-    super::pda::{
-        get_delegation_record_address,
-        get_pool_config_address,
-        get_pool_reward_custody_address,
-        get_slash_event_address,
+    super::{
+        helper_functions::get_default_slash_custody,
+        pda::{
+            get_delegation_record_address,
+            get_pool_config_address,
+            get_pool_reward_custody_address,
+            get_slash_event_address,
+        },
     },
     crate::{
-        solana::utils::fetch_account_data,
+        solana::{
+            instructions::initialize_ata,
+            utils::fetch_account_data,
+        },
         staking::pda::{
             get_config_address,
             get_stake_account_custody_address,
@@ -42,7 +48,10 @@ use {
         signer::Signer,
         transaction::Transaction,
     },
-    staking::state::stake_account::StakeAccountMetadataV2,
+    staking::state::{
+        global_config::GlobalConfig,
+        stake_account::StakeAccountMetadataV2,
+    },
     std::convert::TryInto,
 };
 
@@ -94,6 +103,8 @@ pub fn create_pool_data_account(
     let pool_data_space: u64 = PoolData::LEN.try_into().unwrap();
     let global_config = get_config_address();
 
+    let global_config_data = fetch_account_data::<GlobalConfig>(svm, &global_config);
+
     let rent = svm.minimum_balance_for_rent_exemption(pool_data_space.try_into().unwrap());
 
     let create_pool_data_acc_ix = create_account(
@@ -111,11 +122,24 @@ pub fn create_pool_data_account(
         y: YIELD,
     };
 
+    // initialize slash custody
+    initialize_ata(
+        svm,
+        payer,
+        global_config_data.pyth_token_mint,
+        reward_program_authority,
+    )
+    .unwrap();
+
     let initialize_pool_accs = integrity_pool::accounts::InitializePool {
         payer:          payer.pubkey(),
         pool_data:      pool_data_keypair.pubkey(),
         pool_config:    pool_config_pubkey,
         config_account: global_config,
+        slash_custody:  get_default_slash_custody(
+            &reward_program_authority,
+            &global_config_data.pyth_token_mint,
+        ),
         system_program: system_program::ID,
     };
 

diff --git a/staking/integration-tests/tests/integrity_pool_slash.rs b/staking/integration-tests/tests/integrity_pool_slash.rs
@@ -3,6 +3,7 @@ use {
     integration_tests::{
         assert_anchor_program_error,
         integrity_pool::{
+            helper_functions::get_default_slash_custody,
             instructions::{
                 advance,
                 advance_delegation_record,
@@ -73,7 +74,13 @@ fn test_create_slash_event() {
     });
     let publisher_index = maybe_publisher_index.unwrap();
 
-    let slash_custody = create_token_account(&mut svm, &payer, &pyth_token_mint.pubkey()).pubkey();
+    let slash_custody = get_default_slash_custody(
+        &reward_program_authority.pubkey(),
+        &pyth_token_mint.pubkey(),
+    );
+    let bad_slash_custody =
+        create_token_account(&mut svm, &payer, &pyth_token_mint.pubkey()).pubkey();
+
     let slashed_publisher = publisher_keypair.pubkey();
 
     assert_anchor_program_error!(
@@ -91,6 +98,21 @@ fn test_create_slash_event() {
         0
     );
 
+    assert_anchor_program_error!(
+        create_slash_event(
+            &mut svm,
+            &payer,
+            &reward_program_authority,
+            0,
+            FRAC_64_MULTIPLIER / 2,
+            bad_slash_custody,
+            slashed_publisher,
+            pool_data_pubkey,
+        ),
+        IntegrityPoolError::InvalidSlashCustodyAccount,
+        0
+    );
+
     create_slash_event(
         &mut svm,
         &payer,
@@ -226,7 +248,10 @@ fn test_slash() {
     });
     let publisher_index = maybe_publisher_index.unwrap();
 
-    let slash_custody = create_token_account(&mut svm, &payer, &pyth_token_mint.pubkey()).pubkey();
+    let slash_custody = get_default_slash_custody(
+        &reward_program_authority.pubkey(),
+        &pyth_token_mint.pubkey(),
+    );
 
     let stake_account_positions =
         initialize_new_stake_account(&mut svm, &payer, &pyth_token_mint, true, true);

diff --git a/staking/programs/integrity-pool/src/context.rs b/staking/programs/integrity-pool/src/context.rs
@@ -42,6 +42,11 @@ pub struct InitializePool<'info> {
     #[account(init, payer = payer, seeds = [POOL_CONFIG.as_bytes()], space = PoolConfig::LEN, bump)]
     pub pool_config: Account<'info, PoolConfig>,
 
+    #[account(
+        token::mint = config_account.pyth_token_mint,
+    )]
+    pub slash_custody: Account<'info, TokenAccount>,
+
     pub system_program: Program<'info, System>,
 }
 
@@ -358,6 +363,7 @@ pub struct CreateSlashEvent<'info> {
         bump,
         has_one = reward_program_authority @ IntegrityPoolError::InvalidRewardProgramAuthority,
         has_one = pool_data @ IntegrityPoolError::InvalidPoolDataAccount,
+        has_one = slash_custody @ IntegrityPoolError::InvalidSlashCustodyAccount,
     )]
     pub pool_config: Account<'info, PoolConfig>,
 

diff --git a/staking/programs/integrity-pool/src/error.rs b/staking/programs/integrity-pool/src/error.rs
@@ -32,4 +32,5 @@ pub enum IntegrityPoolError {
     InvalidPublisher,
     #[msg("Y should not be greater than 1%")]
     InvalidY,
+    InvalidSlashCustodyAccount,
 }
diff --git a/staking/programs/integrity-pool/src/lib.rs b/staking/programs/integrity-pool/src/lib.rs
@@ -44,6 +44,7 @@ pub mod integrity_pool {
         pool_config.reward_program_authority = reward_program_authority;
         pool_config.pyth_token_mint = global_config.pyth_token_mint;
         pool_config.y = y;
+        pool_config.slash_custody = ctx.accounts.slash_custody.key();
 
         let mut pool_data = ctx.accounts.pool_data.load_init()?;
         pool_data.last_updated_epoch = get_current_epoch()? - 1;

diff --git a/staking/programs/integrity-pool/src/state/pool.rs b/staking/programs/integrity-pool/src/state/pool.rs
@@ -495,6 +495,7 @@ pub struct PoolConfig {
     pub reward_program_authority: Pubkey,
     pub pyth_token_mint:          Pubkey,
     pub y:                        frac64,
+    pub slash_custody:            Pubkey,
 }
 
 impl PoolConfig {