Limit the maximum number of tokens that can be locked

contracts/hydro/src/contract.rs

@@ -14,9 +14,9 @@ use crate::error::ContractError;
 use crate::msg::{ExecuteMsg, InstantiateMsg};
 use crate::query::{QueryMsg, RoundProposalsResponse, UserLockupsResponse};
 use crate::state::{
-    Constants, CovenantParams, LockEntry, Proposal, Tranche, Vote, CONSTANTS, LOCKS_MAP, LOCK_ID,
-    PROPOSAL_MAP, PROPS_BY_SCORE, PROP_ID, TOTAL_ROUND_POWER, TOTAL_VOTED_POWER, TRANCHE_MAP,
-    VOTE_MAP, WHITELIST, WHITELIST_ADMINS,
+    Constants, CovenantParams, LockEntry, Proposal, Tranche, Vote, CONSTANTS, LOCKED_TOKENS,
+    LOCKS_MAP, LOCK_ID, PROPOSAL_MAP, PROPS_BY_SCORE, PROP_ID, TOTAL_ROUND_POWER,
+    TOTAL_VOTED_POWER, TRANCHE_MAP, VOTE_MAP, WHITELIST, WHITELIST_ADMINS,
 };
 use cw_utils::must_pay;
 
@@ -55,9 +55,11 @@ pub fn instantiate(
         round_length: msg.round_length,
         lock_epoch_length: msg.lock_epoch_length,
         first_round_start: msg.first_round_start,
+        max_locked_tokens: msg.max_locked_tokens,
     };
 
     CONSTANTS.save(deps.storage, &state)?;
+    LOCKED_TOKENS.save(deps.storage, &0)?;
     LOCK_ID.save(deps.storage, &0)?;
     PROP_ID.save(deps.storage, &0)?;
 
@@ -125,6 +127,9 @@ pub fn execute(
         ExecuteMsg::RemoveFromWhitelist { covenant_params } => {
             remove_from_whitelist(deps, env, info, covenant_params)
         }
+        ExecuteMsg::UpdateMaxLockedTokens { max_locked_tokens } => {
+            update_max_locked_tokens(deps, info, max_locked_tokens)
+        }
     }
 }
 
@@ -143,6 +148,15 @@ fn lock_tokens(
     validate_lock_duration(constants.lock_epoch_length, lock_duration)?;
     must_pay(&info, &constants.denom)?;
 
+    let amount_to_lock = info.funds[0].amount.u128();
+    let locked_tokens = LOCKED_TOKENS.load(deps.storage)?;
+
+    if locked_tokens + amount_to_lock > constants.max_locked_tokens {
+        return Err(ContractError::Std(StdError::generic_err(
+            "The limit for locking tokens has been reached. No more tokens can be locked.",
+        )));
+    }
+
     // validate that the user does not have too many locks
     if get_lock_count(deps.as_ref(), info.sender.clone()) >= MAX_LOCK_ENTRIES {
         return Err(ContractError::Std(StdError::generic_err(format!(
@@ -161,6 +175,7 @@ fn lock_tokens(
     let lock_id = LOCK_ID.load(deps.storage)?;
     LOCK_ID.save(deps.storage, &(lock_id + 1))?;
     LOCKS_MAP.save(deps.storage, (info.sender, lock_id), &lock_entry)?;
+    LOCKED_TOKENS.save(deps.storage, &(locked_tokens + amount_to_lock))?;
 
     // Calculate and update the total voting power info for current and all
     // future rounds in which the user will have voting power greather than 0
@@ -305,6 +320,13 @@ fn unlock_tokens(deps: DepsMut, env: Env, info: MessageInfo) -> Result<Response,
     let mut response = Response::new().add_attribute("action", "unlock_tokens");
 
     if !send.amount.is_zero() {
+        LOCKED_TOKENS.update(
+            deps.storage,
+            |locked_tokens| -> Result<u128, ContractError> {
+                Ok(locked_tokens - send.amount.u128())
+            },
+        )?;
+
         response = response.add_message(BankMsg::Send {
             to_address: info.sender.to_string(),
             amount: vec![send],
@@ -574,11 +596,7 @@ fn add_to_whitelist(
     info: MessageInfo,
     covenant_params: CovenantParams,
 ) -> Result<Response, ContractError> {
-    // Validate that the sender is a whitelist admin
-    let whitelist_admins = WHITELIST_ADMINS.load(deps.storage)?;
-    if !whitelist_admins.contains(&info.sender) {
-        return Err(ContractError::Unauthorized);
-    }
+    validate_sender_is_whitelist_admin(&deps, &info)?;
 
     // Add covenant_params to whitelist
     let mut whitelist = WHITELIST.load(deps.storage)?;
@@ -602,11 +620,7 @@ fn remove_from_whitelist(
     info: MessageInfo,
     covenant_params: CovenantParams,
 ) -> Result<Response, ContractError> {
-    // Validate that the sender is a whitelist admin
-    let whitelist_admins = WHITELIST_ADMINS.load(deps.storage)?;
-    if !whitelist_admins.contains(&info.sender) {
-        return Err(ContractError::Unauthorized);
-    }
+    validate_sender_is_whitelist_admin(&deps, &info)?;
 
     // Remove covenant_params from whitelist
     let mut whitelist = WHITELIST.load(deps.storage)?;
@@ -616,6 +630,40 @@ fn remove_from_whitelist(
     Ok(Response::new().add_attribute("action", "remove_from_whitelist"))
 }
 
+fn update_max_locked_tokens(
+    deps: DepsMut,
+    info: MessageInfo,
+    max_locked_tokens: u128,
+) -> Result<Response, ContractError> {
+    validate_sender_is_whitelist_admin(&deps, &info)?;
+
+    CONSTANTS.update(
+        deps.storage,
+        |mut constants| -> Result<Constants, ContractError> {
+            constants.max_locked_tokens = max_locked_tokens;
+
+            Ok(constants)
+        },
+    )?;
+
+    Ok(Response::new()
+        .add_attribute("action", "update_max_locked_tokens")
+        .add_attribute("sender", info.sender.clone())
+        .add_attribute("max_locked_tokens", max_locked_tokens.to_string()))
+}
+
+fn validate_sender_is_whitelist_admin(
+    deps: &DepsMut,
+    info: &MessageInfo,
+) -> Result<(), ContractError> {
+    let whitelist_admins = WHITELIST_ADMINS.load(deps.storage)?;
+    if !whitelist_admins.contains(&info.sender) {
+        return Err(ContractError::Unauthorized);
+    }
+
+    Ok(())
+}
+
 #[cfg_attr(not(feature = "library"), entry_point)]
 pub fn query(deps: Deps, env: Env, msg: QueryMsg) -> StdResult<Binary> {
     match msg {
@@ -677,6 +725,7 @@ pub fn query(deps: Deps, env: Env, msg: QueryMsg) -> StdResult<Binary> {
         )?),
         QueryMsg::Whitelist {} => to_json_binary(&query_whitelist(deps)?),
         QueryMsg::WhitelistAdmins {} => to_json_binary(&query_whitelist_admins(deps)?),
+        QueryMsg::TotalLockedTokens => to_json_binary(&LOCKED_TOKENS.load(deps.storage)?),
     }
 }
 

contracts/hydro/src/msg.rs

@@ -11,6 +11,7 @@ pub struct InstantiateMsg {
     pub lock_epoch_length: u64,
     pub tranches: Vec<Tranche>,
     pub first_round_start: Timestamp,
+    pub max_locked_tokens: u128,
     pub whitelist_admins: Vec<String>,
     pub initial_whitelist: Vec<CovenantParams>,
 }
@@ -42,4 +43,7 @@ pub enum ExecuteMsg {
     RemoveFromWhitelist {
         covenant_params: CovenantParams,
     },
+    UpdateMaxLockedTokens {
+        max_locked_tokens: u128,
+    },
 }

contracts/hydro/src/query.rs

@@ -50,6 +50,7 @@ pub enum QueryMsg {
     },
     Whitelist {},
     WhitelistAdmins {},
+    TotalLockedTokens,
 }
 
 #[derive(Serialize, Deserialize, Clone, PartialEq, JsonSchema, Debug, Default)]

contracts/hydro/src/state.rs

@@ -10,8 +10,11 @@ pub struct Constants {
     pub round_length: u64,
     pub lock_epoch_length: u64,
     pub first_round_start: Timestamp,
+    pub max_locked_tokens: u128,
 }
 
+pub const LOCKED_TOKENS: Item<u128> = Item::new("locked_tokens");
+
 pub const LOCK_ID: Item<u64> = Item::new("lock_id");
 
 pub const PROP_ID: Item<u64> = Item::new("prop_id");

contracts/hydro/src/testing.rs

@@ -29,6 +29,7 @@ pub fn get_default_instantiate_msg() -> InstantiateMsg {
             metadata: "tranche 1".to_string(),
         }],
         first_round_start: mock_env().block.time,
+        max_locked_tokens: 1000000,
         initial_whitelist: vec![get_default_covenant_params()],
         whitelist_admins: vec![],
     }
@@ -865,3 +866,83 @@ fn test_too_many_locks() {
         }
     }
 }
+
+#[test]
+fn max_locked_tokens_test() {
+    let (mut deps, mut env, mut info) =
+        (mock_dependencies(), mock_env(), mock_info("addr0000", &[]));
+
+    let mut msg = get_default_instantiate_msg();
+    msg.max_locked_tokens = 2000;
+    msg.whitelist_admins = vec!["addr0001".to_string()];
+
+    let res = instantiate(deps.as_mut(), env.clone(), info.clone(), msg.clone());
+    assert!(res.is_ok());
+
+    // total tokens locked after this action will be 1500
+    info = mock_info("addr0000", &[Coin::new(1500, STATOM.to_string())]);
+    let mut lock_msg = ExecuteMsg::LockTokens {
+        lock_duration: ONE_MONTH_IN_NANO_SECONDS,
+    };
+    let res = execute(deps.as_mut(), env.clone(), info.clone(), lock_msg.clone());
+    assert!(res.is_ok());
+
+    // total tokens locked after this action would be 3000, which is not allowed
+    info = mock_info("addr0000", &[Coin::new(1500, STATOM.to_string())]);
+    let res = execute(deps.as_mut(), env.clone(), info.clone(), lock_msg.clone());
+    assert!(res.is_err());
+    assert!(res
+        .unwrap_err()
+        .to_string()
+        .contains("The limit for locking tokens has been reached. No more tokens can be locked."));
+
+    // total tokens locked after this action will be 2000, which is the cap
+    info = mock_info("addr0000", &[Coin::new(500, STATOM.to_string())]);
+    lock_msg = ExecuteMsg::LockTokens {
+        lock_duration: THREE_MONTHS_IN_NANO_SECONDS,
+    };
+    let res = execute(deps.as_mut(), env.clone(), info.clone(), lock_msg.clone());
+    assert!(res.is_ok());
+
+    // advance the chain by one month plus one nanosecond and unlock the first lockup
+    env.block.time = env.block.time.plus_nanos(ONE_MONTH_IN_NANO_SECONDS + 1);
+    let res = execute(
+        deps.as_mut(),
+        env.clone(),
+        info.clone(),
+        ExecuteMsg::UnlockTokens {},
+    );
+    assert!(res.is_ok());
+
+    // now a user can lock new 1500 tokens
+    info = mock_info("addr0000", &[Coin::new(1500, STATOM.to_string())]);
+    let res = execute(deps.as_mut(), env.clone(), info.clone(), lock_msg.clone());
+    assert!(res.is_ok());
+
+    // a privileged user can update the maximum allowed locked tokens
+    info = mock_info("addr0001", &[]);
+    let update_max_locked_tokens_msg = ExecuteMsg::UpdateMaxLockedTokens {
+        max_locked_tokens: 3000,
+    };
+    let res = execute(
+        deps.as_mut(),
+        env.clone(),
+        info.clone(),
+        update_max_locked_tokens_msg,
+    );
+    assert!(res.is_ok());
+
+    // now a user can lock up to additional 1000 tokens
+    info = mock_info("addr0002", &[Coin::new(1000, STATOM.to_string())]);
+    let res = execute(deps.as_mut(), env.clone(), info.clone(), lock_msg.clone());
+    assert!(res.is_ok());
+
+    // but no more than the cap of 3000 tokens
+    info = mock_info("addr0002", &[Coin::new(1, STATOM.to_string())]);
+    let res = execute(deps.as_mut(), env.clone(), info.clone(), lock_msg.clone());
+    assert!(res.is_err());
+    assert!(res
+        .unwrap_err()
+        .to_string()
+        .contains("The limit for locking tokens has been reached. No more tokens can be locked."));
+}