Check for deposit funds length at entry_point

quasar-finance · Aug 23, 2024 · 8dd4a28 · 8dd4a28
1 parent d4ac47a
commit 8dd4a28
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 41 deletions.
diff --git a/smart-contracts/osmosis/contracts/cl-vault/src/error.rs b/smart-contracts/osmosis/contracts/cl-vault/src/error.rs
@@ -167,6 +167,13 @@ pub enum ContractError {
     TryFromIntError(#[from] TryFromIntError),
 }
 
+pub fn assert_deposits(funds: &[Coin]) -> Result<(), ContractError> {
+    if funds.len() != 2 && funds.len() != 1 {
+        return Err(ContractError::IncorrectAmountFunds);
+    }
+    Ok(())
+}
+
 pub fn assert_admin(storage: &dyn Storage, caller: &Addr) -> Result<(), ContractError> {
     if ADMIN_ADDRESS.load(storage)? != caller {
         return Err(ContractError::Unauthorized {});

diff --git a/smart-contracts/osmosis/contracts/cl-vault/src/helpers/assert.rs b/smart-contracts/osmosis/contracts/cl-vault/src/helpers/assert.rs
@@ -5,10 +5,6 @@ pub(crate) fn must_pay_one_or_two(
     funds: &[Coin],
     denoms: (String, String),
 ) -> Result<(Coin, Coin), ContractError> {
-    if funds.len() != 2 && funds.len() != 1 {
-        return Err(ContractError::IncorrectAmountFunds);
-    }
-
     let token0 = funds
         .iter()
         .find(|coin| coin.denom == denoms.0)
@@ -24,29 +20,6 @@ pub(crate) fn must_pay_one_or_two(
     Ok((token0, token1))
 }
 
-pub(crate) fn must_pay_one_or_two_from_balance(
-    funds: Vec<Coin>,
-    denoms: (String, String),
-) -> Result<(Coin, Coin), ContractError> {
-    if funds.len() < 2 {
-        return Err(ContractError::IncorrectAmountFunds);
-    }
-
-    let token0 = funds
-        .clone()
-        .into_iter()
-        .find(|coin| coin.denom == denoms.0)
-        .unwrap_or(coin(0, denoms.0));
-
-    let token1 = funds
-        .clone()
-        .into_iter()
-        .find(|coin| coin.denom == denoms.1)
-        .unwrap_or(coin(0, denoms.1));
-
-    Ok((token0, token1))
-}
-
 #[cfg(test)]
 mod tests {
 
@@ -76,15 +49,6 @@ mod tests {
         assert_eq!(expected1, token1);
     }
 
-    #[test]
-    fn must_pay_one_or_two_rejects_three() {
-        let expected0 = coin(100, "uatom");
-        let expected1 = coin(200, "uosmo");
-        let funds = vec![expected0, expected1, coin(200, "uqsr")];
-        let _err =
-            must_pay_one_or_two(&funds, ("uatom".to_string(), "uosmo".to_string())).unwrap_err();
-    }
-
     #[test]
     fn must_pay_one_or_two_accepts_second_token() {
         let funds = vec![coin(200, "uosmo")];

diff --git a/smart-contracts/osmosis/contracts/cl-vault/src/vault/autocompound.rs b/smart-contracts/osmosis/contracts/cl-vault/src/vault/autocompound.rs
@@ -6,9 +6,8 @@ use osmosis_std::types::cosmos::bank::v1beta1::{Input, MsgMultiSend, Output};
 use osmosis_std::types::osmosis::concentratedliquidity::v1beta1::ConcentratedliquidityQuerier;
 use osmosis_std::types::osmosis::concentratedliquidity::v1beta1::MsgCreatePositionResponse;
 
-use crate::helpers::assert::must_pay_one_or_two_from_balance;
 use crate::helpers::coinlist::CoinList;
-use crate::helpers::getters::get_unused_balances;
+use crate::helpers::getters::get_unused_pair_balances;
 use crate::msg::{ExecuteMsg, MergePositionMsg};
 use crate::reply::Replies;
 #[allow(deprecated)]
@@ -38,11 +37,11 @@ pub fn execute_autocompound(
         .position
         .ok_or(ContractError::PositionNotFound)?;
 
-    let balance = get_unused_balances(&deps.querier, env)?;
     let pool = POOL_CONFIG.load(deps.storage)?;
+    let balance = get_unused_pair_balances(&deps, env, &pool)?;
 
-    let (token0, token1) =
-        must_pay_one_or_two_from_balance(balance.coins(), (pool.token0, pool.token1))?;
+    let token0 = balance[0].clone();
+    let token1 = balance[1].clone();
 
     // Create coins_to_send with no zero amounts
     let mut coins_to_send = vec![];

diff --git a/smart-contracts/osmosis/contracts/cl-vault/src/vault/deposit.rs b/smart-contracts/osmosis/contracts/cl-vault/src/vault/deposit.rs
@@ -1,4 +1,5 @@
 use crate::{
+    error::assert_deposits,
     helpers::{
         getters::{
             get_depositable_tokens, get_single_sided_deposit_0_to_1_swap_amount,
@@ -30,6 +31,7 @@ pub(crate) fn execute_exact_deposit(
     info: MessageInfo,
     recipient: Option<String>,
 ) -> Result<Response, ContractError> {
+    assert_deposits(&info.funds)?;
     let recipient = recipient.map_or(Ok(info.sender.clone()), |x| deps.api.addr_validate(&x))?;
     let pool_config = POOL_CONFIG.load(deps.storage)?;
     let deposit_info = get_depositable_tokens(&deps, &info.funds, &pool_config)?;
@@ -44,6 +46,7 @@ pub(crate) fn execute_any_deposit(
     recipient: Option<String>,
     max_slippage: Decimal,
 ) -> Result<Response, ContractError> {
+    assert_deposits(&info.funds)?;
     let recipient = recipient.map_or(Ok(info.sender.clone()), |x| deps.api.addr_validate(&x))?;
 
     let pool_config = POOL_CONFIG.load(deps.storage)?;