validate token vault

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

@@ -10,35 +10,44 @@ pub mod seeds {
     pub const TOKEN_VAULT_PREFIX: &[u8] = b"token_vault";
 }
 
-pub static WHITELISTED_ACTIONS: [(Pubkey, &[u8]); 7] = [
+// (program_id, instruction, index_of_token_vault_account)
+// TODO should find a way to avoid hardcoding index of token_vault_account
+pub static WHITELISTED_ACTIONS: [(Pubkey, &[u8], usize); 7] = [
     // damm v2
     (
         damm_v2::ID,
         damm_v2::client::args::ClaimPositionFee::DISCRIMINATOR,
+        4,
     ),
     (
         damm_v2::ID,
         damm_v2::client::args::ClaimReward::DISCRIMINATOR,
+        5,
     ),
     // DBC
     (
         dynamic_bonding_curve::ID,
         dynamic_bonding_curve::client::args::CreatorWithdrawSurplus::DISCRIMINATOR,
+        3,
     ),
     (
         dynamic_bonding_curve::ID,
         dynamic_bonding_curve::client::args::ClaimCreatorTradingFee::DISCRIMINATOR,
+        3,
     ),
     (
         dynamic_bonding_curve::ID,
         dynamic_bonding_curve::client::args::PartnerWithdrawSurplus::DISCRIMINATOR,
+        3,
     ),
     (
         dynamic_bonding_curve::ID,
         dynamic_bonding_curve::client::args::ClaimTradingFee::DISCRIMINATOR,
+        4,
     ),
     (
         dynamic_bonding_curve::ID,
         dynamic_bonding_curve::client::args::WithdrawMigrationFee::DISCRIMINATOR,
+        3,
     ),
 ];

programs/dynamic-fee-sharing/src/instructions/ix_fund_by_claiming_fee.rs

@@ -22,10 +22,17 @@ pub struct FundByClaimingFeeCtx<'info> {
     pub source_program: UncheckedAccount<'info>,
 }
 
-pub fn is_support_action(source_program: &Pubkey, discriminator: &[u8]) -> bool {
-    for &(program, disc) in WHITELISTED_ACTIONS.iter() {
+pub fn is_support_action<'info>(
+    source_program: &Pubkey,
+    discriminator: &[u8],
+    token_vault: Pubkey,
+    remaining_accounts: &[AccountInfo<'info>],
+) -> bool {
+    for &(program, disc, token_vault_index) in WHITELISTED_ACTIONS.iter() {
         if program.eq(source_program) && disc.eq(discriminator) {
-            return true;
+            if let Some(token_vault_account) = remaining_accounts.get(token_vault_index) {
+                return token_vault.eq(token_vault_account.key);
+            }
         }
     }
     false
@@ -37,7 +44,12 @@ pub fn handle_fund_by_claiming_fee(
 ) -> Result<()> {
     let discriminator = &payload[..8]; // first 8 bytes is discriminator
     require!(
-        is_support_action(ctx.accounts.source_program.key, discriminator),
+        is_support_action(
+            ctx.accounts.source_program.key,
+            discriminator,
+            ctx.accounts.token_vault.key(),
+            ctx.remaining_accounts
+        ),
         FeeVaultError::InvalidAction
     );