feat: add whitelist executors for collecting fees

src/contract.rs

@@ -4,7 +4,7 @@ use crate::msg::{
     CollectFeeRequirement, ConfigResponse, DistributeTargetsResponse, ExecuteMsg, InstantiateMsg,
     MigrateMsg, QueryMsg,
 };
-use crate::state::{Config, DistributeTarget, CONFIG, DISTRIBUTION_TARGETS};
+use crate::state::{Config, DistributeTarget, CONFIG, DISTRIBUTION_TARGETS, EXECUTORS};
 use crate::ContractError;
 use cosmos_sdk_proto::cosmos::authz::v1beta1::MsgExec;
 use cosmos_sdk_proto::cosmos::bank::v1beta1::MsgSend;
@@ -59,6 +59,17 @@ pub fn instantiate(
         .collect::<Result<Vec<DistributeTarget>, ContractError>>()?;
 
     DISTRIBUTION_TARGETS.save(deps.storage, &valid_distribute_targets)?;
+
+    let valid_executors = msg
+        .executors
+        .iter()
+        .map(|executor| deps.api.addr_validate(executor.as_ref()))
+        .collect::<StdResult<Vec<Addr>>>()?;
+
+    valid_executors.iter().for_each(|executor| {
+        EXECUTORS.save(deps.storage, executor, &true).unwrap();
+    });
+
     Ok(Response::default())
 }
 
@@ -83,6 +94,23 @@ pub fn execute(
         ExecuteMsg::CollectFees {
             collect_fee_requirements,
         } => execute_collect_fees(deps, env, info, collect_fee_requirements),
+        ExecuteMsg::UpdateExecutors {
+            executors,
+            permission,
+        } => {
+            let valid_executors = executors
+                .iter()
+                .map(|executor| deps.api.addr_validate(executor.as_str()))
+                .collect::<StdResult<Vec<Addr>>>()?;
+
+            valid_executors.iter().for_each(|executor| {
+                EXECUTORS.save(deps.storage, executor, &permission).unwrap();
+            });
+
+            Ok(Response::new()
+                .add_attribute("action", "update_executors")
+                .add_attribute("permission", permission.to_string()))
+        }
     }
 }
 
@@ -174,11 +202,20 @@ fn execute_distribute(
 pub fn execute_collect_fees(
     deps: DepsMut,
     env: Env,
-    _info: MessageInfo,
+    info: MessageInfo,
     collect_fee_requirements: Vec<CollectFeeRequirement>,
 ) -> Result<Response, ContractError> {
+    if !EXECUTORS
+        .load(deps.storage, &info.sender)
+        .unwrap_or_default()
+    {
+        return Err(ContractError::Unauthorized {});
+    }
+
     let config = CONFIG.load(deps.storage)?;
+
     let mut messages: Vec<CosmosMsg> = vec![];
+
     if config.router.is_none() {
         return Err(ContractError::RouterAndApproverNotSet {});
     }
@@ -364,6 +401,9 @@ pub fn query(deps: Deps, _env: Env, msg: QueryMsg) -> StdResult<Binary> {
         QueryMsg::DistributeTargets {} => to_json_binary(&DistributeTargetsResponse(
             DISTRIBUTION_TARGETS.load(deps.storage)?,
         )),
+        QueryMsg::IsExecutor { addr } => {
+            to_json_binary(&EXECUTORS.load(deps.storage, &addr).unwrap_or_default())
+        }
     }
 }
 
@@ -405,6 +445,7 @@ mod tests {
             distribute_token: Addr::unchecked("distribute_token"),
             init_distribution_targets: init_distribution_targets.clone(),
             router: Some(Addr::unchecked("router")),
+            executors: vec![Addr::unchecked("owner"), Addr::unchecked("executor")],
         };
 
         let mock_info = mock_info("owner", &[]);

src/msg.rs

@@ -10,20 +10,31 @@ pub struct InstantiateMsg {
     pub distribute_token: Addr,
     pub router: Option<Addr>,
     pub init_distribution_targets: Vec<DistributeTarget>,
+    pub executors: Vec<Addr>,
 }
 
 #[cw_serde]
 pub enum ExecuteMsg {
+    /////////////////
+    /// Owner API ///
+    ////////////////
     UpdateConfig {
         owner: Option<Addr>,
         distribute_token: Option<Addr>,
     },
     UpdateDistributeTarget {
         distribute_targets: Vec<DistributeTarget>,
     },
+    UpdateExecutors {
+        executors: Vec<Addr>,
+        permission: bool,
+    },
     Distribute {
         amount_distribute: Uint128,
     },
+    /////////////////
+    ///Executors////
+    ///////////////
     CollectFees {
         collect_fee_requirements: Vec<CollectFeeRequirement>,
     },
@@ -43,6 +54,8 @@ pub enum QueryMsg {
     Config {},
     #[returns(DistributeTargetsResponse)]
     DistributeTargets {},
+    #[returns(bool)]
+    IsExecutor { addr: Addr },
 }
 
 #[cw_serde]

src/multitest/contract.rs

@@ -40,6 +40,7 @@ impl TreasuryContract {
                 distribute_token: distribute_token.clone(),
                 init_distribution_targets,
                 router: Some(router.clone()),
+                executors: vec![owner.clone()],
             },
             &[],
             "treasury contract",

src/multitest/tests.rs

@@ -1,10 +1,11 @@
-use crate::contract::execute_collect_fees;
-use crate::msg::{CollectFeeRequirement, ExecuteMsg};
-use crate::state::{Config, CONFIG};
+use crate::contract::{execute_collect_fees, query};
+use crate::msg::{CollectFeeRequirement, ExecuteMsg, QueryMsg};
+use crate::state::{Config, CONFIG, EXECUTORS};
 use crate::{state::DistributeTarget, ContractError};
 use cosmwasm_std::testing::{mock_dependencies, mock_env, mock_info, MockApi};
 use cosmwasm_std::{
-    coin, to_json_binary, Addr, Empty, Event, GovMsg, IbcMsg, IbcQuery, MemoryStorage, Uint128,
+    coin, from_json, to_json_binary, Addr, Empty, Event, GovMsg, IbcMsg, IbcQuery, MemoryStorage,
+    Uint128,
 };
 use cw20::{BalanceResponse, Cw20ExecuteMsg};
 use cw_multi_test::{
@@ -221,6 +222,21 @@ fn test_execute_collect_fees_router_approver_not_set() {
             },
         )
         .unwrap();
+    EXECUTORS
+        .save(deps.as_mut().storage, &Addr::unchecked("sender"), &true)
+        .unwrap();
+
+    let res_binary = query(
+        deps.as_ref(),
+        mock_env(),
+        QueryMsg::IsExecutor {
+            addr: Addr::unchecked("sender"),
+        },
+    )
+    .unwrap();
+
+    assert!(from_json::<bool>(&res_binary).unwrap());
+
     let result = execute_collect_fees(
         deps.as_mut(),
         mock_env(),
@@ -233,7 +249,11 @@ fn test_execute_collect_fees_router_approver_not_set() {
     )
     .unwrap_err();
     assert_eq!(result, ContractError::RouterAndApproverNotSet {});
+}
 
+#[test]
+fn test_execute_collect_fees_unauthorize() {
+    let mut deps = mock_dependencies();
     CONFIG
         .save(
             deps.as_mut().storage,
@@ -248,15 +268,16 @@ fn test_execute_collect_fees_router_approver_not_set() {
     let result = execute_collect_fees(
         deps.as_mut(),
         mock_env(),
-        mock_info("sender", &vec![]),
+        mock_info("owner", &[]),
         vec![CollectFeeRequirement {
             approver: Addr::unchecked("owner"),
             swap_operations: vec![],
             minimum_receive: None,
         }],
     )
     .unwrap_err();
-    assert_eq!(result, ContractError::RouterAndApproverNotSet {});
+
+    assert_eq!(result, ContractError::Unauthorized {});
 }
 
 #[test]

src/state.rs

@@ -1,6 +1,6 @@
 use cosmwasm_schema::cw_serde;
 use cosmwasm_std::{Addr, Binary};
-use cw_storage_plus::Item;
+use cw_storage_plus::{Item, Map};
 
 const CONFIG_KEY: &str = "config";
 const DISTRIBUTION_TARGET: &str = "distribution_target";
@@ -21,3 +21,4 @@ pub struct DistributeTarget {
 
 pub const CONFIG: Item<Config> = Item::new(CONFIG_KEY);
 pub const DISTRIBUTION_TARGETS: Item<Vec<DistributeTarget>> = Item::new(DISTRIBUTION_TARGET);
+pub const EXECUTORS: Map<&Addr, bool> = Map::new("executors");