fix: item 3 - potential DoS/gas issue

- added a limit so only can transfer N/20 tokens at a time
- iterated via REWARD tokens.

contracts/treasurechest-contract/src/contract.rs

@@ -73,7 +73,7 @@ pub fn execute(
         ExecuteMsg::ChangeTokenFactory {
             token_factory_type,
         } => change_token_factory(deps, info.sender, &token_factory_type),
-        ExecuteMsg::ReturnDust {} => return_dust(deps, env, info.sender),
+        ExecuteMsg::ReturnDust {limit} => return_dust(deps, env, info.sender, limit),
     }
 }
 

contracts/treasurechest-contract/src/executions.rs

@@ -1,9 +1,13 @@
 use std::{ops::Mul, str::FromStr};
+use std::collections::HashMap;
+
+const DEFAULT_SIZE:u32 = 20;
 
 use cosmwasm_std::{
     Addr, BankMsg, Coin, CosmosMsg, DepsMut, Env, Event, MessageInfo, Order, Response, StdResult,
     Uint128,
 };
+
 use treasurechest::{errors::ContractError, tf::tokenfactory::TokenFactoryType};
 
 use crate::state::{CONFIG, TOTAL_REWARDS};
@@ -76,27 +80,32 @@ pub fn change_token_factory(
         .add_attribute("action", "treasurechest/change_token_factory"))
 }
 
-pub fn return_dust(deps: DepsMut, env: Env, sender: Addr) -> Result<Response, ContractError> {
+pub fn return_dust(deps: DepsMut, env: Env, sender: Addr, limit: Option<u32>) -> Result<Response, ContractError> {
     cw_ownable::assert_owner(deps.storage, &sender)?;
     let config = CONFIG.load(deps.storage)?;
     let balances = deps
         .querier
         .query_all_balances(env.contract.address)?
         .into_iter()
-        .filter(|x| x.denom != config.denom)
-        .collect::<Vec<Coin>>();
+        .filter(|x| x.denom != config.denom).map(|coin| (coin.denom,coin.amount))
+        .collect::<HashMap<String,Uint128>>();
     let mut balances_out = vec![];
 
-    for entry in balances {
-        if let Some(one_amt) = TOTAL_REWARDS.may_load(deps.storage, entry.denom.clone())? {
-            if one_amt.to_uint_floor() > entry.amount {
-                TOTAL_REWARDS.remove(deps.storage, entry.denom.clone());
-                balances_out.push(entry)
+    let rewards = TOTAL_REWARDS.range(deps.storage, None,None,Order::Ascending).take(limit.unwrap_or(DEFAULT_SIZE).try_into()?).collect::<StdResult<Vec<_>>>()?;
+    for reward in rewards {
+        let reward_amt = reward.1.to_uint_floor();
+        if let Some(token_balance) = balances.get( &reward.0) {
+            if &reward_amt  > token_balance{
+                TOTAL_REWARDS.remove(deps.storage, reward.0.clone());
+                balances_out.push(Coin{denom: reward.0, amount: token_balance.clone()})
             }
         }
     }
     // balances_out should only contain the dust now.
     // TOTAL rewards should no longer show that token
+    if balances_out.is_empty() {
+       return Ok(Response::new().add_attribute("action", "treasurechest/return_dust").add_attribute("dust","no-dust"))
+    }
 
     let msg = CosmosMsg::Bank(BankMsg::Send {
         to_address: sender.to_string(),

integration/src/treasurechest/contract.rs

@@ -42,8 +42,8 @@ impl TreasureChestContract {
         })
     }
 
-    pub fn return_dust(&self) -> StdResult<CosmosMsg> {
-        self.call(ExecuteMsg::ReturnDust {})
+    pub fn return_dust(&self, limit:Option<u32>) -> StdResult<CosmosMsg> {
+        self.call(ExecuteMsg::ReturnDust {limit})
     }
 
     pub fn config(&self, app: &App) -> StdResult<ConfigResponse> {

integration/src/treasurechest/tests.rs

@@ -335,7 +335,7 @@ pub fn test_withdraw() -> AnyResult<()> {
     assert_eq!(state.outstanding, Uint128::zero());
 
     app.update_block(|f| f.height += 1);
-    let res_admin = app.execute(Addr::unchecked(ADMIN), chest.return_dust()?)?;
+    let res_admin = app.execute(Addr::unchecked(ADMIN), chest.return_dust(None)?)?;
     for e in get_events("transfer", &res_admin.events) {
         if let Some(amount) = get_attribute("amount", &e.attributes) {
             let coins = amount.split(',').map(|x| x.to_string()).collect::<Vec<_>>();

packages/treasurechest/src/chest.rs

@@ -24,7 +24,7 @@ pub enum ExecuteMsg {
     /// Withdraw pending rewards
     Withdraw {},
     /// If balance is below >1< tickets worth (ADMIN only)
-    ReturnDust {},
+    ReturnDust { limit: Option<u32>},
     /// change token factory type (ADMIN only)
     ChangeTokenFactory {
         token_factory_type: String,

packages/treasurechest/src/errors.rs

@@ -1,3 +1,4 @@
+use std::num::TryFromIntError;
 use cosmwasm_std::{CheckedFromRatioError, DivideByZeroError, OverflowError, StdError};
 use cw_ownable::OwnershipError;
 use thiserror::Error;
@@ -10,6 +11,8 @@ pub enum ContractError {
     Ownership(#[from] OwnershipError),
     #[error("{0}")]
     OverflowError(#[from] OverflowError),
+    #[error("{0}")]
+    TryFromIntError(#[from] TryFromIntError),
 
     #[error("{0}")]
     DivideByZeroError(#[from] DivideByZeroError),