Skip to content

Commit

Permalink
feat: Add rate lmited withdrawals
Browse files Browse the repository at this point in the history
    Add rate limited withdrawals to asset manager and multi sig actions to XCall manager
  • Loading branch information
AntonAndell committed Mar 14, 2024
1 parent 811d521 commit cf2b929
Show file tree
Hide file tree
Showing 12 changed files with 335 additions and 15 deletions.
146 changes: 133 additions & 13 deletions contracts/core-contracts/cw-asset-manager/src/contract.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ use cw_common::asset_manager_msg::{ExecuteMsg, InstantiateMsg, MigrateMsg, Query
use cw_common::helpers::{get_fee, get_protocols, verify_protocol};
use cw_common::network_address::IconAddressValidation;
use cw_common::network_address::NetworkAddress;
use cw_common::rate_limit::RateLimit;
use cw_common::x_call_msg::XCallMsg;
use cw_common::xcall_data_types::Deposit;

Expand Down Expand Up @@ -82,6 +83,26 @@ pub fn execute(
ensure_eq!(owner, info.sender, ContractError::OnlyOwner);
exec::setup_native_token(deps, native_token_address, native_token_manager)
}
ExecuteMsg::ConfigureRateLimit {
asset,
period,
percentage,
} => {
let owner = OWNER.load(deps.storage).map_err(ContractError::Std)?;
ensure_eq!(owner, info.sender, ContractError::OnlyOwner);
let limit = RateLimit {
period,
percentage,
last_update: env.block.time.seconds(),
current_limit: 0,
};

RATE_LIMITS
.save(deps.storage, asset, &limit)
.map_err(ContractError::Std)?;

Ok(Response::default())
}
ExecuteMsg::DepositDenom { denom, to, data } => {
ensure!(
info.funds.len() == 2 || info.funds.len() == 1,
Expand Down Expand Up @@ -185,15 +206,16 @@ pub fn execute(
}

mod exec {
use std::str::FromStr;
use std::{str::FromStr, u128};

use cosmwasm_std::{BankMsg, Coin, CosmosMsg};
use cw_ibc_rlp_lib::rlp::Encodable;

use cw_common::{helpers::query_network_address, xcall_data_types::DepositRevert};
use cw_xcall_lib::network_address::NetId;

use super::*;
use cw_common::{
helpers::query_network_address, rate_limit::RateLimited, xcall_data_types::DepositRevert,
};
use cw_xcall_lib::network_address::NetId;

pub fn setup(
deps: DepsMut,
Expand Down Expand Up @@ -383,14 +405,13 @@ mod exec {

pub fn handle_xcall_msg(
deps: DepsMut,
_env: Env,
env: Env,
info: MessageInfo,
from: String,
data: Vec<u8>,
) -> Result<Response, ContractError> {
let x_call_addr = SOURCE_XCALL.load(deps.storage)?;
let x_network = X_CALL_NETWORK_ADDRESS.load(deps.storage)?;

if info.sender != x_call_addr {
return Err(ContractError::OnlyXcallService);
}
Expand All @@ -407,7 +428,7 @@ mod exec {
let account = data.account;
let amount = Uint128::from(data.amount);

transfer_tokens(deps, account, token_address, amount)?
transfer_tokens(deps, env, account, token_address, amount)?
}

DecodedStruct::WithdrawTo(data_struct) => {
Expand All @@ -420,7 +441,7 @@ mod exec {
let account = data_struct.user_address;
let amount = Uint128::from(data_struct.amount);

transfer_tokens(deps, account, token_address, amount)?
transfer_tokens(deps, env, account, token_address, amount)?
}

DecodedStruct::WithdrawNativeTo(data_struct) => {
Expand All @@ -433,7 +454,7 @@ mod exec {
let account = data_struct.user_address;
let amount = Uint128::from(data_struct.amount);

swap_to_native(deps, account, token_address, amount)?
swap_to_native(deps, env, account, token_address, amount)?
}
};

Expand All @@ -443,14 +464,17 @@ mod exec {
//internal function to transfer tokens from contract to account
fn transfer_tokens(
deps: DepsMut,
env: Env,
account: String,
token_address: String,
amount: Uint128,
) -> Result<Response, ContractError> {
deps.api.addr_validate(&account)?;
let addr = deps.api.addr_validate(&token_address);
let is_denom = addr.is_err() || !is_contract(deps.querier, &addr.unwrap());

if addr.is_err() || !is_contract(deps.querier, &addr.unwrap()) {
verify_withdraw(deps, env, token_address.clone(), amount.u128(), is_denom)?;
if is_denom {
let coin = Coin {
denom: token_address,
amount,
Expand Down Expand Up @@ -487,6 +511,7 @@ mod exec {
#[cfg(feature = "archway")]
fn swap_to_native(
deps: DepsMut,
env: Env,
account: String,
token_address: String,
amount: Uint128,
Expand All @@ -502,6 +527,8 @@ mod exec {
.query_wasm_smart::<ConfigResponse>(manager, &query_msg)?;
let swap_contract = query_resp.swap_contract_addr;

verify_withdraw(deps, env, token_address.clone(), amount.u128(), false)?;

let hook = &Cw20HookMsg::Swap {
belief_price: None,
max_spread: None,
Expand All @@ -528,14 +555,41 @@ mod exec {
Ok(Response::new().add_submessage(sub_msg))
}

pub fn verify_withdraw(
deps: DepsMut,
env: Env,
asset: String,
amount: u128,
is_denom: bool,
) -> Result<(), ContractError> {
let limit_res = RATE_LIMITS.load(deps.storage, asset.clone());
if limit_res.is_err() {
return Ok(());
}

let limit = limit_res?;
let new_rate_limit =
match limit.verify_withdraw(&deps, env, asset.clone(), amount, is_denom) {
Ok(rl) => rl,
Err(_e) => return Err(ContractError::RateLimit),
};

RATE_LIMITS
.save(deps.storage, asset, &new_rate_limit)
.map_err(ContractError::Std)?;

Ok(())
}

#[cfg(not(any(feature = "archway")))]
fn swap_to_native(
deps: DepsMut,
env: Env,
account: String,
token_address: String,
amount: Uint128,
) -> Result<Response, ContractError> {
transfer_tokens(deps, account, token_address, amount)
transfer_tokens(deps, env, account, token_address, amount)
}

pub fn calculate_denom_funds(
Expand Down Expand Up @@ -630,14 +684,15 @@ mod query {
mod tests {
use cosmwasm_std::{
testing::{mock_dependencies, mock_env, mock_info, MockApi, MockQuerier},
Coin, ContractInfoResponse, ContractResult, MemoryStorage, OwnedDeps, SystemError,
SystemResult, Uint128, WasmQuery,
BlockInfo, Coin, ContractInfoResponse, ContractResult, MemoryStorage, OwnedDeps,
SystemError, SystemResult, Uint128, WasmQuery,
};
use cw_common::xcall_manager_msg::QueryMsg::GetProtocols;
use cw_xcall_multi::msg::QueryMsg::GetNetworkAddress;

use cw_ibc_rlp_lib::rlp::Encodable;
use std::vec;
use tests::exec::verify_withdraw;

use cw_common::{asset_manager_msg::InstantiateMsg, xcall_data_types::WithdrawTo};
use cw_common::{xcall_data_types::DepositRevert, xcall_manager_msg::ProtocolConfig};
Expand Down Expand Up @@ -1040,6 +1095,71 @@ mod tests {
assert!(resp.is_ok());
}

#[test]
fn test_rate_limit() {
let (mut deps, env, _, _) = test_setup();

let token = "denom/ibc-ics-20/token";
let owner = OWNER.load(&deps.storage).unwrap();

// 10% every 100 seconds || protect 90% of fund in a timeframe of 100 seconds
let exe_msg = ExecuteMsg::ConfigureRateLimit {
asset: token.to_string(),
period: 100,
percentage: 9000,
};
let mock_info = mock_info(&owner.to_string(), &[]);

let resp = execute(deps.as_mut(), env.clone(), mock_info.clone(), exe_msg);
assert!(resp.is_ok());

deps.querier.update_balance(
env.clone().contract.address,
vec![Coin {
denom: token.to_string(),
amount: Uint128::new(1000),
}],
);
let res = verify_withdraw(deps.as_mut(), env.clone(), token.to_string(), 100, true);
assert!(res.is_ok());
deps.querier.update_balance(
env.clone().contract.address,
vec![Coin {
denom: token.to_string(),
amount: Uint128::new(900),
}],
);

let limit = RATE_LIMITS
.load(deps.as_mut().storage, token.to_string())
.unwrap();
assert_eq!(limit.current_limit, 900);
assert_eq!(limit.last_update, env.block.time.seconds());
let res = verify_withdraw(deps.as_mut(), env.clone(), token.to_string(), 100, true);
assert!(res.is_err());

let block_info = BlockInfo {
height: env.block.height,
time: env.block.time.plus_seconds(70),
chain_id: env.block.chain_id,
};
let mock_env = Env {
block: block_info,
transaction: env.transaction,
contract: env.contract,
};
// 50 seconds passes we should be able to withdraw up to 63 tokens
let res = verify_withdraw(deps.as_mut(), mock_env.clone(), token.to_string(), 50, true);
assert!(res.is_ok());

let limit = RATE_LIMITS
.load(deps.as_mut().storage, token.to_string())
.unwrap();
// Limit should be 900-63 = 837
assert_eq!(limit.current_limit, 837);
assert_eq!(limit.last_update, mock_env.block.time.seconds());
}

#[cfg(feature = "archway")]
#[test]
fn test_withdraw_native_archway() {
Expand Down
3 changes: 3 additions & 0 deletions contracts/core-contracts/cw-asset-manager/src/error.rs
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,9 @@ pub enum ContractError {

#[error("Invalid funds")]
InvalidFunds,

#[error("RateLimit exceeded")]
RateLimit,
}

impl From<DecoderError> for ContractError {
Expand Down
4 changes: 4 additions & 0 deletions contracts/core-contracts/cw-asset-manager/src/state.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
use cosmwasm_std::Addr;
use cw_common::rate_limit::RateLimit;
use cw_storage_plus::Item;
use cw_storage_plus::Map;

use cw_common::network_address::{NetId, NetworkAddress};

Expand All @@ -16,3 +18,5 @@ pub const NATIVE_TOKEN_ADDRESS: Item<Addr> = Item::new("native_token_address");
pub const NATIVE_TOKEN_MANAGER: Item<Addr> = Item::new("native_token_manager");

pub const X_CALL_MANAGER: Item<Addr> = Item::new("xcall_manager");

pub const RATE_LIMITS: Map<String, RateLimit> = Map::new("rate_limits");
24 changes: 24 additions & 0 deletions contracts/core-contracts/cw-xcall-manager/src/contract.rs
Original file line number Diff line number Diff line change
Expand Up @@ -87,6 +87,24 @@ pub fn execute(
PROPOSER.save(deps.storage, &proposer)?;
Ok(Response::new())
}
ExecuteMsg::WhitelistAction { action } => {
let current_proposer = PROPOSER.load(deps.storage)?;
if info.sender != current_proposer {
return Err(ContractError::OnlyProposer);

Check warning on line 93 in contracts/core-contracts/cw-xcall-manager/src/contract.rs

View check run for this annotation

Codecov / codecov/patch

contracts/core-contracts/cw-xcall-manager/src/contract.rs#L93

Added line #L93 was not covered by tests
}

WHITELISTED_ACTIONS.save(deps.storage, action, &true)?;
Ok(Response::new())
}
ExecuteMsg::RemoveAction { action } => {
let current_proposer = PROPOSER.load(deps.storage)?;
if info.sender != current_proposer {
return Err(ContractError::OnlyProposer);
}

WHITELISTED_ACTIONS.remove(deps.storage, action);
Ok(Response::new())

Check warning on line 106 in contracts/core-contracts/cw-xcall-manager/src/contract.rs

View check run for this annotation

Codecov / codecov/patch

contracts/core-contracts/cw-xcall-manager/src/contract.rs#L99-L106

Added lines #L99 - L106 were not covered by tests
}
}
}

Expand Down Expand Up @@ -116,6 +134,12 @@ pub fn handle_call_message(
return Err(res.err().unwrap());
}

let whitelisted = WHITELISTED_ACTIONS.load(deps.storage, data.clone());
if whitelisted.is_err() || !whitelisted.unwrap() {
return Err(ContractError::ActionNotWhitelisted);

Check warning on line 139 in contracts/core-contracts/cw-xcall-manager/src/contract.rs

View check run for this annotation

Codecov / codecov/patch

contracts/core-contracts/cw-xcall-manager/src/contract.rs#L139

Added line #L139 was not covered by tests
}
WHITELISTED_ACTIONS.remove(deps.storage, data.clone());

match method.as_str() {
CONFIGURE_PROTOCOLS => {
let configure_protocols: ConfigureProtocols = decode(&data).unwrap();
Expand Down
3 changes: 3 additions & 0 deletions contracts/core-contracts/cw-xcall-manager/src/error.rs
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,9 @@ pub enum ContractError {
#[error("Invalid protocols")]
InvalidProtocol,

#[error("Action is not whitelisted")]
ActionNotWhitelisted,

#[error("Rlp Error: {error}")]
DecoderError { error: DecoderError },
}
Expand Down
4 changes: 3 additions & 1 deletion contracts/core-contracts/cw-xcall-manager/src/state.rs
Original file line number Diff line number Diff line change
@@ -1,9 +1,11 @@
use cosmwasm_std::Addr;
use cw_common::xcall_manager_msg::ProtocolConfig;
use cw_storage_plus::Item;
use cw_storage_plus::{Item, Map};

pub const X_CALL: Item<String> = Item::new("xcall_address");
pub const ICON_GOVERNANCE: Item<String> = Item::new("icon_governance_network_address");
pub const PROPOSER: Item<Addr> = Item::new("admin_wallet");
pub const PROTOCOLS: Item<ProtocolConfig> = Item::new("protocols");
pub const PROPOSED_REMOVAL: Item<String> = Item::new("proposed_removal");

pub const WHITELISTED_ACTIONS: Map<Vec<u8>, bool> = Map::new("whitelisted_actions");
Loading

0 comments on commit cf2b929

Please sign in to comment.