Merge pull request #12 from apollodao/dev/support-other-reward-tokens

feat: Support other reward token types
apollodao · Nov 3, 2023 · 7b9fd59 · 7b9fd59
2 parents 09c7aef + 12f470a
commit 7b9fd59
Show file tree

Hide file tree

Showing 10 changed files with 340 additions and 104 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/contracts/reward-distributor/src/contract.rs b/contracts/reward-distributor/src/contract.rs
@@ -7,7 +7,7 @@ use cw_dex::astroport::AstroportPool;
 use cw_vault_standard::{VaultContract, VaultContractUnchecked};
 use neutron_astroport_reward_distributor::{
     Config, ConfigUnchecked, ContractError, ExecuteMsg, InstantiateMsg, InternalMsg, QueryMsg,
-    StateResponse, CONFIG, LAST_DISTRIBUTED, REWARD_POOL, REWARD_VAULT,
+    RewardInfo, RewardType, StateResponse, CONFIG, LAST_DISTRIBUTED, REWARD_TOKEN,
 };
 
 use crate::execute;
@@ -25,24 +25,41 @@ pub fn instantiate(
     cw2::set_contract_version(deps.storage, CONTRACT_NAME, CONTRACT_VERSION)?;
     cw_ownable::initialize_owner(deps.storage, deps.api, Some(&msg.owner))?;
 
-    let reward_vault: VaultContract =
-        VaultContractUnchecked::new(&msg.reward_vault_addr).check(deps.as_ref())?;
+    let reward_token = match msg.reward_token_info {
+        RewardInfo::VaultAddr(reward_vault_addr) => {
+            let reward_vault: VaultContract =
+                VaultContractUnchecked::new(&reward_vault_addr).check(deps.as_ref())?;
 
-    // Validate reward vault base token as CW20 Astroport LP token
-    let reward_lp_token = deps
-        .api
-        .addr_validate(&reward_vault.base_token)
-        .map_err(|_| StdError::generic_err("Invalid base token of reward vault"))?;
+            // Validate reward vault base token as CW20 Astroport LP token
+            let reward_lp_token = deps
+                .api
+                .addr_validate(&reward_vault.base_token)
+                .map_err(|_| StdError::generic_err("Invalid base token of reward vault"))?;
 
-    // Query minter of LP token to get reward pool address
-    let minter_res: MinterResponse = deps.querier.query(&QueryRequest::Wasm(WasmQuery::Smart {
-        contract_addr: reward_lp_token.to_string(),
-        msg: to_binary(&Cw20QueryMsg::Minter {})?,
-    }))?;
-    let reward_pool_addr = deps.api.addr_validate(&minter_res.minter)?;
+            // Query minter of LP token to get reward pool address
+            let minter_res: MinterResponse =
+                deps.querier.query(&QueryRequest::Wasm(WasmQuery::Smart {
+                    contract_addr: reward_lp_token.to_string(),
+                    msg: to_binary(&Cw20QueryMsg::Minter {})?,
+                }))?;
+            let reward_pool_addr = deps.api.addr_validate(&minter_res.minter)?;
 
-    // Query reward pool for pool info to create pool object
-    let reward_pool = AstroportPool::new(deps.as_ref(), reward_pool_addr)?;
+            // Query reward pool for pool info to create pool object
+            let reward_pool = AstroportPool::new(deps.as_ref(), reward_pool_addr)?;
+
+            RewardType::Vault {
+                vault: reward_vault,
+                pool: reward_pool,
+            }
+        }
+        RewardInfo::AstroportPoolAddr(pool_addr) => {
+            let reward_pool =
+                AstroportPool::new(deps.as_ref(), deps.api.addr_validate(&pool_addr)?)?;
+
+            RewardType::LP(reward_pool)
+        }
+        RewardInfo::NativeCoin(reward_coin_denom) => RewardType::Coin(reward_coin_denom),
+    };
 
     // Create config
     let config: Config = ConfigUnchecked {
@@ -54,8 +71,7 @@ pub fn instantiate(
 
     CONFIG.save(deps.storage, &config)?;
     LAST_DISTRIBUTED.save(deps.storage, &env.block.time.seconds())?;
-    REWARD_POOL.save(deps.storage, &reward_pool)?;
-    REWARD_VAULT.save(deps.storage, &reward_vault)?;
+    REWARD_TOKEN.save(deps.storage, &reward_token)?;
 
     Ok(Response::default())
 }
@@ -103,14 +119,11 @@ pub fn query(deps: Deps, _env: Env, msg: QueryMsg) -> StdResult<Binary> {
         }
         QueryMsg::State {} => {
             let config = CONFIG.load(deps.storage)?;
-            let reward_pool = REWARD_POOL.load(deps.storage)?;
-            let reward_vault = REWARD_VAULT.load(deps.storage)?;
             let last_distributed = LAST_DISTRIBUTED.load(deps.storage)?;
 
             to_binary(&StateResponse {
                 config,
-                reward_pool,
-                reward_vault,
+                reward_token: REWARD_TOKEN.load(deps.storage)?,
                 last_distributed,
             })
         }

diff --git a/contracts/reward-distributor/src/execute.rs b/contracts/reward-distributor/src/execute.rs
@@ -1,8 +1,10 @@
 use apollo_cw_asset::{Asset, AssetInfo, AssetList};
-use cosmwasm_std::{Deps, DepsMut, Env, Event, MessageInfo, Response, Uint128};
+use cosmwasm_std::{
+    coins, BankMsg, CosmosMsg, Deps, DepsMut, Env, Event, MessageInfo, Response, Uint128,
+};
 use cw_dex::traits::Pool as PoolTrait;
 use neutron_astroport_reward_distributor::{
-    ConfigUpdates, ContractError, InternalMsg, CONFIG, LAST_DISTRIBUTED, REWARD_POOL, REWARD_VAULT,
+    ConfigUpdates, ContractError, InternalMsg, RewardType, CONFIG, LAST_DISTRIBUTED, REWARD_TOKEN,
 };
 
 pub fn execute_distribute(deps: DepsMut, env: Env) -> Result<Response, ContractError> {
@@ -22,55 +24,77 @@ pub fn execute_distribute(deps: DepsMut, env: Env) -> Result<Response, ContractE
 
     // Calculate amount of rewards to be distributed
     let time_elapsed = current_time.saturating_sub(last_distributed.max(config.rewards_start_time));
-    let redeem_amount = config.emission_per_second * Uint128::from(time_elapsed);
-
-    // Query the vault to see how many base tokens would be returned after
-    // redeeming. If zero we return Ok, so that update_config does not fail when
-    // trying to distribute.
-    let base_token_amount = REWARD_VAULT
-        .load(deps.storage)?
-        .query_convert_to_assets(&deps.querier, redeem_amount)?;
-    if base_token_amount.is_zero() {
-        return Ok(Response::new());
-    }
-
-    // Check contract's balance of vault tokens and error if not enough. This is
-    // just so we get a clearer error message rather than the confusing "cannot
-    // sub 0 with x".
-    let reward_vault = REWARD_VAULT.load(deps.storage)?;
-    let vault_token_balance = deps
-        .querier
-        .query_balance(&env.contract.address, &reward_vault.vault_token)?;
-    if vault_token_balance.amount < redeem_amount {
-        return Err(ContractError::InsufficientVaultTokenBalance {
-            vault_token_balance: vault_token_balance.amount,
-            redeem_amount,
-        });
+    let reward_amount = config.emission_per_second * Uint128::from(time_elapsed);
+
+    let reward_token = REWARD_TOKEN.load(deps.storage)?;
+
+    let mut res = Response::new();
+
+    match reward_token {
+        RewardType::Vault { vault, pool: _ } => {
+            // Query the vault to see how many base tokens would be returned after
+            // redeeming. If zero we return Ok, so that update_config does not fail when
+            // trying to distribute.
+            let base_token_amount = vault.query_convert_to_assets(&deps.querier, reward_amount)?;
+            if base_token_amount.is_zero() {
+                return Ok(Response::new());
+            }
+
+            // Check contract's balance of vault tokens and error if not enough. This is
+            // just so we get a clearer error message rather than the confusing "cannot
+            // sub 0 with x".
+            let vault_token_balance = deps
+                .querier
+                .query_balance(&env.contract.address, &vault.vault_token)?;
+            if vault_token_balance.amount < reward_amount {
+                return Err(ContractError::InsufficientVaultTokenBalance {
+                    vault_token_balance: vault_token_balance.amount,
+                    redeem_amount: reward_amount,
+                });
+            }
+
+            // Redeem rewards from the vault
+            let redeem_msg = vault.redeem(reward_amount, None)?;
+
+            // Create internal callback msg
+            let callback_msg = InternalMsg::VaultTokensRedeemed {}.into_cosmos_msg(&env)?;
+
+            res = res.add_message(redeem_msg).add_message(callback_msg);
+        }
+        RewardType::LP(pool) => {
+            // Create message to withdraw liquidity from pool
+            let lp_tokens = Asset::new(AssetInfo::Cw20(pool.lp_token_addr.clone()), reward_amount);
+            res = pool.withdraw_liquidity(deps.as_ref(), &env, lp_tokens, AssetList::new())?;
+
+            // Create internal callback msg
+            let callback_msg = InternalMsg::LpRedeemed {}.into_cosmos_msg(&env)?;
+            res = res.add_message(callback_msg);
+        }
+        RewardType::Coin(reward_coin_denom) => {
+            // Create message to send coins to distribution address
+            let send_msg: CosmosMsg = BankMsg::Send {
+                to_address: config.distribution_addr.to_string(),
+                amount: coins(reward_amount.u128(), reward_coin_denom),
+            }
+            .into();
+            res = res.add_message(send_msg);
+        }
     }
 
     // Set last distributed time to current time
     LAST_DISTRIBUTED.save(deps.storage, &current_time)?;
 
-    // Redeem rewards from the vault
-    let redeem_msg = reward_vault.redeem(redeem_amount, None)?;
-
-    // Create internal callback msg
-    let callback_msg = InternalMsg::VaultTokensRedeemed {}.into_cosmos_msg(&env)?;
-
     let event = Event::new("apollo/neutron-astroport-reward-distributor/execute_distribute")
-        .add_attribute("vault_tokens_redeemed", redeem_amount);
+        .add_attribute("vault_tokens_redeemed", reward_amount);
 
-    Ok(Response::default()
-        .add_message(redeem_msg)
-        .add_message(callback_msg)
-        .add_event(event))
+    Ok(res.add_event(event))
 }
 
 pub fn execute_internal_vault_tokens_redeemed(
     deps: Deps,
     env: Env,
 ) -> Result<Response, ContractError> {
-    let reward_pool = REWARD_POOL.load(deps.storage)?;
+    let reward_pool = REWARD_TOKEN.load(deps.storage)?.into_pool()?;
 
     // Query lp token balance
     let reward_lp_token = AssetInfo::Cw20(reward_pool.lp_token_addr.clone());
@@ -93,7 +117,7 @@ pub fn execute_internal_vault_tokens_redeemed(
 
 pub fn execute_internal_lp_redeemed(deps: Deps, env: Env) -> Result<Response, ContractError> {
     let config = CONFIG.load(deps.storage)?;
-    let reward_pool = REWARD_POOL.load(deps.storage)?;
+    let reward_pool = REWARD_TOKEN.load(deps.storage)?.into_pool()?;
 
     // Query contracts balances of pool assets
     let pool_asset_balances: AssetList = AssetList::query_asset_info_balances(

diff --git a/contracts/reward-distributor/tests/access_control.rs b/contracts/reward-distributor/tests/access_control.rs
@@ -9,7 +9,7 @@ use locked_astroport_vault_test_helpers::robot::LockedAstroportVaultRobot;
 use neutron_astroport_reward_distributor::{ConfigUpdates, ExecuteMsg, InternalMsg};
 use neutron_astroport_reward_distributor_test_helpers as test_helpers;
 
-use test_helpers::robot::RewardDistributorRobot;
+use test_helpers::robot::{RewardDistributorRobot, TestRewardType};
 
 use crate::common::{DEPS_PATH, UNOPTIMIZED_PATH};
 
@@ -30,6 +30,7 @@ fn update_ownership_can_only_be_called_by_admin() {
         DEPS_PATH,
         UNOPTIMIZED_PATH,
         treasury_addr.address(),
+        TestRewardType::VaultToken,
         &admin,
         emission_per_second,
         rewards_start_time,
@@ -67,6 +68,7 @@ fn update_config_can_only_be_called_by_admin() {
         DEPS_PATH,
         UNOPTIMIZED_PATH,
         treasury_addr.address(),
+        TestRewardType::VaultToken,
         &admin,
         emission_per_second,
         rewards_start_time,
@@ -103,6 +105,7 @@ fn internal_msg_can_only_be_called_by_contract() {
         DEPS_PATH,
         UNOPTIMIZED_PATH,
         treasury_addr.address(),
+        TestRewardType::VaultToken,
         &admin,
         emission_per_second,
         rewards_start_time,