Voting power limits for dao-voting-native-staked

contracts/voting/dao-voting-native-staked/src/contract.rs

@@ -18,7 +18,8 @@ use crate::msg::{
     QueryMsg, StakerBalanceResponse,
 };
 use crate::state::{
-    Config, ACTIVE_THRESHOLD, CLAIMS, CONFIG, DAO, HOOKS, MAX_CLAIMS, STAKED_BALANCES, STAKED_TOTAL,
+    Config, ACTIVE_THRESHOLD, CLAIMS, CONFIG, DAO, HOOKS, LIMITS, MAX_CLAIMS, STAKED_BALANCES,
+    STAKED_TOTAL,
 };
 
 pub(crate) const CONTRACT_NAME: &str = "crates.io:dao-voting-native-staked";
@@ -115,6 +116,7 @@ pub fn execute(
         }
         ExecuteMsg::AddHook { addr } => execute_add_hook(deps, env, info, addr),
         ExecuteMsg::RemoveHook { addr } => execute_remove_hook(deps, env, info, addr),
+        ExecuteMsg::UpdateLimit { addr, limit } => execute_update_limit(deps, info, addr, limit),
     }
 }
 
@@ -125,6 +127,13 @@ pub fn execute_stake(
 ) -> Result<Response, ContractError> {
     let config = CONFIG.load(deps.storage)?;
     let amount = must_pay(&info, &config.denom)?;
+    let limit = LIMITS.may_load(deps.storage, &info.sender)?;
+
+    if let Some(limit) = limit {
+        if amount > limit {
+            return Err(ContractError::LimitExceeded { limit });
+        }
+    }
 
     STAKED_BALANCES.update(
         deps.storage,
@@ -322,6 +331,36 @@ pub fn execute_remove_hook(
         .add_attribute("hook", addr))
 }
 
+pub fn execute_update_limit(
+    deps: DepsMut,
+    info: MessageInfo,
+    addr: String,
+    limit: Option<Uint128>,
+) -> Result<Response, ContractError> {
+    let dao = DAO.load(deps.storage)?;
+    if info.sender != dao {
+        return Err(ContractError::Unauthorized {});
+    }
+
+    let addr = deps.api.addr_validate(&addr)?;
+
+    if let Some(limit) = limit {
+        LIMITS.save(deps.storage, &addr, &limit)?;
+    } else {
+        LIMITS.remove(deps.storage, &addr);
+    }
+
+    Ok(Response::new()
+        .add_attribute("action", "update_limit")
+        .add_attribute("addr", addr.to_string())
+        .add_attribute(
+            "limit",
+            limit
+                .as_ref()
+                .map_or("None".to_owned(), ToString::to_string),
+        ))
+}
+
 #[cfg_attr(not(feature = "library"), entry_point)]
 pub fn query(deps: Deps, env: Env, msg: QueryMsg) -> StdResult<Binary> {
     match msg {
@@ -342,6 +381,9 @@ pub fn query(deps: Deps, env: Env, msg: QueryMsg) -> StdResult<Binary> {
         QueryMsg::IsActive {} => query_is_active(deps),
         QueryMsg::ActiveThreshold {} => query_active_threshold(deps),
         QueryMsg::GetHooks {} => to_binary(&query_hooks(deps)?),
+        QueryMsg::Limit { addr } => {
+            to_binary(&LIMITS.may_load(deps.storage, &deps.api.addr_validate(&addr)?)?)
+        }
     }
 }
 

contracts/voting/dao-voting-native-staked/src/error.rs

@@ -1,4 +1,4 @@
-use cosmwasm_std::StdError;
+use cosmwasm_std::{StdError, Uint128};
 use cw_utils::PaymentError;
 use thiserror::Error;
 
@@ -39,4 +39,7 @@ pub enum ContractError {
 
     #[error("Amount being unstaked must be non-zero")]
     ZeroUnstake {},
+
+    #[error("Limit cannot be exceeded")]
+    LimitExceeded { limit: Uint128 },
 }

contracts/voting/dao-voting-native-staked/src/msg.rs

@@ -1,7 +1,9 @@
 use cosmwasm_schema::{cw_serde, QueryResponses};
 use cosmwasm_std::Uint128;
 use cw_utils::Duration;
-use dao_dao_macros::{active_query, voting_module_query};
+use dao_dao_macros::{
+    active_query, limitable_voting_module, limitable_voting_module_query, voting_module_query,
+};
 use dao_voting::threshold::{ActiveThreshold, ActiveThresholdResponse};
 
 #[cw_serde]
@@ -15,6 +17,7 @@ pub struct InstantiateMsg {
     pub active_threshold: Option<ActiveThreshold>,
 }
 
+#[limitable_voting_module]
 #[cw_serde]
 pub enum ExecuteMsg {
     /// Stakes tokens with the contract to get voting power in the DAO
@@ -37,6 +40,7 @@ pub enum ExecuteMsg {
     RemoveHook { addr: String },
 }
 
+#[limitable_voting_module_query]
 #[voting_module_query]
 #[active_query]
 #[cw_serde]

contracts/voting/dao-voting-native-staked/src/state.rs

@@ -2,7 +2,7 @@ use cosmwasm_schema::cw_serde;
 use cosmwasm_std::{Addr, Uint128};
 use cw_controllers::Claims;
 use cw_hooks::Hooks;
-use cw_storage_plus::{Item, SnapshotItem, SnapshotMap, Strategy};
+use cw_storage_plus::{Item, Map, SnapshotItem, SnapshotMap, Strategy};
 use cw_utils::Duration;
 use dao_voting::threshold::ActiveThreshold;
 
@@ -26,6 +26,9 @@ pub const STAKED_BALANCES: SnapshotMap<&Addr, Uint128> = SnapshotMap::new(
     Strategy::EveryBlock,
 );
 
+/// Keeps track of limits by address
+pub const LIMITS: Map<&Addr, Uint128> = Map::new("limits");
+
 /// Keeps track of staked total over time
 pub const STAKED_TOTAL: SnapshotItem<Uint128> = SnapshotItem::new(
     "total_staked",