Repository: DA0-DA0/dao-contracts (rust)

Issue:
Check that `ActiveThreshold::AbsoluteCount` isn't larger than total token supply.
In `cw20_balances_voting` we currently assert that the active threshold percent is <= 100%. We should also ensure that the active threshold is not larger than the total number of outstanding tokens.

Pull request:
Check absolute count active threshold is not greater than total supply
Closes #269 
Validate absolute count against the token total supply to prevent potential bricking.

Feedback appreciated.

Patch:
diff --git a/contracts/cw20-staked-balance-voting/src/contract.rs b/contracts/cw20-staked-balance-voting/src/contract.rs
index 0facc355..0a293c3e 100644
--- a/contracts/cw20-staked-balance-voting/src/contract.rs
+++ b/contracts/cw20-staked-balance-voting/src/contract.rs
@@ -1,7 +1,7 @@
 #[cfg(not(feature = "library"))]
 use cosmwasm_std::entry_point;
 use cosmwasm_std::{
-    to_binary, Binary, Decimal, Deps, DepsMut, Env, MessageInfo, Reply, Response, StdResult,
+    to_binary, Addr, Binary, Decimal, Deps, DepsMut, Env, MessageInfo, Reply, Response, StdResult,
     SubMsg, Uint128, Uint256, WasmMsg,
 };
 use cw2::set_contract_version;
@@ -39,7 +39,7 @@ pub fn instantiate(
 ) -> Result<Response, ContractError> {
     set_contract_version(deps.storage, CONTRACT_NAME, CONTRACT_VERSION)?;
     DAO.save(deps.storage, &info.sender)?;
-    if let Some(active_threshold) = msg.active_threshold {
+    if let Some(active_threshold) = msg.active_threshold.clone() {
         if let ActiveThreshold::Percentage { percent } = active_threshold {
             if percent > Decimal::percent(100) {
                 return Err(ContractError::InvalidActivePercentage {});
@@ -55,6 +55,9 @@ pub fn instantiate(
         } => {
             let address = deps.api.addr_validate(&address)?;
             TOKEN.save(deps.storage, &address)?;
+            if let Some(ActiveThreshold::AbsoluteCount { count }) = msg.active_threshold {
+                assert_valid_absolute_count_threshold(deps.as_ref(), address.clone(), count)?;
+            }
 
             match staking_contract {
                 StakingInfo::Existing {
@@ -163,6 +166,20 @@ pub fn instantiate(
     }
 }
 
+pub fn assert_valid_absolute_count_threshold(
+    deps: Deps,
+    token_addr: Addr,
+    count: Uint128,
+) -> Result<(), ContractError> {
+    let token_info: cw20::TokenInfoResponse = deps
+        .querier
+        .query_wasm_smart(token_addr, &cw20_base::msg::QueryMsg::TokenInfo {})?;
+    if count > token_info.total_supply {
+        return Err(ContractError::InvalidAbsoluteCount {});
+    }
+    Ok(())
+}
+
 #[cfg_attr(not(feature = "library"), entry_point)]
 pub fn execute(
     deps: DepsMut,
@@ -189,9 +206,15 @@ pub fn execute_update_active_threshold(
     }
 
     if let Some(active_threshold) = new_active_threshold {
-        if let ActiveThreshold::Percentage { percent } = active_threshold {
-            if percent > Decimal::percent(100) {
-                return Err(ContractError::InvalidActivePercentage {});
+        match active_threshold {
+            ActiveThreshold::Percentage { percent } => {
+                if percent > Decimal::percent(100) {
+                    return Err(ContractError::InvalidActivePercentage {});
+                }
+            }
+            ActiveThreshold::AbsoluteCount { count } => {
+                let token = TOKEN.load(deps.storage)?;
+                assert_valid_absolute_count_threshold(deps.as_ref(), token, count)?;
             }
         }
         ACTIVE_THRESHOLD.save(deps.storage, &active_threshold)?;
@@ -331,6 +354,12 @@ pub fn reply(deps: DepsMut, env: Env, msg: Reply) -> Result<Response, ContractEr
                     }
                     let token = deps.api.addr_validate(&res.contract_address)?;
                     TOKEN.save(deps.storage, &token)?;
+
+                    let active_threshold = ACTIVE_THRESHOLD.may_load(deps.storage)?;
+                    if let Some(ActiveThreshold::AbsoluteCount { count }) = active_threshold {
+                        assert_valid_absolute_count_threshold(deps.as_ref(), token.clone(), count)?;
+                    }
+
                     let staking_contract_code_id = STAKING_CONTRACT_CODE_ID.load(deps.storage)?;
                     let unstaking_duration =
                         STAKING_CONTRACT_UNSTAKING_DURATION.load(deps.storage)?;
diff --git a/contracts/cw20-staked-balance-voting/src/error.rs b/contracts/cw20-staked-balance-voting/src/error.rs
index 4ca98fe1..10d2392a 100644
--- a/contracts/cw20-staked-balance-voting/src/error.rs
+++ b/contracts/cw20-staked-balance-voting/src/error.rs
@@ -29,4 +29,7 @@ pub enum ContractError {
 
     #[error("Active threshold percentage must be greater than 0 and less than 1")]
     InvalidActivePercentage {},
+
+    #[error("Absolute count threshold cannot be greater than the total token supply")]
+    InvalidAbsoluteCount {},
 }


Test patch:
diff --git a/contracts/cw20-staked-balance-voting/src/tests.rs b/contracts/cw20-staked-balance-voting/src/tests.rs
index 956c4521..3575ed61 100644
--- a/contracts/cw20-staked-balance-voting/src/tests.rs
+++ b/contracts/cw20-staked-balance-voting/src/tests.rs
@@ -1186,7 +1186,7 @@ fn test_update_active_threshold() {
 
 #[test]
 #[should_panic(expected = "Active threshold percentage must be greater than 0 and less than 1")]
-fn test_active_threshold_invalid() {
+fn test_active_threshold_percentage_invalid() {
     let mut app = App::default();
     let cw20_id = app.store_code(cw20_contract());
     let voting_id = app.store_code(staked_balance_voting_contract());
@@ -1217,3 +1217,37 @@ fn test_active_threshold_invalid() {
         },
     );
 }
+
+#[test]
+#[should_panic(expected = "Absolute count threshold cannot be greater than the total token supply")]
+fn test_active_threshold_absolute_count_invalid() {
+    let mut app = App::default();
+    let cw20_id = app.store_code(cw20_contract());
+    let voting_id = app.store_code(staked_balance_voting_contract());
+    let staking_contract_id = app.store_code(staking_contract());
+
+    instantiate_voting(
+        &mut app,
+        voting_id,
+        InstantiateMsg {
+            token_info: crate::msg::TokenInfo::New {
+                code_id: cw20_id,
+                label: "DAO DAO voting".to_string(),
+                name: "DAO DAO".to_string(),
+                symbol: "DAO".to_string(),
+                decimals: 6,
+                initial_balances: vec![Cw20Coin {
+                    address: CREATOR_ADDR.to_string(),
+                    amount: Uint128::from(200u64),
+                }],
+                marketing: None,
+                unstaking_duration: None,
+                staking_code_id: staking_contract_id,
+                initial_dao_balance: Some(Uint128::from(100u64)),
+            },
+            active_threshold: Some(ActiveThreshold::AbsoluteCount {
+                count: Uint128::new(10000),
+            }),
+        },
+    );
+}