Merge pull request #11 from oraichain/feat/claim-protocol-fee

Feat/claim protocol fee

contracts/oraiswap-v3/src/contract.rs

@@ -45,6 +45,9 @@ pub fn execute(
     match msg {
         ExecuteMsg::ChangeAdmin { new_admin } => change_admin(deps, info, new_admin),
         ExecuteMsg::WithdrawProtocolFee { pool_key } => withdraw_protocol_fee(deps, info, pool_key),
+        ExecuteMsg::WithdrawAllProtocolFee { receiver } => {
+            withdraw_all_protocol_fee(deps, info, receiver)
+        }
         ExecuteMsg::ChangeProtocolFee { protocol_fee } => {
             change_protocol_fee(deps, info, protocol_fee)
         }

contracts/oraiswap-v3/src/entrypoints/execute.rs

@@ -2,7 +2,9 @@ use crate::state::{self, CONFIG, POOLS};
 use oraiswap_v3_common::asset::{Asset, AssetInfo};
 use oraiswap_v3_common::error::ContractError;
 use oraiswap_v3_common::incentives_fund_manager;
-use oraiswap_v3_common::interface::{CalculateSwapResult, Cw721ReceiveMsg, SwapHop};
+use oraiswap_v3_common::interface::{
+    CalculateSwapResult, Cw721ReceiveMsg, PoolWithPoolKey, SwapHop,
+};
 use oraiswap_v3_common::math::fee_growth::FeeGrowth;
 use oraiswap_v3_common::math::liquidity::Liquidity;
 use oraiswap_v3_common::math::percentage::Percentage;
@@ -17,7 +19,8 @@ use super::{
     transfer_nft, update_approvals, TimeStampExt,
 };
 use cosmwasm_std::{
-    attr, wasm_execute, Addr, Attribute, Binary, DepsMut, Env, MessageInfo, Response,
+    attr, wasm_execute, Addr, Attribute, Binary, DepsMut, Env, MessageInfo, Order, Response,
+    StdResult,
 };
 use cw20::Expiration;
 use decimal::Decimal;
@@ -54,6 +57,69 @@ pub fn change_admin(
     Ok(Response::new().add_attributes(event_attributes))
 }
 
+/// Allows an fee receiver to withdraw collected fees.
+///
+///
+/// # Errors
+/// - Reverts the call when the caller is an unauthorized receiver.
+pub fn withdraw_all_protocol_fee(
+    deps: DepsMut,
+    mut info: MessageInfo,
+    receiver: Option<Addr>,
+) -> Result<Response, ContractError> {
+    let pools: Vec<PoolWithPoolKey> = POOLS
+        .range_raw(deps.storage, None, None, Order::Ascending)
+        .map(|item| {
+            let (raw_key, pool) = item?;
+            Ok(PoolWithPoolKey {
+                pool_key: PoolKey::from_bytes(&raw_key)?,
+                pool,
+            })
+        })
+        .collect::<StdResult<_>>()?;
+    let mut attrs: Vec<Attribute> = vec![
+        attr("action", "withdraw_protocol_fee"),
+        attr("receiver", info.sender.as_str()),
+    ];
+    let mut msgs = vec![];
+    let sender = info.sender.clone();
+    if let Some(receiver) = receiver {
+        info.sender = receiver;
+    }
+
+    for mut pool_info in pools {
+        if pool_info.pool.fee_receiver != sender {
+            continue;
+        }
+        let pool_key_db = pool_info.pool_key.key();
+        let (fee_protocol_token_x, fee_protocol_token_y) = pool_info.pool.withdraw_protocol_fee();
+        POOLS.save(deps.storage, &pool_key_db, &pool_info.pool)?;
+
+        let asset_0 = Asset {
+            info: AssetInfo::from_denom(deps.api, pool_info.pool_key.token_x.as_str()),
+            amount: fee_protocol_token_x.into(),
+        };
+
+        let asset_1 = Asset {
+            info: AssetInfo::from_denom(deps.api, pool_info.pool_key.token_y.as_str()),
+            amount: fee_protocol_token_y.into(),
+        };
+
+        asset_0.transfer(&mut msgs, &info)?;
+        asset_1.transfer(&mut msgs, &info)?;
+
+        let mut event_attributes = vec![
+            attr("pool_key", pool_info.pool_key.to_string()),
+            attr("token_x", fee_protocol_token_x.to_string()),
+            attr("token_y", fee_protocol_token_y.to_string()),
+        ];
+
+        attrs.append(&mut event_attributes);
+    }
+
+    Ok(Response::new().add_messages(msgs).add_attributes(attrs))
+}
+
 /// Allows an fee receiver to withdraw collected fees.
 ///
 /// # Parameters

contracts/oraiswap-v3/src/tests/helper.rs

@@ -184,6 +184,20 @@ impl MockApp {
         )
     }
 
+    pub fn withdraw_all_protocol_fee(
+        &mut self,
+        sender: &str,
+        dex: &str,
+        receiver: Option<Addr>,
+    ) -> MockResult<ExecuteResponse> {
+        self.execute(
+            Addr::unchecked(sender),
+            Addr::unchecked(dex),
+            &oraiswap_v3_msg::ExecuteMsg::WithdrawAllProtocolFee { receiver },
+            &[],
+        )
+    }
+
     pub fn change_fee_receiver(
         &mut self,
         sender: &str,
@@ -1194,6 +1208,13 @@ pub mod macros {
     }
     pub(crate) use withdraw_protocol_fee;
 
+    macro_rules! withdraw_all_protocol_fee {
+        ($app:ident, $dex_address:expr,$receiver:expr, $caller:tt) => {{
+            $app.withdraw_all_protocol_fee($caller, $dex_address.as_str(), $receiver)
+        }};
+    }
+    pub(crate) use withdraw_all_protocol_fee;
+
     macro_rules! change_fee_receiver {
         ($app:ident,  $dex_address:expr, $pool_key:expr, $fee_receiver:tt, $caller:tt) => {{
             $app.change_fee_receiver($caller, $dex_address.as_str(), &$pool_key, $fee_receiver)

contracts/oraiswap-v3/src/tests/protocol_fee.rs

@@ -1,4 +1,4 @@
-use cosmwasm_std::coins;
+use cosmwasm_std::{coins, Addr};
 use decimal::*;
 
 use crate::tests::helper::{macros::*, MockApp, FEE_DENOM};
@@ -48,6 +48,86 @@ fn test_protocol_fee() {
     );
 }
 
+#[test]
+fn test_withdraw_all_protocol_fee() {
+    let (mut app, accounts) = MockApp::new(&[
+        ("alice", &coins(100_000_000_000, FEE_DENOM)),
+        ("bob", &coins(100_000_000_000, FEE_DENOM)),
+    ]);
+    let alice = &accounts[0];
+    let bob = &accounts[1];
+
+    let (dex, token_x, token_y) = init_dex_and_tokens!(app, alice);
+    init_basic_pool!(app, dex, token_x, token_y, alice);
+    init_basic_position!(app, dex, token_x, token_y, alice);
+    init_basic_swap!(app, dex, token_x, token_y, alice, bob);
+
+    let fee_tier = FeeTier::new(Percentage::from_scale(6, 3), 10).unwrap();
+
+    withdraw_all_protocol_fee!(app, dex, None, alice).unwrap();
+
+    let amount_x = balance_of!(app, token_x, alice);
+    let amount_y = balance_of!(app, token_y, alice);
+    assert_eq!(amount_x, 9999999501);
+    assert_eq!(amount_y, 9999999000);
+
+    let amount_x = balance_of!(app, token_x, dex);
+    let amount_y = balance_of!(app, token_y, dex);
+    assert_eq!(amount_x, 1499);
+    assert_eq!(amount_y, 7);
+
+    let pool_after_withdraw = get_pool!(app, dex, token_x, token_y, fee_tier).unwrap();
+    assert_eq!(
+        pool_after_withdraw.fee_protocol_token_x,
+        TokenAmount::new(0)
+    );
+    assert_eq!(
+        pool_after_withdraw.fee_protocol_token_y,
+        TokenAmount::new(0)
+    );
+}
+
+#[test]
+fn test_withdraw_all_protocol_fee_with_receiver() {
+    let (mut app, accounts) = MockApp::new(&[
+        ("alice", &coins(100_000_000_000, FEE_DENOM)),
+        ("bob", &coins(100_000_000_000, FEE_DENOM)),
+        ("charlie", &coins(100_000_000_000, FEE_DENOM)),
+    ]);
+    let alice = &accounts[0];
+    let bob = &accounts[1];
+    let charlie = &accounts[2];
+
+    let (dex, token_x, token_y) = init_dex_and_tokens!(app, alice);
+    init_basic_pool!(app, dex, token_x, token_y, alice);
+    init_basic_position!(app, dex, token_x, token_y, alice);
+    init_basic_swap!(app, dex, token_x, token_y, alice, bob);
+
+    let fee_tier = FeeTier::new(Percentage::from_scale(6, 3), 10).unwrap();
+
+    withdraw_all_protocol_fee!(app, dex, Some(Addr::unchecked(charlie)), alice).unwrap();
+
+    let amount_x = balance_of!(app, token_x, charlie);
+    let amount_y = balance_of!(app, token_y, charlie);
+    assert_eq!(amount_x, 1);
+    assert_eq!(amount_y, 0);
+
+    let amount_x = balance_of!(app, token_x, dex);
+    let amount_y = balance_of!(app, token_y, dex);
+    assert_eq!(amount_x, 1499);
+    assert_eq!(amount_y, 7);
+
+    let pool_after_withdraw = get_pool!(app, dex, token_x, token_y, fee_tier).unwrap();
+    assert_eq!(
+        pool_after_withdraw.fee_protocol_token_x,
+        TokenAmount::new(0)
+    );
+    assert_eq!(
+        pool_after_withdraw.fee_protocol_token_y,
+        TokenAmount::new(0)
+    );
+}
+
 #[test]
 fn test_protocol_fee_not_admin() {
     let (mut app, accounts) = MockApp::new(&[
@@ -78,6 +158,38 @@ fn test_protocol_fee_not_admin() {
         .contains(&ContractError::Unauthorized {}.to_string()));
 }
 
+#[test]
+fn test_withdraw_all_protocol_fee_not_admin() {
+    let (mut app, accounts) = MockApp::new(&[
+        ("alice", &coins(100_000_000_000, FEE_DENOM)),
+        ("bob", &coins(100_000_000_000, FEE_DENOM)),
+    ]);
+    let alice = &accounts[0];
+    let bob = &accounts[1];
+    let (dex, token_x, token_y) = init_dex_and_tokens!(app, alice);
+    init_basic_pool!(app, dex, token_x, token_y, alice);
+    init_basic_position!(app, dex, token_x, token_y, alice);
+    init_basic_swap!(app, dex, token_x, token_y, alice, bob);
+    let fee_tier = FeeTier::new(Percentage::from_scale(6, 3), 10).unwrap();
+
+    withdraw_all_protocol_fee!(app, dex, Some(Addr::unchecked(alice)), bob).unwrap();
+
+    let amount_x = balance_of!(app, token_x, alice);
+    let amount_y = balance_of!(app, token_y, alice);
+    assert_eq!(amount_x, 9999999500);
+    assert_eq!(amount_y, 9999999000);
+
+    let pool_after_withdraw = get_pool!(app, dex, token_x, token_y, fee_tier).unwrap();
+    assert_eq!(
+        pool_after_withdraw.fee_protocol_token_x,
+        TokenAmount::new(1)
+    );
+    assert_eq!(
+        pool_after_withdraw.fee_protocol_token_y,
+        TokenAmount::new(0)
+    );
+}
+
 #[test]
 fn test_withdraw_fee_not_deployer() {
     let (mut app, accounts) = MockApp::new(&[

packages/oraiswap-v3-common/src/oraiswap_v3_msg.rs

@@ -1,14 +1,21 @@
 #![allow(unused_imports)]
+use crate::asset::{Asset, AssetInfo};
 use cosmwasm_schema::{cw_serde, QueryResponses};
 use cosmwasm_std::{Addr, Binary, Uint64};
 use cw20::Expiration;
-use crate::asset::{Asset, AssetInfo};
 
-use crate::{interface::{
-    AllNftInfoResponse, ApprovedForAllResponse, NftInfoResponse, NumTokensResponse, OwnerOfResponse, PoolWithPoolKey, PositionTick, QuoteResult, SwapHop, TokensResponse
-}, math::{liquidity::Liquidity, percentage::Percentage, sqrt_price::SqrtPrice, token_amount::TokenAmount}, storage::{FeeTier, LiquidityTick, Pool, PoolKey, Position, Tick}};
+use crate::{
+    interface::{
+        AllNftInfoResponse, ApprovedForAllResponse, NftInfoResponse, NumTokensResponse,
+        OwnerOfResponse, PoolWithPoolKey, PositionTick, QuoteResult, SwapHop, TokensResponse,
+    },
+    math::{
+        liquidity::Liquidity, percentage::Percentage, sqrt_price::SqrtPrice,
+        token_amount::TokenAmount,
+    },
+    storage::{FeeTier, LiquidityTick, Pool, PoolKey, Position, Tick},
+};
 #[allow(unused_imports)]
-
 #[cw_serde]
 pub struct InstantiateMsg {
     pub protocol_fee: Percentage,
@@ -33,6 +40,9 @@ pub enum ExecuteMsg {
     WithdrawProtocolFee {
         pool_key: PoolKey,
     },
+    WithdrawAllProtocolFee {
+        receiver: Option<Addr>,
+    },
     ChangeProtocolFee {
         protocol_fee: Percentage,
     },