post-review fixes

programs/strategy/src/instructions/deposit.rs

@@ -39,14 +39,14 @@ pub fn handle_deposit<'info>(
         return Err(ErrorCode::MaxDepositReached.into());
     }
 
+    strategy.deposit(amount)?;
+    strategy.save_changes(&mut &mut ctx.accounts.strategy.try_borrow_mut_data()?[8..])?;
+
     transfer(
         ctx.accounts.token_program.to_account_info(),
         ctx.accounts.vault_token_account.to_account_info(), 
         ctx.accounts.underlying_token_account.to_account_info(), 
         ctx.accounts.signer.to_account_info(), 
         amount
-    )?;
-
-    strategy.deposit(amount)?;
-    strategy.save_changes(&mut &mut ctx.accounts.strategy.try_borrow_mut_data()?[8..])
+    )
 }

programs/tokenized_vault/src/errors.rs

@@ -62,9 +62,6 @@ pub enum ErrorCode {
     #[msg("Exceed withdraw limit")]
     ExceedWithdrawLimit,
 
-    #[msg("Invalid strategy / token account pairs")]
-    InvalidAccountPairs,
-
     #[msg("Loss is too high")]
     TooMuchLoss,
 

programs/tokenized_vault/src/instructions/add_strategy.rs

@@ -4,11 +4,12 @@ use access_control::{
     program::AccessControl,
    state::{UserRole, Role}
 };
+use strategy::program::Strategy;
 
 use crate::errors::ErrorCode;
 use crate::constants::STRATEGY_DATA_SEED;
 use crate::state::{ StrategyData, Vault};
-use crate::utils::strategy;
+use crate::utils::strategy as strategy_utils;
 
 #[derive(Accounts)]
 pub struct AddStrategy<'info> {
@@ -29,7 +30,7 @@ pub struct AddStrategy<'info> {
     pub vault: AccountLoader<'info, Vault>,
 
     /// CHECK: can be any strategy
-    #[account()]
+    #[account(constraint = *strategy.owner == strategy_program.key())]
     pub strategy: UncheckedAccount<'info>,
 
     #[account(
@@ -47,11 +48,12 @@ pub struct AddStrategy<'info> {
     pub signer: Signer<'info>,
 
     pub access_control: Program<'info, AccessControl>,
-    pub system_program: Program<'info, System>
+    pub system_program: Program<'info, System>,
+    pub strategy_program: Program<'info, Strategy>,
 }
 
 pub fn handle_add_strategy(ctx: Context<AddStrategy>, max_debt: u64) -> Result<()> {
-    let strategy_vault = strategy::get_vault(&ctx.accounts.strategy.to_account_info())?;
+    let strategy_vault = strategy_utils::get_vault(&ctx.accounts.strategy.to_account_info())?;
 
     if strategy_vault != *ctx.accounts.vault.to_account_info().key {
         return Err(ErrorCode::InvalidStrategyToAdd.into());

programs/tokenized_vault/src/instructions/update_debt.rs

@@ -77,14 +77,14 @@ pub fn handle_update_debt<'a, 'b, 'c, 'info>(
 ) -> Result<()> {
     let (total_idle, total_debt, new_debt) = handle_internal(&mut ctx, new_debt)?;
 
-    let vaut_mut = &mut ctx.accounts.vault.load_mut()?;
-    vaut_mut.total_idle = total_idle;
-    vaut_mut.total_debt = total_debt;
+    let vault_mut = &mut ctx.accounts.vault.load_mut()?;
+    vault_mut.total_idle = total_idle;
+    vault_mut.total_debt = total_debt;
 
     ctx.accounts.strategy_data.update_strategy_current_debt(new_debt)?;
 
     emit!(UpdatedCurrentDebtForStrategyEvent {
-        vault_key: vaut_mut.key,
+        vault_key: vault_mut.key,
         strategy_key: ctx.accounts.strategy.key(),
         total_idle: total_idle,
         total_debt: total_debt,

programs/tokenized_vault/src/instructions/withdraw.rs

@@ -103,6 +103,8 @@ pub fn handle_withdraw<'info>(
         }
     }
 
+    ctx.accounts.vault.load_mut()?.handle_withdraw(assets_to_transfer, shares_to_burn);
+
     token::burn(
         ctx.accounts.token_program.to_account_info(),
         ctx.accounts.shares_mint.to_account_info(),
@@ -120,8 +122,7 @@ pub fn handle_withdraw<'info>(
         &ctx.accounts.vault.load()?.seeds()
     )?;
 
-    let mut vault = ctx.accounts.vault.load_mut()?;
-    vault.handle_withdraw(assets_to_transfer, shares_to_burn);
+    let vault = ctx.accounts.vault.load()?;
 
     emit!(VaultWithdrawlEvent {
         vault_key: vault.key,