MP-2543. Deposit cap improvement (#425)

* Deposit cap check ommited if amount after TX <= before TX for Swap and ProvideLiquidity.

* Update test cases.

* Update schema.

* Fix typo.

* Improve comments in tests.

* Add test to cover missing line coverage.

* Review comments.

* Update comment.

contracts/credit-manager/src/deposit.rs

@@ -1,6 +1,6 @@
-use std::collections::BTreeSet;
+use std::collections::BTreeMap;
 
-use cosmwasm_std::{Coin, Coins, Deps, DepsMut, Response};
+use cosmwasm_std::{Coin, Coins, Deps, DepsMut, Response, Uint128};
 use mars_types::params::TotalDepositResponse;
 
 use crate::{
@@ -48,12 +48,15 @@ fn assert_sent_fund(expected: &Coin, received_coins: &Coins) -> ContractResult<(
 /// Given a list of denoms, assert that the total deposited amount of each denom
 /// across Red Bank and Rover does not exceed its deposit cap recorded in the
 /// params contract.
-pub fn assert_deposit_caps(deps: Deps, denoms: BTreeSet<String>) -> ContractResult<Response> {
+pub fn assert_deposit_caps(
+    deps: Deps,
+    denom_deposits: BTreeMap<String, Option<Uint128>>,
+) -> ContractResult<Response> {
     let params = PARAMS.load(deps.storage)?;
 
     let mut response = Response::new().add_attribute("action", "callback/assert_deposit_caps");
 
-    for denom in denoms {
+    for (denom, deposited_opt) in denom_deposits {
         // Asset is not found (not whitelisted) and it doesn't count towards the cap and Health Factor
         let params_opt = params.query_asset_params(&deps.querier, &denom)?;
         if params_opt.is_none() {
@@ -66,6 +69,16 @@ pub fn assert_deposit_caps(deps: Deps, denoms: BTreeSet<String>) -> ContractResu
             cap,
         } = params.query_total_deposit(&deps.querier, &denom)?;
 
+        // - If there is None in the map, it means that the deposit cap should be enforced. It is related to the Deposit action.
+        // - If there is Some in the map, it means that the deposit amount should be compared (value before and after the TX).
+        // It is related to the SwapExactIn and ProvideLiquidity actions.
+        if let Some(deposited) = deposited_opt {
+            // amount is lower than or equal to the previous deposit amount so it is fine
+            if amount <= deposited {
+                continue;
+            }
+        }
+
         if amount > cap {
             return Err(ContractError::AboveAssetDepositCap {
                 new_value: Coin {
@@ -84,3 +97,61 @@ pub fn assert_deposit_caps(deps: Deps, denoms: BTreeSet<String>) -> ContractResu
 
     Ok(response)
 }
+
+/// Update the total deposit amount for the asset in the denom_deposits map
+/// The function either resets the deposit amount to None for Deposit actions
+/// or updates the deposit amount based on the received coins and existing parameters.
+pub fn update_or_reset_denom_deposits(
+    deps: Deps,
+    denom_deposits: &mut BTreeMap<String, Option<Uint128>>,
+    denom: &str,
+    received_coins: &Coins,
+    deposit_action: bool,
+) -> ContractResult<()> {
+    // Overwrite the previous amount for Deposit action.
+    // It means that we don't compare deposits before and after the TX.
+    //
+    // Strictly enforce the deposit cap to prevent any increase in total value.
+    // Even if we use the funds for operations within the account (such as liquidation),
+    // and withdraw them at the end (resulting in zero net inflow), temporary funds
+    // could still be used for malicious actions.
+    if deposit_action {
+        denom_deposits.insert(denom.to_string(), None);
+
+        return Ok(());
+    }
+
+    // Check if the denomination is already in the list.
+    // This ensures that a Deposit action (which does not have an associated amount)
+    // is not overwritten by a subsequent Swap or ProvideLiquidity action.
+    // By confirming the existence of the denomination in the list, we maintain
+    // the integrity of the original Deposit action.
+    if denom_deposits.contains_key(denom) {
+        return Ok(());
+    }
+
+    // Load the params
+    let params = PARAMS.load(deps.storage)?;
+
+    // Asset is not found (not whitelisted) and it doesn't count towards the cap and Health Factor
+    let params_opt = params.query_asset_params(&deps.querier, denom)?;
+    if params_opt.is_none() {
+        return Ok(());
+    }
+
+    // Check total deposit amount for the asset
+    let total_deposit_amt = params.query_total_deposit(&deps.querier, denom)?.amount;
+
+    // Check if the asset was sent in the TX
+    let received_amt = received_coins.amount_of(denom);
+
+    // Total deposit amount is the sum of all deposits for the asset across Red Bank and Rover.
+    // If the asset was sent in the TX, the Credit Manager balance already includes the deposited amount
+    // so we need to subtract it from the total deposit amount to see previous state.
+    let new_total_deposit_amt = total_deposit_amt.checked_sub(received_amt).unwrap_or_default();
+
+    // Update the total deposit amount for the asset
+    denom_deposits.insert(denom.to_string(), Some(new_total_deposit_amt));
+
+    Ok(())
+}

contracts/credit-manager/src/execute.rs

@@ -1,7 +1,8 @@
-use std::collections::BTreeSet;
+use std::collections::BTreeMap;
 
 use cosmwasm_std::{
-    to_json_binary, Addr, Coins, CosmosMsg, DepsMut, Env, MessageInfo, Response, StdResult, WasmMsg,
+    to_json_binary, Addr, Coins, CosmosMsg, DepsMut, Env, MessageInfo, Response, StdResult,
+    Uint128, WasmMsg,
 };
 use mars_types::{
     account_nft::ExecuteMsg as NftExecuteMsg,
@@ -15,7 +16,7 @@ use crate::{
     borrow::borrow,
     claim_astro_lp_rewards::claim_lp_rewards,
     claim_rewards::claim_rewards,
-    deposit::{assert_deposit_caps, deposit},
+    deposit::{assert_deposit_caps, deposit, update_or_reset_denom_deposits},
     error::{ContractError, ContractResult},
     health::{assert_max_ltv, query_health_state},
     hls::assert_hls_rules,
@@ -151,13 +152,10 @@ pub fn dispatch_actions(
         None
     };
 
-    // We use a Set to record all denoms whose deposited amount may go up as the
+    // We use a Map to record all denoms whose deposited amount may go up as the
     // result of any action. We invoke the AssertDepositCaps callback in the end
     // to make sure that none of the deposit cap is exceeded.
     //
-    // Additionally, we use a BTreeSet (instead of a Vec or HashSet) to ensure
-    // uniqueness and determininism.
-    //
     // There are a few actions that may result in an asset's deposit amount
     // going up:
     // - Deposit: we check the deposited denom
@@ -172,14 +170,22 @@ pub fn dispatch_actions(
     // Note that Borrow/Lend/Reclaim does not impact total deposit amount,
     // because they simply move assets between Red Bank and Rover. We don't
     // check these actions.
-    let mut denoms_for_cap_check = BTreeSet::new();
+    // If there is None in the map, it means that the deposit cap should be enforced,
+    // otherwise it should compare deposit amount before and after the TX.
+    let mut denoms_for_cap_check: BTreeMap<String, Option<Uint128>> = BTreeMap::new();
 
     for action in actions {
         match action {
             Action::Deposit(coin) => {
                 response = deposit(&mut deps, response, account_id, &coin, &mut received_coins)?;
-                // check the deposit cap of the deposited denom
-                denoms_for_cap_check.insert(coin.denom);
+                // add the denom to the map to check the deposit cap in the end of the TX
+                update_or_reset_denom_deposits(
+                    deps.as_ref(),
+                    &mut denoms_for_cap_check,
+                    &coin.denom,
+                    &received_coins,
+                    true,
+                )?;
             }
             Action::Withdraw(coin) => callbacks.push(CallbackMsg::Withdraw {
                 account_id: account_id.to_string(),
@@ -285,8 +291,14 @@ pub fn dispatch_actions(
                     min_receive,
                     route,
                 });
-                // check the deposit cap of the swap output denom
-                denoms_for_cap_check.insert(denom_out);
+                // add the output denom to the map to check the deposit cap in the end of the TX
+                update_or_reset_denom_deposits(
+                    deps.as_ref(),
+                    &mut denoms_for_cap_check,
+                    &denom_out,
+                    &received_coins,
+                    false,
+                )?;
             }
             Action::ExitVault {
                 vault,
@@ -324,8 +336,14 @@ pub fn dispatch_actions(
                     slippage,
                 });
 
-                // check the deposit cap of the LP output denom
-                denoms_for_cap_check.insert(lp_token_out);
+                // add the LP output denom to the map to check the deposit cap in the end of the TX
+                update_or_reset_denom_deposits(
+                    deps.as_ref(),
+                    &mut denoms_for_cap_check,
+                    &lp_token_out,
+                    &received_coins,
+                    false,
+                )?;
             }
             Action::WithdrawLiquidity {
                 lp_token,