program: cleanup some vault protocol things (#59)

* rm extra msg * add missing validate to protocol_request_withdraw * make sure manager shares dont underflow * another invariant * add another invariant * fix invariant check * rm incorrect invariant check --------- Co-authored-by: wphan <[email protected]>
drift-labs · Oct 29, 2024 · 0f3f82b · 0f3f82b
1 parent dc0a85f
commit 0f3f82b
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 4 deletions.
diff --git a/programs/drift_vaults/src/state/vault.rs b/programs/drift_vaults/src/state/vault.rs
@@ -310,6 +310,15 @@ impl Vault {
             self.last_fee_update_ts = now;
         }
 
+        validate!(
+            self.total_shares >= self.user_shares,
+            ErrorCode::InvalidVaultSharesDetected,
+            "total_shares must be >= user_shares"
+        )?;
+
+        // this will underflow if there is an issue with protocol fee calc
+        self.get_manager_shares(vault_protocol)?;
+
         Ok(VaultFee {
             management_fee_payment: management_fee_payment.cast::<i64>()?,
             management_fee_shares: management_fee_shares.cast::<i64>()?,
@@ -801,6 +810,13 @@ impl Vault {
             ErrorCode::InvalidVaultWithdrawSize,
             "Requested n_shares = 0"
         )?;
+        validate!(
+            protocol_shares_before >= n_shares,
+            ErrorCode::InvalidVaultWithdrawSize,
+            "Requested n_shares={} > protocol shares={}",
+            n_shares,
+            protocol_shares_before,
+        )?;
 
         let total_vault_shares_before = self.total_shares;
         let user_vault_shares_before = self.user_shares;

diff --git a/programs/drift_vaults/src/state/vault_depositor.rs b/programs/drift_vaults/src/state/vault_depositor.rs
@@ -362,8 +362,6 @@ impl VaultDepositor {
         } = vault.apply_fee(vault_protocol, vault_equity, now)?;
         let (manager_profit_share, protocol_profit_share) =
             self.apply_profit_share(vault_equity, vault, vault_protocol)?;
-        msg!("manager_profit_share: {}", manager_profit_share);
-        msg!("protocol_profit_share: {}", protocol_profit_share);
 
         let (withdraw_value, n_shares) = withdraw_unit.get_withdraw_value_and_shares(
             withdraw_amount,
@@ -772,8 +770,6 @@ impl VaultDepositor {
         )?;
         let withdraw_amount = self.last_withdraw_request.value.min(shares_value);
 
-        msg!("withdraw amount: {}", withdraw_amount);
-
         let mut spot_market = spot_market_map.get_ref_mut(&vault.spot_market_index)?;
 
         // Save relevant data before updating balances