feat: address comments

Author: bangyro

programs/dynamic-fee-sharing/src/error.rs

@@ -43,4 +43,7 @@ pub enum FeeVaultError {
 
     #[msg("Type cast error")]
     TypeCastFailed,
+
+    #[msg("Invalid account discriminator")]
+    InvalidDiscriminator,
 }

programs/dynamic-fee-sharing/src/instructions/operator/ix_remove_user.rs

@@ -2,7 +2,7 @@ use crate::constants::seeds::USER_UNCLAIMED_FEE_PREFIX;
 use crate::event::EvtRemoveUser;
 use crate::math::SafeMath;
 use crate::state::{FeeVault, UserUnclaimedFee};
-use crate::utils::account::create_pda_account_with_anchor_discriminator;
+use crate::utils::account::{create_pda_account_with_anchor_discriminator, load_account_data_mut};
 use anchor_lang::prelude::*;
 
 #[event_cpi]
@@ -59,7 +59,7 @@ pub fn handle_remove_user(ctx: Context<RemoveUserCtx>, index: u8) -> Result<()>
 
         let mut data = user_unclaimed_fee.try_borrow_mut_data()?;
 
-        let user_unclaimed_fee = bytemuck::from_bytes_mut::<UserUnclaimedFee>(&mut data[8..]);
+        let user_unclaimed_fee = load_account_data_mut::<UserUnclaimedFee>(&mut data)?;
         user_unclaimed_fee.unclaimed_fee =
             user_unclaimed_fee.unclaimed_fee.safe_add(unclaimed_fee)?;
     }

programs/dynamic-fee-sharing/src/state/fee_vault.rs

@@ -220,11 +220,12 @@ impl FeeVault {
             .iter()
             .filter(|u| u.address.ne(&Pubkey::default()))
             .count();
+        // user_count include the user being removed. after removal user count should be at least MIN_USER
         require!(user_count > MIN_USER, FeeVaultError::InvalidNumberOfUsers);
 
-        let unclaimed_fee = self.users[index].get_total_pending_fee(self.fee_per_share)?;
+        let unclaimed_fee = user.get_total_pending_fee(self.fee_per_share)?;
 
-        self.total_share = self.total_share.safe_sub(self.users[index].share)?;
+        self.total_share = self.total_share.safe_sub(user.share)?;
 
         // shift users to the left
         for i in index..MAX_USER - 1 {

programs/dynamic-fee-sharing/src/utils/account.rs

@@ -1,3 +1,4 @@
+use crate::error::FeeVaultError;
 use anchor_lang::prelude::*;
 use anchor_lang::solana_program::program::{invoke, invoke_signed};
 use anchor_lang::solana_program::system_instruction;
@@ -86,3 +87,12 @@ pub fn create_pda_account_with_anchor_discriminator<'a, T: Discriminator + Space
 
     Ok(())
 }
+
+/// Validates the Anchor discriminator and returns a mutable bytemuck reference.
+pub fn load_account_data_mut<T: Discriminator + bytemuck::Pod>(
+    data: &mut [u8],
+) -> Result<&mut T> {
+    let (disc, rest) = data.split_at_mut(T::DISCRIMINATOR.len());
+    require!(disc == T::DISCRIMINATOR, FeeVaultError::InvalidDiscriminator);
+    Ok(bytemuck::from_bytes_mut::<T>(rest))
+}