feat: Improvements from audit feedback

src/consts.rs

@@ -1,4 +1,3 @@
-use solana_program::pubkey;
 use solana_program::pubkey::Pubkey;
 
 /// The delegation session fees (extracted in percentage from the delegation PDAs rent on closure).
@@ -11,10 +10,4 @@ pub const PROTOCOL_FEES_PERCENTAGE: u8 = 10;
 pub const EXTERNAL_UNDELEGATE_DISCRIMINATOR: [u8; 8] = [196, 28, 41, 206, 48, 37, 51, 167];
 
 /// The program ID of the delegation program.
-pub const DELEGATION_PROGRAM_ID: Pubkey = crate::id();
-
-/// The admin pubkey of the authority allowed to whitelist validators.
-#[cfg(feature = "unit_test_config")]
-pub const ADMIN_PUBKEY: Pubkey = pubkey!("tEsT3eV6RFCWs1BZ7AXTzasHqTtMnMLCB2tjQ42TDXD");
-#[cfg(not(feature = "unit_test_config"))]
-pub const ADMIN_PUBKEY: Pubkey = pubkey!("3FwNxjbCqdD7G6MkrAdwTd5Zf6R3tHoapam4Pv1X2KBB");
+pub const DELEGATION_PROGRAM_ID: Pubkey = crate::id();

src/instruction_builder/close_validator_fees_vault.rs

@@ -1,5 +1,5 @@
 use solana_program::instruction::Instruction;
-use solana_program::{instruction::AccountMeta, pubkey::Pubkey};
+use solana_program::{bpf_loader_upgradeable, instruction::AccountMeta, pubkey::Pubkey};
 
 use crate::discriminator::DlpDiscriminator;
 use crate::pda::validator_fees_vault_pda_from_validator;
@@ -12,11 +12,14 @@ pub fn close_validator_fees_vault(
     validator_identity: Pubkey,
 ) -> Instruction {
     let validator_fees_vault_pda = validator_fees_vault_pda_from_validator(&validator_identity);
+    let delegation_program_data =
+        Pubkey::find_program_address(&[crate::ID.as_ref()], &bpf_loader_upgradeable::id()).0;
     Instruction {
         program_id: crate::id(),
         accounts: vec![
             AccountMeta::new(payer, true),
             AccountMeta::new(admin, true),
+            AccountMeta::new_readonly(delegation_program_data, false),
             AccountMeta::new(validator_identity, false),
             AccountMeta::new(validator_fees_vault_pda, false),
         ],

src/instruction_builder/init_validator_fees_vault.rs

@@ -1,5 +1,5 @@
 use solana_program::instruction::Instruction;
-use solana_program::system_program;
+use solana_program::{bpf_loader_upgradeable, system_program};
 use solana_program::{instruction::AccountMeta, pubkey::Pubkey};
 
 use crate::discriminator::DlpDiscriminator;
@@ -13,11 +13,14 @@ pub fn init_validator_fees_vault(
     validator_identity: Pubkey,
 ) -> Instruction {
     let validator_fees_vault_pda = validator_fees_vault_pda_from_validator(&validator_identity);
+    let delegation_program_data =
+        Pubkey::find_program_address(&[crate::ID.as_ref()], &bpf_loader_upgradeable::id()).0;
     Instruction {
         program_id: crate::id(),
         accounts: vec![
             AccountMeta::new(payer, true),
             AccountMeta::new(admin, true),
+            AccountMeta::new_readonly(delegation_program_data, false),
             AccountMeta::new(validator_identity, false),
             AccountMeta::new(validator_fees_vault_pda, false),
             AccountMeta::new_readonly(system_program::id(), false),

src/instruction_builder/protocol_claim_fees.rs

@@ -1,5 +1,5 @@
 use solana_program::instruction::Instruction;
-use solana_program::{instruction::AccountMeta, pubkey::Pubkey};
+use solana_program::{bpf_loader_upgradeable, instruction::AccountMeta, pubkey::Pubkey};
 
 use crate::discriminator::DlpDiscriminator;
 use crate::pda::fees_vault_pda;
@@ -8,11 +8,14 @@ use crate::pda::fees_vault_pda;
 /// See [crate::processor::process_protocol_claim_fees] for docs.
 pub fn protocol_claim_fees(admin: Pubkey) -> Instruction {
     let fees_vault_pda = fees_vault_pda();
+    let delegation_program_data =
+        Pubkey::find_program_address(&[crate::ID.as_ref()], &bpf_loader_upgradeable::id()).0;
     Instruction {
         program_id: crate::id(),
         accounts: vec![
             AccountMeta::new(admin, true),
             AccountMeta::new(fees_vault_pda, false),
+            AccountMeta::new_readonly(delegation_program_data, false),
         ],
         data: DlpDiscriminator::ProtocolClaimFees.to_vec(),
     }

src/instruction_builder/whitelist_validator_for_program.rs

@@ -20,6 +20,8 @@ pub fn whitelist_validator_for_program(
     let args = WhitelistValidatorForProgramArgs { insert };
     let program_data =
         Pubkey::find_program_address(&[program.as_ref()], &bpf_loader_upgradeable::id()).0;
+    let delegation_program_data =
+        Pubkey::find_program_address(&[crate::ID.as_ref()], &bpf_loader_upgradeable::id()).0;
     let program_config_pda = program_config_from_program_id(&program);
     Instruction {
         program_id: crate::id(),
@@ -28,6 +30,7 @@ pub fn whitelist_validator_for_program(
             AccountMeta::new_readonly(validator_identity, false),
             AccountMeta::new_readonly(program, false),
             AccountMeta::new_readonly(program_data, false),
+            AccountMeta::new_readonly(delegation_program_data, false),
             AccountMeta::new(program_config_pda, false),
             AccountMeta::new_readonly(system_program::id(), false),
         ],

src/processor/close_validator_fees_vault.rs

@@ -2,9 +2,10 @@ use solana_program::msg;
 use solana_program::program_error::ProgramError;
 use solana_program::{account_info::AccountInfo, entrypoint::ProgramResult, pubkey::Pubkey};
 
-use crate::consts::ADMIN_PUBKEY;
 use crate::error::DlpError::Unauthorized;
-use crate::processor::utils::loaders::{load_initialized_pda, load_signer};
+use crate::processor::utils::loaders::{
+    load_initialized_pda, load_program_upgrade_authority, load_signer,
+};
 use crate::processor::utils::pda::close_pda;
 use crate::validator_fees_vault_seeds_from_validator;
 
@@ -30,7 +31,9 @@ pub fn process_close_validator_fees_vault(
     _data: &[u8],
 ) -> ProgramResult {
     // Load Accounts
-    let [payer, admin, validator_identity, validator_fees_vault] = accounts else {
+    let [payer, admin, delegation_program_data, validator_identity, validator_fees_vault] =
+        accounts
+    else {
         return Err(ProgramError::NotEnoughAccountKeys);
     };
 
@@ -39,10 +42,12 @@ pub fn process_close_validator_fees_vault(
     load_signer(admin, "admin")?;
 
     // Check if the admin is the correct one
-    if !admin.key.eq(&ADMIN_PUBKEY) {
+    let admin_pubkey =
+        load_program_upgrade_authority(&crate::ID, delegation_program_data)?.ok_or(Unauthorized)?;
+    if !admin.key.eq(&admin_pubkey) {
         msg!(
             "Expected admin pubkey: {} but got {}",
-            ADMIN_PUBKEY,
+            admin_pubkey,
             admin.key
         );
         return Err(Unauthorized.into());

src/processor/commit_state.rs

@@ -130,14 +130,14 @@ pub(crate) fn process_commit_state_internal(
     // If the slot is less, we simply do not commit.
     // Since commit instructions are typically bundled, we return without error
     // so that correct commits are executed.
-    if args.commit_record_slot < delegation_metadata.last_update_external_slot {
-        msg!(
-            "Slot {} is outdated, previous slot is {}. Skipping commit",
-            args.commit_record_slot,
-            delegation_metadata.last_update_external_slot
-        );
-        return Ok(());
-    }
+    // if args.commit_record_slot < delegation_metadata.last_update_external_slot {
+    //     msg!(
+    //         "Slot {} is outdated, previous slot is {}. Skipping commit",
+    //         args.commit_record_slot,
+    //         delegation_metadata.last_update_external_slot
+    //     );
+    //     return Ok(());
+    // }
 
     // Once the account is marked as undelegatable, any subsequent commit should fail
     if delegation_metadata.is_undelegatable {

src/processor/init_validator_fees_vault.rs

@@ -4,9 +4,10 @@ use solana_program::{
     account_info::AccountInfo, entrypoint::ProgramResult, pubkey::Pubkey, system_program,
 };
 
-use crate::consts::ADMIN_PUBKEY;
 use crate::error::DlpError::Unauthorized;
-use crate::processor::utils::loaders::{load_program, load_signer, load_uninitialized_pda};
+use crate::processor::utils::loaders::{
+    load_program, load_program_upgrade_authority, load_signer, load_uninitialized_pda,
+};
 use crate::processor::utils::pda::create_pda;
 use crate::validator_fees_vault_seeds_from_validator;
 
@@ -35,7 +36,9 @@ pub fn process_init_validator_fees_vault(
     _data: &[u8],
 ) -> ProgramResult {
     // Load Accounts
-    let [payer, admin, validator_identity, validator_fees_vault, system_program] = accounts else {
+    let [payer, admin, delegation_program_data, validator_identity, validator_fees_vault, system_program] =
+        accounts
+    else {
         return Err(ProgramError::NotEnoughAccountKeys);
     };
 
@@ -45,10 +48,12 @@ pub fn process_init_validator_fees_vault(
     load_program(system_program, system_program::id(), "system program")?;
 
     // Check if the admin is the correct one
-    if !admin.key.eq(&ADMIN_PUBKEY) {
+    let admin_pubkey =
+        load_program_upgrade_authority(&crate::ID, delegation_program_data)?.ok_or(Unauthorized)?;
+    if !admin.key.eq(&admin_pubkey) {
         msg!(
             "Expected admin pubkey: {} but got {}",
-            ADMIN_PUBKEY,
+            admin_pubkey,
             admin.key
         );
         return Err(Unauthorized.into());

src/processor/protocol_claim_fees.rs

@@ -1,6 +1,7 @@
-use crate::consts::ADMIN_PUBKEY;
 use crate::error::DlpError::Unauthorized;
-use crate::processor::utils::loaders::{load_initialized_protocol_fees_vault, load_signer};
+use crate::processor::utils::loaders::{
+    load_initialized_protocol_fees_vault, load_program_upgrade_authority, load_signer,
+};
 use solana_program::msg;
 use solana_program::program_error::ProgramError;
 use solana_program::rent::Rent;
@@ -27,7 +28,7 @@ pub fn process_protocol_claim_fees(
     _data: &[u8],
 ) -> ProgramResult {
     // Load Accounts
-    let [admin, fees_vault] = accounts else {
+    let [admin, fees_vault, delegation_program_data] = accounts else {
         return Err(ProgramError::NotEnoughAccountKeys);
     };
 
@@ -36,10 +37,12 @@ pub fn process_protocol_claim_fees(
     load_initialized_protocol_fees_vault(fees_vault, true)?;
 
     // Check if the admin is the correct one
-    if !admin.key.eq(&ADMIN_PUBKEY) {
+    let admin_pubkey =
+        load_program_upgrade_authority(&crate::ID, delegation_program_data)?.ok_or(Unauthorized)?;
+    if !admin.key.eq(&admin_pubkey) {
         msg!(
             "Expected admin pubkey: {} but got {}",
-            ADMIN_PUBKEY,
+            admin_pubkey,
             admin.key
         );
         return Err(Unauthorized.into());

src/processor/utils/loaders.rs

@@ -6,9 +6,10 @@ use crate::{
     delegation_record_seeds_from_delegated_account, fees_vault_seeds,
     program_config_seeds_from_program_id, validator_fees_vault_seeds_from_validator,
 };
+use solana_program::bpf_loader_upgradeable::UpgradeableLoaderState;
 use solana_program::{
-    account_info::AccountInfo, msg, program_error::ProgramError, pubkey::Pubkey, system_program,
-    sysvar,
+    account_info::AccountInfo, bpf_loader_upgradeable, msg, program_error::ProgramError,
+    pubkey::Pubkey, system_program, sysvar,
 };
 
 /// Errors if:
@@ -193,6 +194,46 @@ pub fn load_program(info: &AccountInfo, key: Pubkey, label: &str) -> Result<(),
     Ok(())
 }
 
+/// Get the program upgrade authority for a given program
+pub fn load_program_upgrade_authority(
+    program: &Pubkey,
+    program_data: &AccountInfo,
+) -> Result<Option<Pubkey>, ProgramError> {
+    let program_data_address =
+        Pubkey::find_program_address(&[program.as_ref()], &bpf_loader_upgradeable::id()).0;
+
+    /// During tests, the upgrade authority is a test pubkey
+    #[cfg(feature = "unit_test_config")]
+    if program.eq(&crate::ID) {
+        return Ok(Some(solana_program::pubkey!(
+            "tEsT3eV6RFCWs1BZ7AXTzasHqTtMnMLCB2tjQ42TDXD"
+        )));
+    }
+
+    if !program_data_address.eq(program_data.key) {
+        msg!(
+            "Expected program data address to be {}, but got {}",
+            program_data_address,
+            program_data.key
+        );
+        return Err(ProgramError::InvalidAccountData);
+    }
+
+    let program_account_data = program_data.try_borrow_data()?;
+    if let UpgradeableLoaderState::ProgramData {
+        upgrade_authority_address,
+        ..
+    } = bincode::deserialize(&program_account_data).map_err(|_| {
+        msg!("Unable to deserialize ProgramData {}", program);
+        ProgramError::InvalidAccountData
+    })? {
+        Ok(upgrade_authority_address)
+    } else {
+        msg!("Expected program account {} to hold ProgramData", program);
+        Err(ProgramError::InvalidAccountData)
+    }
+}
+
 /// Load fee vault PDA
 /// - Protocol fees vault PDA
 pub fn load_initialized_protocol_fees_vault(

src/processor/whitelist_validator_for_program.rs

@@ -1,17 +1,16 @@
 use crate::args::WhitelistValidatorForProgramArgs;
-use crate::consts::ADMIN_PUBKEY;
 use crate::error::DlpError::Unauthorized;
-use crate::processor::utils::loaders::{load_pda, load_program, load_signer};
+use crate::processor::utils::loaders::{
+    load_pda, load_program, load_program_upgrade_authority, load_signer,
+};
 use crate::processor::utils::pda::{create_pda, resize_pda};
 use crate::program_config_seeds_from_program_id;
 use crate::state::ProgramConfig;
 use borsh::BorshDeserialize;
-use solana_program::bpf_loader_upgradeable::UpgradeableLoaderState;
 use solana_program::msg;
 use solana_program::program_error::ProgramError;
 use solana_program::{
-    account_info::AccountInfo, bpf_loader_upgradeable, entrypoint::ProgramResult, pubkey::Pubkey,
-    system_program,
+    account_info::AccountInfo, entrypoint::ProgramResult, pubkey::Pubkey, system_program,
 };
 
 /// Whitelist a validator for a program
@@ -44,14 +43,14 @@ pub fn process_whitelist_validator_for_program(
     let args = WhitelistValidatorForProgramArgs::try_from_slice(data)?;
 
     // Load Accounts
-    let [authority, validator_identity, program, program_data, program_config_account, system_program] =
+    let [authority, validator_identity, program, program_data, delegation_program_data, program_config_account, system_program] =
         accounts
     else {
         return Err(ProgramError::NotEnoughAccountKeys);
     };
 
     load_signer(authority, "authority")?;
-    validate_authority(authority, program, program_data)?;
+    validate_authority(authority, program, program_data, delegation_program_data)?;
     load_program(system_program, system_program::id(), "system program")?;
 
     let program_config_bump = load_pda(
@@ -104,53 +103,22 @@ fn validate_authority(
     authority: &AccountInfo,
     program: &AccountInfo,
     program_data: &AccountInfo,
+    delegation_program_data: &AccountInfo,
 ) -> Result<(), ProgramError> {
-    if authority.key.eq(&ADMIN_PUBKEY)
+    let admin_pubkey =
+        load_program_upgrade_authority(&crate::ID, delegation_program_data)?.ok_or(Unauthorized)?;
+    if authority.key.eq(&admin_pubkey)
         || authority
             .key
-            .eq(&get_program_upgrade_authority(program, program_data)?.ok_or(Unauthorized)?)
+            .eq(&load_program_upgrade_authority(program.key, program_data)?.ok_or(Unauthorized)?)
     {
         Ok(())
     } else {
         msg!(
             "Expected authority to be {} or program upgrade authority, but got {}",
-            ADMIN_PUBKEY,
+            admin_pubkey,
             authority.key
         );
         Err(Unauthorized.into())
     }
 }
-
-/// Get the program upgrade authority for a given program
-fn get_program_upgrade_authority(
-    program: &AccountInfo,
-    program_data: &AccountInfo,
-) -> Result<Option<Pubkey>, ProgramError> {
-    let program_data_address =
-        Pubkey::find_program_address(&[program.key.as_ref()], &bpf_loader_upgradeable::id()).0;
-
-    if !program_data_address.eq(program_data.key) {
-        msg!(
-            "Expected program data address to be {}, but got {}",
-            program_data_address,
-            program_data.key
-        );
-        return Err(ProgramError::InvalidAccountData);
-    }
-
-    let program_account_data = program_data.try_borrow_data()?;
-    if let UpgradeableLoaderState::ProgramData {
-        upgrade_authority_address,
-        ..
-    } =
-        bincode::deserialize(&program_account_data).map_err(|_| ProgramError::InvalidAccountData)?
-    {
-        Ok(upgrade_authority_address)
-    } else {
-        msg!(
-            "Expected program account {} to hold ProgramData",
-            program.key
-        );
-        Err(ProgramError::InvalidAccountData)
-    }
-}