From 0a2656cf9efe6297843b4f9d5ff1634497585235 Mon Sep 17 00:00:00 2001 From: Nicolas Lara Date: Mon, 15 Aug 2022 08:58:11 +0200 Subject: [PATCH] added ibc-rate-limiting contract --- x/ibc-rate-limit/.beaker/state.json | 1 + x/ibc-rate-limit/.gitignore | 21 + x/ibc-rate-limit/Cargo.toml | 16 + x/ibc-rate-limit/contracts/.gitkeep | 0 .../contracts/rate-limiter/.cargo/config | 3 + .../contracts/rate-limiter/.gitignore | 15 + .../contracts/rate-limiter/Cargo.toml | 53 ++ .../contracts/rate-limiter/src/contract.rs | 501 ++++++++++++++++++ .../contracts/rate-limiter/src/error.rs | 20 + .../contracts/rate-limiter/src/helpers.rs | 52 ++ .../rate-limiter/src/integration_tests.rs | 273 ++++++++++ .../contracts/rate-limiter/src/lib.rs | 9 + .../contracts/rate-limiter/src/management.rs | 238 +++++++++ .../contracts/rate-limiter/src/msg.rs | 69 +++ .../contracts/rate-limiter/src/state.rs | 149 ++++++ 15 files changed, 1420 insertions(+) create mode 100644 x/ibc-rate-limit/.beaker/state.json create mode 100644 x/ibc-rate-limit/.gitignore create mode 100644 x/ibc-rate-limit/Cargo.toml create mode 100644 x/ibc-rate-limit/contracts/.gitkeep create mode 100644 x/ibc-rate-limit/contracts/rate-limiter/.cargo/config create mode 100644 x/ibc-rate-limit/contracts/rate-limiter/.gitignore create mode 100644 x/ibc-rate-limit/contracts/rate-limiter/Cargo.toml create mode 100644 x/ibc-rate-limit/contracts/rate-limiter/src/contract.rs create mode 100644 x/ibc-rate-limit/contracts/rate-limiter/src/error.rs create mode 100644 x/ibc-rate-limit/contracts/rate-limiter/src/helpers.rs create mode 100644 x/ibc-rate-limit/contracts/rate-limiter/src/integration_tests.rs create mode 100644 x/ibc-rate-limit/contracts/rate-limiter/src/lib.rs create mode 100644 x/ibc-rate-limit/contracts/rate-limiter/src/management.rs create mode 100644 x/ibc-rate-limit/contracts/rate-limiter/src/msg.rs create mode 100644 x/ibc-rate-limit/contracts/rate-limiter/src/state.rs diff --git a/x/ibc-rate-limit/.beaker/state.json b/x/ibc-rate-limit/.beaker/state.json new file mode 100644 index 00000000000..9e26dfeeb6e --- /dev/null +++ b/x/ibc-rate-limit/.beaker/state.json @@ -0,0 +1 @@ +{} \ No newline at end of file diff --git a/x/ibc-rate-limit/.gitignore b/x/ibc-rate-limit/.gitignore new file mode 100644 index 00000000000..0814c1f8964 --- /dev/null +++ b/x/ibc-rate-limit/.gitignore @@ -0,0 +1,21 @@ +# Generated by Cargo +# will have compiled files and executables +debug/ +target/ + +# Generated by rust-optimizer +artifacts/ + +# Remove Cargo.lock from gitignore if creating an executable, leave it for libraries +# More information here https://doc.rust-lang.org/cargo/guide/cargo-toml-vs-cargo-lock.html +Cargo.lock + +# These are backup files generated by rustfmt +**/*.rs.bk + +# MSVC Windows builds of rustc generate these, which store debugging information +*.pdb + + +# Ignores local beaker state +**/state.local.json diff --git a/x/ibc-rate-limit/Cargo.toml b/x/ibc-rate-limit/Cargo.toml new file mode 100644 index 00000000000..9e4bf04d415 --- /dev/null +++ b/x/ibc-rate-limit/Cargo.toml @@ -0,0 +1,16 @@ +[workspace] + +members = [ + 'contracts/*', +] + +[profile.release] +codegen-units = 1 +debug = false +debug-assertions = false +incremental = false +lto = true +opt-level = 3 +overflow-checks = true +panic = 'abort' +rpath = false diff --git a/x/ibc-rate-limit/contracts/.gitkeep b/x/ibc-rate-limit/contracts/.gitkeep new file mode 100644 index 00000000000..e69de29bb2d diff --git a/x/ibc-rate-limit/contracts/rate-limiter/.cargo/config b/x/ibc-rate-limit/contracts/rate-limiter/.cargo/config new file mode 100644 index 00000000000..f31de6c2a75 --- /dev/null +++ b/x/ibc-rate-limit/contracts/rate-limiter/.cargo/config @@ -0,0 +1,3 @@ +[alias] +wasm = "build --release --target wasm32-unknown-unknown" +unit-test = "test --lib" diff --git a/x/ibc-rate-limit/contracts/rate-limiter/.gitignore b/x/ibc-rate-limit/contracts/rate-limiter/.gitignore new file mode 100644 index 00000000000..dfdaaa6bcda --- /dev/null +++ b/x/ibc-rate-limit/contracts/rate-limiter/.gitignore @@ -0,0 +1,15 @@ +# Build results +/target + +# Cargo+Git helper file (https://github.com/rust-lang/cargo/blob/0.44.1/src/cargo/sources/git/utils.rs#L320-L327) +.cargo-ok + +# Text file backups +**/*.rs.bk + +# macOS +.DS_Store + +# IDEs +*.iml +.idea diff --git a/x/ibc-rate-limit/contracts/rate-limiter/Cargo.toml b/x/ibc-rate-limit/contracts/rate-limiter/Cargo.toml new file mode 100644 index 00000000000..a94d596a72c --- /dev/null +++ b/x/ibc-rate-limit/contracts/rate-limiter/Cargo.toml @@ -0,0 +1,53 @@ +[package] +name = "rate-limiter" +version = "0.1.0" +authors = ["Nicolas Lara "] +edition = "2018" + +exclude = [ + # Those files are rust-optimizer artifacts. You might want to commit them for convenience but they should not be part of the source code publication. + "contract.wasm", + "hash.txt", +] + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[lib] +crate-type = ["cdylib", "rlib"] + +[profile.release] +opt-level = 3 +debug = false +rpath = false +lto = true +debug-assertions = false +codegen-units = 1 +panic = 'abort' +incremental = false +overflow-checks = true + +[features] +# for more explicit tests, cargo test --features=backtraces +backtraces = ["cosmwasm-std/backtraces"] +# use library feature to disable all instantiate/execute/query exports +library = [] + +[package.metadata.scripts] +optimize = """docker run --rm -v "$(pwd)":/code \ + --mount type=volume,source="$(basename "$(pwd)")_cache",target=/code/target \ + --mount type=volume,source=registry_cache,target=/usr/local/cargo/registry \ + cosmwasm/rust-optimizer:0.12.6 +""" + +[dependencies] +cosmwasm-std = "1.0.0" +cosmwasm-storage = "1.0.0" +cw-storage-plus = "0.13.2" +cw2 = "0.13.2" +schemars = "0.8.8" +serde = { version = "1.0.137", default-features = false, features = ["derive"] } +thiserror = { version = "1.0.31" } + +[dev-dependencies] +cosmwasm-schema = "1.0.0" +cw-multi-test = "0.13.2" diff --git a/x/ibc-rate-limit/contracts/rate-limiter/src/contract.rs b/x/ibc-rate-limit/contracts/rate-limiter/src/contract.rs new file mode 100644 index 00000000000..821fc2ad514 --- /dev/null +++ b/x/ibc-rate-limit/contracts/rate-limiter/src/contract.rs @@ -0,0 +1,501 @@ +#[cfg(not(feature = "library"))] +use cosmwasm_std::entry_point; +use cosmwasm_std::{ + to_binary, Addr, Binary, Deps, DepsMut, Env, MessageInfo, Response, StdResult, Timestamp, +}; +use cw2::set_contract_version; + +use crate::error::ContractError; +use crate::management::{ + add_new_channels, try_add_channel, try_remove_channel, try_reset_channel_quota, +}; +use crate::msg::{ExecuteMsg, InstantiateMsg, QueryMsg}; +use crate::state::{ChannelFlow, FlowType, CHANNEL_FLOWS, GOVMODULE, IBCMODULE}; + +// version info for migration info +const CONTRACT_NAME: &str = "crates.io:rate-limiter"; +const CONTRACT_VERSION: &str = env!("CARGO_PKG_VERSION"); + +#[cfg_attr(not(feature = "library"), entry_point)] +pub fn instantiate( + deps: DepsMut, + env: Env, + _info: MessageInfo, + msg: InstantiateMsg, +) -> Result { + set_contract_version(deps.storage, CONTRACT_NAME, CONTRACT_VERSION)?; + IBCMODULE.save(deps.storage, &msg.ibc_module)?; + GOVMODULE.save(deps.storage, &msg.gov_module)?; + + add_new_channels(deps, msg.channels, env.block.time)?; + + Ok(Response::new() + .add_attribute("method", "instantiate") + .add_attribute("ibc_module", msg.ibc_module.to_string()) + .add_attribute("gov_module", msg.gov_module.to_string())) +} + +#[cfg_attr(not(feature = "library"), entry_point)] +pub fn execute( + deps: DepsMut, + env: Env, + info: MessageInfo, + msg: ExecuteMsg, +) -> Result { + match msg { + ExecuteMsg::SendPacket { + channel_id, + channel_value, + funds, + } => try_transfer( + deps, + info.sender, + channel_id, + channel_value, + funds, + FlowType::Out, + env.block.time, + ), + ExecuteMsg::RecvPacket { + channel_id, + channel_value, + funds, + } => try_transfer( + deps, + info.sender, + channel_id, + channel_value, + funds, + FlowType::In, + env.block.time, + ), + ExecuteMsg::AddChannel { channel_id, quotas } => { + try_add_channel(deps, info.sender, channel_id, quotas, env.block.time) + } + ExecuteMsg::RemoveChannel { channel_id } => { + try_remove_channel(deps, info.sender, channel_id) + } + ExecuteMsg::ResetChannelQuota { + channel_id, + quota_id, + } => try_reset_channel_quota(deps, info.sender, channel_id, quota_id, env.block.time), + } +} + +pub struct ChannelFlowResponse { + pub channel_flow: ChannelFlow, + pub used: u128, + pub max: u128, +} + +pub fn try_transfer( + deps: DepsMut, + sender: Addr, + channel_id: String, + channel_value: u128, + funds: u128, + direction: FlowType, + now: Timestamp, +) -> Result { + // Only the IBCMODULE can execute transfers + let ibc_module = IBCMODULE.load(deps.storage)?; + if sender != ibc_module { + return Err(ContractError::Unauthorized {}); + } + + let channels = CHANNEL_FLOWS.may_load(deps.storage, &channel_id)?; + + let configured = match channels { + None => false, + Some(ref x) if x.len() == 0 => false, + _ => true, + }; + + if !configured { + // No Quota configured for the current channel. Allowing all messages. + return Ok(Response::new() + .add_attribute("method", "try_transfer") + .add_attribute("channel_id", channel_id) + .add_attribute("quota", "none")); + } + + let mut channels = channels.unwrap(); + + let results: Result, _> = channels + .iter_mut() + .map(|channel| { + let max = channel.quota.capacity_at(&channel_value, &direction); + if channel.flow.is_expired(now) { + channel.flow.expire(now, channel.quota.duration) + } + channel.flow.add_flow(direction.clone(), funds); + + let balance = channel.flow.balance(); + if balance > max { + return Err(ContractError::RateLimitExceded { + channel: channel_id.to_string(), + reset: channel.flow.period_end, + }); + }; + Ok(ChannelFlowResponse { + channel_flow: ChannelFlow { + quota: channel.quota.clone(), + flow: channel.flow.clone(), + }, + used: balance, + max, + }) + }) + .collect(); + let results = results?; + + CHANNEL_FLOWS.save( + deps.storage, + &channel_id, + &results.iter().map(|r| r.channel_flow.clone()).collect(), + )?; + + let response = Response::new() + .add_attribute("method", "try_transfer") + .add_attribute("channel_id", channel_id); + + // Adding the attributes from each quota to the response + results.iter().fold(Ok(response), |acc, result| { + Ok(acc? + .add_attribute( + format!("{}_used", result.channel_flow.quota.name), + result.used.to_string(), + ) + .add_attribute( + format!("{}_max", result.channel_flow.quota.name), + result.max.to_string(), + ) + .add_attribute( + format!("{}_period_end", result.channel_flow.quota.name), + result.channel_flow.flow.period_end.to_string(), + )) + }) +} + +#[cfg_attr(not(feature = "library"), entry_point)] +pub fn query(deps: Deps, _env: Env, msg: QueryMsg) -> StdResult { + match msg { + QueryMsg::GetQuotas { channel_id } => get_quotas(deps, channel_id), + } +} + +fn get_quotas(deps: Deps, channel_id: impl Into) -> StdResult { + to_binary(&CHANNEL_FLOWS.load(deps.storage, &channel_id.into())?) +} + +#[cfg(test)] +mod tests { + use super::*; + use cosmwasm_std::testing::{mock_dependencies, mock_env, mock_info}; + use cosmwasm_std::{from_binary, Addr, Attribute}; + + use crate::helpers::tests::verify_query_response; + use crate::msg::{Channel, QuotaMsg}; + use crate::state::RESET_TIME_WEEKLY; + + const IBC_ADDR: &str = "IBC_MODULE"; + const GOV_ADDR: &str = "GOV_MODULE"; + + #[test] + fn proper_instantiation() { + let mut deps = mock_dependencies(); + + let msg = InstantiateMsg { + gov_module: Addr::unchecked(GOV_ADDR), + ibc_module: Addr::unchecked(IBC_ADDR), + channels: vec![], + }; + let info = mock_info(IBC_ADDR, &vec![]); + + // we can just call .unwrap() to assert this was a success + let res = instantiate(deps.as_mut(), mock_env(), info, msg).unwrap(); + assert_eq!(0, res.messages.len()); + + // The ibc and gov modules are properly stored + assert_eq!(IBCMODULE.load(deps.as_ref().storage).unwrap(), IBC_ADDR); + assert_eq!(GOVMODULE.load(deps.as_ref().storage).unwrap(), GOV_ADDR); + } + + #[test] + fn permissions() { + let mut deps = mock_dependencies(); + + let quota = QuotaMsg::new("Weekly", RESET_TIME_WEEKLY, 10, 10); + let msg = InstantiateMsg { + gov_module: Addr::unchecked(GOV_ADDR), + ibc_module: Addr::unchecked(IBC_ADDR), + channels: vec![Channel { + name: "channel".to_string(), + quotas: vec![quota], + }], + }; + let info = mock_info(IBC_ADDR, &vec![]); + instantiate(deps.as_mut(), mock_env(), info.clone(), msg).unwrap(); + + let msg = ExecuteMsg::SendPacket { + channel_id: "channel".to_string(), + channel_value: 3_000, + funds: 300, + }; + + // This succeeds + execute(deps.as_mut(), mock_env(), info.clone(), msg).unwrap(); + + let info = mock_info("SomeoneElse", &vec![]); + + let msg = ExecuteMsg::SendPacket { + channel_id: "channel".to_string(), + channel_value: 3_000, + funds: 300, + }; + let err = execute(deps.as_mut(), mock_env(), info.clone(), msg).unwrap_err(); + assert!(matches!(err, ContractError::Unauthorized { .. })); + } + + #[test] + fn consume_allowance() { + let mut deps = mock_dependencies(); + + let quota = QuotaMsg::new("weekly", RESET_TIME_WEEKLY, 10, 10); + let msg = InstantiateMsg { + gov_module: Addr::unchecked(GOV_ADDR), + ibc_module: Addr::unchecked(IBC_ADDR), + channels: vec![Channel { + name: "channel".to_string(), + quotas: vec![quota], + }], + }; + let info = mock_info(GOV_ADDR, &vec![]); + let _res = instantiate(deps.as_mut(), mock_env(), info.clone(), msg).unwrap(); + + let msg = ExecuteMsg::SendPacket { + channel_id: "channel".to_string(), + channel_value: 3_000, + funds: 300, + }; + let info = mock_info(IBC_ADDR, &vec![]); + let res = execute(deps.as_mut(), mock_env(), info.clone(), msg).unwrap(); + let Attribute { key, value } = &res.attributes[2]; + assert_eq!(key, "weekly_used"); + assert_eq!(value, "300"); + + let msg = ExecuteMsg::SendPacket { + channel_id: "channel".to_string(), + channel_value: 3_000, + funds: 300, + }; + let err = execute(deps.as_mut(), mock_env(), info, msg).unwrap_err(); + assert!(matches!(err, ContractError::RateLimitExceded { .. })); + } + + #[test] + fn symetric_flows_dont_consume_allowance() { + let mut deps = mock_dependencies(); + + let quota = QuotaMsg::new("weekly", RESET_TIME_WEEKLY, 10, 10); + let msg = InstantiateMsg { + gov_module: Addr::unchecked(GOV_ADDR), + ibc_module: Addr::unchecked(IBC_ADDR), + channels: vec![Channel { + name: "channel".to_string(), + quotas: vec![quota], + }], + }; + let info = mock_info(GOV_ADDR, &vec![]); + let _res = instantiate(deps.as_mut(), mock_env(), info.clone(), msg).unwrap(); + + let info = mock_info(IBC_ADDR, &vec![]); + let send_msg = ExecuteMsg::SendPacket { + channel_id: "channel".to_string(), + channel_value: 3_000, + funds: 300, + }; + let recv_msg = ExecuteMsg::RecvPacket { + channel_id: "channel".to_string(), + channel_value: 3_000, + funds: 300, + }; + + let res = execute(deps.as_mut(), mock_env(), info.clone(), send_msg.clone()).unwrap(); + let Attribute { key, value } = &res.attributes[2]; + assert_eq!(key, "weekly_used"); + assert_eq!(value, "300"); + + let res = execute(deps.as_mut(), mock_env(), info.clone(), recv_msg.clone()).unwrap(); + let Attribute { key, value } = &res.attributes[2]; + assert_eq!(key, "weekly_used"); + assert_eq!(value, "0"); + + // We can still use the channel. Even if we have sent more than the + // allowance through the channel (900 > 3000*.1), the current "balance" + // of inflow vs outflow is still lower than the channel's capacity/quota + let res = execute(deps.as_mut(), mock_env(), info.clone(), recv_msg.clone()).unwrap(); + let Attribute { key, value } = &res.attributes[2]; + assert_eq!(key, "weekly_used"); + assert_eq!(value, "300"); + + let err = execute(deps.as_mut(), mock_env(), info.clone(), recv_msg.clone()).unwrap_err(); + + assert!(matches!(err, ContractError::RateLimitExceded { .. })); + //assert_eq!(18, value.count); + } + + #[test] + fn asymetric_quotas() { + let mut deps = mock_dependencies(); + + let quota = QuotaMsg::new("weekly", RESET_TIME_WEEKLY, 10, 1); + let msg = InstantiateMsg { + gov_module: Addr::unchecked(GOV_ADDR), + ibc_module: Addr::unchecked(IBC_ADDR), + channels: vec![Channel { + name: "channel".to_string(), + quotas: vec![quota], + }], + }; + let info = mock_info(GOV_ADDR, &vec![]); + let _res = instantiate(deps.as_mut(), mock_env(), info.clone(), msg).unwrap(); + + // Sending 2% + let msg = ExecuteMsg::SendPacket { + channel_id: "channel".to_string(), + channel_value: 3_000, + funds: 60, + }; + let info = mock_info(IBC_ADDR, &vec![]); + let res = execute(deps.as_mut(), mock_env(), info.clone(), msg).unwrap(); + let Attribute { key, value } = &res.attributes[2]; + assert_eq!(key, "weekly_used"); + assert_eq!(value, "60"); + + // Sending 1% more. Allowed, as sending has a 10% allowance + let msg = ExecuteMsg::SendPacket { + channel_id: "channel".to_string(), + channel_value: 3_000, + funds: 30, + }; + + let info = mock_info(IBC_ADDR, &vec![]); + let res = execute(deps.as_mut(), mock_env(), info.clone(), msg).unwrap(); + let Attribute { key, value } = &res.attributes[2]; + assert_eq!(key, "weekly_used"); + assert_eq!(value, "90"); + + // Receiving 1% should fail. 3% already executed through the channel + let recv_msg = ExecuteMsg::RecvPacket { + channel_id: "channel".to_string(), + channel_value: 3_000, + funds: 30, + }; + + let err = execute(deps.as_mut(), mock_env(), info.clone(), recv_msg.clone()).unwrap_err(); + assert!(matches!(err, ContractError::RateLimitExceded { .. })); + } + + #[test] + fn query_state() { + let mut deps = mock_dependencies(); + + let quota = QuotaMsg::new("weekly", RESET_TIME_WEEKLY, 10, 10); + let msg = InstantiateMsg { + gov_module: Addr::unchecked(GOV_ADDR), + ibc_module: Addr::unchecked(IBC_ADDR), + channels: vec![Channel { + name: "channel".to_string(), + quotas: vec![quota], + }], + }; + let info = mock_info(GOV_ADDR, &vec![]); + let env = mock_env(); + let _res = instantiate(deps.as_mut(), env.clone(), info, msg).unwrap(); + + let query_msg = QueryMsg::GetQuotas { + channel_id: "channel".to_string(), + }; + + let res = query(deps.as_ref(), mock_env(), query_msg.clone()).unwrap(); + let value: Vec = from_binary(&res).unwrap(); + assert_eq!(value[0].quota.name, "weekly"); + assert_eq!(value[0].quota.max_percentage_send, 10); + assert_eq!(value[0].quota.max_percentage_recv, 10); + assert_eq!(value[0].quota.duration, RESET_TIME_WEEKLY); + assert_eq!(value[0].flow.inflow, 0); + assert_eq!(value[0].flow.outflow, 0); + assert_eq!( + value[0].flow.period_end, + env.block.time.plus_seconds(RESET_TIME_WEEKLY) + ); + + let info = mock_info(IBC_ADDR, &vec![]); + let send_msg = ExecuteMsg::SendPacket { + channel_id: "channel".to_string(), + channel_value: 3_000, + funds: 300, + }; + execute(deps.as_mut(), mock_env(), info.clone(), send_msg.clone()).unwrap(); + + let recv_msg = ExecuteMsg::RecvPacket { + channel_id: "channel".to_string(), + channel_value: 3_000, + funds: 30, + }; + execute(deps.as_mut(), mock_env(), info, recv_msg.clone()).unwrap(); + + // Query + let res = query(deps.as_ref(), mock_env(), query_msg.clone()).unwrap(); + let value: Vec = from_binary(&res).unwrap(); + verify_query_response( + &value[0], + "weekly", + (10, 10), + RESET_TIME_WEEKLY, + 30, + 300, + env.block.time.plus_seconds(RESET_TIME_WEEKLY), + ); + } + + #[test] + fn bad_quotas() { + let mut deps = mock_dependencies(); + + let msg = InstantiateMsg { + gov_module: Addr::unchecked(GOV_ADDR), + ibc_module: Addr::unchecked(IBC_ADDR), + channels: vec![Channel { + name: "channel".to_string(), + quotas: vec![QuotaMsg { + name: "bad_quota".to_string(), + duration: 200, + send_recv: (5000, 101), + }], + }], + }; + let info = mock_info(IBC_ADDR, &vec![]); + + // we can just call .unwrap() to assert this was a success + let env = mock_env(); + instantiate(deps.as_mut(), env.clone(), info, msg).unwrap(); + + // If a quota is higher than 100%, we set it to 100% + let query_msg = QueryMsg::GetQuotas { + channel_id: "channel".to_string(), + }; + let res = query(deps.as_ref(), env.clone(), query_msg).unwrap(); + let value: Vec = from_binary(&res).unwrap(); + verify_query_response( + &value[0], + "bad_quota", + (100, 100), + 200, + 0, + 0, + env.block.time.plus_seconds(200), + ); + } +} diff --git a/x/ibc-rate-limit/contracts/rate-limiter/src/error.rs b/x/ibc-rate-limit/contracts/rate-limiter/src/error.rs new file mode 100644 index 00000000000..3c58ce21bc9 --- /dev/null +++ b/x/ibc-rate-limit/contracts/rate-limiter/src/error.rs @@ -0,0 +1,20 @@ +use cosmwasm_std::{StdError, Timestamp}; +use thiserror::Error; + +#[derive(Error, Debug)] +pub enum ContractError { + #[error("{0}")] + Std(#[from] StdError), + + #[error("Unauthorized")] + Unauthorized {}, + + #[error("IBC Rate Limit exceded for channel {channel:?}. Try again after {reset:?}")] + RateLimitExceded { channel: String, reset: Timestamp }, + + #[error("Quota {quota_id} not found for channel {channel_id}")] + QuotaNotFound { + quota_id: String, + channel_id: String, + }, +} diff --git a/x/ibc-rate-limit/contracts/rate-limiter/src/helpers.rs b/x/ibc-rate-limit/contracts/rate-limiter/src/helpers.rs new file mode 100644 index 00000000000..82f4f1168ba --- /dev/null +++ b/x/ibc-rate-limit/contracts/rate-limiter/src/helpers.rs @@ -0,0 +1,52 @@ +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; + +use cosmwasm_std::{to_binary, Addr, CosmosMsg, StdResult, WasmMsg}; + +use crate::msg::ExecuteMsg; + +/// CwTemplateContract is a wrapper around Addr that provides a lot of helpers +/// for working with this. +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema)] +pub struct RateLimitingContract(pub Addr); + +impl RateLimitingContract { + pub fn addr(&self) -> Addr { + self.0.clone() + } + + pub fn call>(&self, msg: T) -> StdResult { + let msg = to_binary(&msg.into())?; + Ok(WasmMsg::Execute { + contract_addr: self.addr().into(), + msg, + funds: vec![], + } + .into()) + } +} + +#[cfg(test)] +pub mod tests { + use cosmwasm_std::Timestamp; + + use crate::state::ChannelFlow; + + pub fn verify_query_response( + value: &ChannelFlow, + quota_name: &str, + send_recv: (u32, u32), + duration: u64, + inflow: u128, + outflow: u128, + period_end: Timestamp, + ) { + assert_eq!(value.quota.name, quota_name); + assert_eq!(value.quota.max_percentage_send, send_recv.0); + assert_eq!(value.quota.max_percentage_recv, send_recv.1); + assert_eq!(value.quota.duration, duration); + assert_eq!(value.flow.inflow, inflow); + assert_eq!(value.flow.outflow, outflow); + assert_eq!(value.flow.period_end, period_end); + } +} diff --git a/x/ibc-rate-limit/contracts/rate-limiter/src/integration_tests.rs b/x/ibc-rate-limit/contracts/rate-limiter/src/integration_tests.rs new file mode 100644 index 00000000000..26f48cea159 --- /dev/null +++ b/x/ibc-rate-limit/contracts/rate-limiter/src/integration_tests.rs @@ -0,0 +1,273 @@ +#[cfg(test)] +mod tests { + use crate::helpers::RateLimitingContract; + use crate::msg::{Channel, InstantiateMsg}; + use cosmwasm_std::{Addr, Coin, Empty, Uint128}; + use cw_multi_test::{App, AppBuilder, Contract, ContractWrapper, Executor}; + + pub fn contract_template() -> Box> { + let contract = ContractWrapper::new( + crate::contract::execute, + crate::contract::instantiate, + crate::contract::query, + ); + Box::new(contract) + } + + const USER: &str = "USER"; + const IBC_ADDR: &str = "IBC_MODULE"; + const GOV_ADDR: &str = "GOV_MODULE"; + const NATIVE_DENOM: &str = "nosmo"; + + fn mock_app() -> App { + AppBuilder::new().build(|router, _, storage| { + router + .bank + .init_balance( + storage, + &Addr::unchecked(USER), + vec![Coin { + denom: NATIVE_DENOM.to_string(), + amount: Uint128::new(1_000), + }], + ) + .unwrap(); + }) + } + + fn proper_instantiate(channels: Vec) -> (App, RateLimitingContract) { + let mut app = mock_app(); + let cw_template_id = app.store_code(contract_template()); + + let msg = InstantiateMsg { + gov_module: Addr::unchecked(GOV_ADDR), + ibc_module: Addr::unchecked(IBC_ADDR), + channels, + }; + + let cw_template_contract_addr = app + .instantiate_contract( + cw_template_id, + Addr::unchecked(GOV_ADDR), + &msg, + &[], + "test", + None, + ) + .unwrap(); + + let cw_template_contract = RateLimitingContract(cw_template_contract_addr); + + (app, cw_template_contract) + } + + mod expiration { + use cosmwasm_std::Attribute; + + use super::*; + use crate::{ + msg::{Channel, ExecuteMsg, QuotaMsg}, + state::{RESET_TIME_DAILY, RESET_TIME_MONTHLY, RESET_TIME_WEEKLY}, + }; + + #[test] + fn expiration() { + let quota = QuotaMsg::new("weekly", RESET_TIME_WEEKLY, 10, 10); + + let (mut app, cw_template_contract) = proper_instantiate(vec![Channel { + name: "channel".to_string(), + quotas: vec![quota], + }]); + + // Using all the allowance + let msg = ExecuteMsg::SendPacket { + channel_id: "channel".to_string(), + channel_value: 3_000, + funds: 300, + }; + let cosmos_msg = cw_template_contract.call(msg).unwrap(); + let res = app.execute(Addr::unchecked(IBC_ADDR), cosmos_msg).unwrap(); + + let Attribute { key, value } = &res.custom_attrs(1)[2]; + assert_eq!(key, "weekly_used"); + assert_eq!(value, "300"); + let Attribute { key, value } = &res.custom_attrs(1)[3]; + assert_eq!(key, "weekly_max"); + assert_eq!(value, "300"); + + // Another packet is rate limited + let msg = ExecuteMsg::SendPacket { + channel_id: "channel".to_string(), + channel_value: 3_000, + funds: 300, + }; + let cosmos_msg = cw_template_contract.call(msg).unwrap(); + let err = app + .execute(Addr::unchecked(IBC_ADDR), cosmos_msg) + .unwrap_err(); + + // TODO: how do we check the error type here? + println!("{err:?}"); + + // ... Time passes + app.update_block(|b| { + b.height += 1000; + b.time = b.time.plus_seconds(RESET_TIME_WEEKLY + 1) + }); + + // Sending the packet should work now + let msg = ExecuteMsg::SendPacket { + channel_id: "channel".to_string(), + channel_value: 3_000, + funds: 300, + }; + + let cosmos_msg = cw_template_contract.call(msg).unwrap(); + let res = app.execute(Addr::unchecked(IBC_ADDR), cosmos_msg).unwrap(); + + let Attribute { key, value } = &res.custom_attrs(1)[2]; + assert_eq!(key, "weekly_used"); + assert_eq!(value, "300"); + let Attribute { key, value } = &res.custom_attrs(1)[3]; + assert_eq!(key, "weekly_max"); + assert_eq!(value, "300"); + } + + #[test] + fn multiple_quotas() { + let quotas = vec![ + QuotaMsg::new("daily", RESET_TIME_DAILY, 1, 1), + QuotaMsg::new("weekly", RESET_TIME_WEEKLY, 5, 5), + QuotaMsg::new("monthly", RESET_TIME_MONTHLY, 5, 5), + ]; + + let (mut app, cw_template_contract) = proper_instantiate(vec![Channel { + name: "channel".to_string(), + quotas, + }]); + + // Sending 1% to use the daily allowance + let msg = ExecuteMsg::SendPacket { + channel_id: "channel".to_string(), + channel_value: 100, + funds: 1, + }; + let cosmos_msg = cw_template_contract.call(msg).unwrap(); + let res = app.execute(Addr::unchecked(IBC_ADDR), cosmos_msg).unwrap(); + + println!("{res:?}"); + + // Another packet is rate limited + let msg = ExecuteMsg::SendPacket { + channel_id: "channel".to_string(), + channel_value: 100, + funds: 1, + }; + let cosmos_msg = cw_template_contract.call(msg).unwrap(); + let _err = app + .execute(Addr::unchecked(IBC_ADDR), cosmos_msg) + .unwrap_err(); + + // ... One day passes + app.update_block(|b| { + b.height += 10; + b.time = b.time.plus_seconds(RESET_TIME_DAILY + 1) + }); + + // Sending the packet should work now + let msg = ExecuteMsg::SendPacket { + channel_id: "channel".to_string(), + channel_value: 100, + funds: 1, + }; + + let cosmos_msg = cw_template_contract.call(msg).unwrap(); + app.execute(Addr::unchecked(IBC_ADDR), cosmos_msg).unwrap(); + + // Do that for 4 more days + for _ in 1..4 { + // ... One day passes + app.update_block(|b| { + b.height += 10; + b.time = b.time.plus_seconds(RESET_TIME_DAILY + 1) + }); + + // Sending the packet should work now + let msg = ExecuteMsg::SendPacket { + channel_id: "channel".to_string(), + channel_value: 100, + funds: 1, + }; + let cosmos_msg = cw_template_contract.call(msg).unwrap(); + app.execute(Addr::unchecked(IBC_ADDR), cosmos_msg).unwrap(); + } + + // ... One day passes + app.update_block(|b| { + b.height += 10; + b.time = b.time.plus_seconds(RESET_TIME_DAILY + 1) + }); + + // We now have exceeded the weekly limit! Even if the daily limit allows us, the weekly doesn't + let msg = ExecuteMsg::SendPacket { + channel_id: "channel".to_string(), + channel_value: 100, + funds: 1, + }; + let cosmos_msg = cw_template_contract.call(msg).unwrap(); + let _err = app + .execute(Addr::unchecked(IBC_ADDR), cosmos_msg) + .unwrap_err(); + + // ... One week passes + app.update_block(|b| { + b.height += 10; + b.time = b.time.plus_seconds(RESET_TIME_WEEKLY + 1) + }); + + // We can still can't send because the weekly and monthly limits are the same + let msg = ExecuteMsg::SendPacket { + channel_id: "channel".to_string(), + channel_value: 100, + funds: 1, + }; + let cosmos_msg = cw_template_contract.call(msg).unwrap(); + let _err = app + .execute(Addr::unchecked(IBC_ADDR), cosmos_msg) + .unwrap_err(); + + // Waiting a week again, doesn't help!! + // ... One week passes + app.update_block(|b| { + b.height += 10; + b.time = b.time.plus_seconds(RESET_TIME_WEEKLY + 1) + }); + + // We can still can't send because the monthly limit hasn't passed + let msg = ExecuteMsg::SendPacket { + channel_id: "channel".to_string(), + channel_value: 100, + funds: 1, + }; + let cosmos_msg = cw_template_contract.call(msg).unwrap(); + let _err = app + .execute(Addr::unchecked(IBC_ADDR), cosmos_msg) + .unwrap_err(); + + // Only after two more weeks we can send again + app.update_block(|b| { + b.height += 10; + b.time = b.time.plus_seconds((RESET_TIME_WEEKLY * 2) + 1) // Two weeks + }); + + println!("{:?}", app.block_info()); + let msg = ExecuteMsg::SendPacket { + channel_id: "channel".to_string(), + channel_value: 100, + funds: 1, + }; + let cosmos_msg = cw_template_contract.call(msg).unwrap(); + let _err = app.execute(Addr::unchecked(IBC_ADDR), cosmos_msg).unwrap(); + } + } +} diff --git a/x/ibc-rate-limit/contracts/rate-limiter/src/lib.rs b/x/ibc-rate-limit/contracts/rate-limiter/src/lib.rs new file mode 100644 index 00000000000..573d76512d7 --- /dev/null +++ b/x/ibc-rate-limit/contracts/rate-limiter/src/lib.rs @@ -0,0 +1,9 @@ +pub mod contract; +mod error; +pub mod helpers; +pub mod integration_tests; +pub mod management; +pub mod msg; +pub mod state; + +pub use crate::error::ContractError; diff --git a/x/ibc-rate-limit/contracts/rate-limiter/src/management.rs b/x/ibc-rate-limit/contracts/rate-limiter/src/management.rs new file mode 100644 index 00000000000..b1ba33ce35b --- /dev/null +++ b/x/ibc-rate-limit/contracts/rate-limiter/src/management.rs @@ -0,0 +1,238 @@ +use crate::msg::{Channel, QuotaMsg}; +use crate::state::{ChannelFlow, Flow, CHANNEL_FLOWS, GOVMODULE, IBCMODULE}; +use crate::ContractError; +use cosmwasm_std::{Addr, DepsMut, Response, Timestamp}; + +pub fn add_new_channels( + deps: DepsMut, + channels: Vec, + now: Timestamp, +) -> Result<(), ContractError> { + for channel in channels { + CHANNEL_FLOWS.save( + deps.storage, + &channel.name, + &channel + .quotas + .iter() + .map(|q| ChannelFlow { + quota: q.into(), + flow: Flow::new(0_u128, 0_u128, now, q.duration), + }) + .collect(), + )? + } + Ok(()) +} + +pub fn try_add_channel( + deps: DepsMut, + sender: Addr, + channel_id: String, + quotas: Vec, + now: Timestamp, +) -> Result { + let ibc_module = IBCMODULE.load(deps.storage)?; + let gov_module = GOVMODULE.load(deps.storage)?; + if sender != ibc_module && sender != gov_module { + return Err(ContractError::Unauthorized {}); + } + add_new_channels( + deps, + vec![Channel { + name: channel_id.to_string(), + quotas, + }], + now, + )?; + + Ok(Response::new() + .add_attribute("method", "try_add_channel") + .add_attribute("channel_id", channel_id)) +} + +pub fn try_remove_channel( + deps: DepsMut, + sender: Addr, + channel_id: String, +) -> Result { + let ibc_module = IBCMODULE.load(deps.storage)?; + let gov_module = GOVMODULE.load(deps.storage)?; + if sender != ibc_module && sender != gov_module { + return Err(ContractError::Unauthorized {}); + } + CHANNEL_FLOWS.remove(deps.storage, &channel_id); + Ok(Response::new() + .add_attribute("method", "try_remove_channel") + .add_attribute("channel_id", channel_id)) +} + +pub fn try_reset_channel_quota( + deps: DepsMut, + sender: Addr, + channel_id: String, + quota_id: String, + now: Timestamp, +) -> Result { + let gov_module = GOVMODULE.load(deps.storage)?; + if sender != gov_module { + return Err(ContractError::Unauthorized {}); + } + + CHANNEL_FLOWS.update( + deps.storage, + &channel_id.clone(), + |maybe_flows| match maybe_flows { + None => Err(ContractError::QuotaNotFound { + quota_id, + channel_id: channel_id.clone(), + }), + Some(mut flows) => { + flows.iter_mut().for_each(|channel| { + if channel.quota.name == channel_id.as_ref() { + channel.flow.expire(now, channel.quota.duration) + } + }); + Ok(flows) + } + }, + )?; + + Ok(Response::new() + .add_attribute("method", "try_reset_channel") + .add_attribute("channel_id", channel_id)) +} + +#[cfg(test)] +mod tests { + + use cosmwasm_std::testing::{mock_dependencies, mock_env, mock_info}; + use cosmwasm_std::{from_binary, Addr, StdError}; + + use crate::contract::{execute, query}; + use crate::helpers::tests::verify_query_response; + use crate::msg::{ExecuteMsg, QueryMsg, QuotaMsg}; + use crate::state::{ChannelFlow, GOVMODULE, IBCMODULE}; + + const IBC_ADDR: &str = "IBC_MODULE"; + const GOV_ADDR: &str = "GOV_MODULE"; + + #[test] + fn management_add_and_remove_channel() { + let mut deps = mock_dependencies(); + IBCMODULE + .save(deps.as_mut().storage, &Addr::unchecked(IBC_ADDR)) + .unwrap(); + GOVMODULE + .save(deps.as_mut().storage, &Addr::unchecked(GOV_ADDR)) + .unwrap(); + + let msg = ExecuteMsg::AddChannel { + channel_id: "channel".to_string(), + quotas: vec![QuotaMsg { + name: "daily".to_string(), + duration: 1600, + send_recv: (3, 5), + }], + }; + let info = mock_info(IBC_ADDR, &vec![]); + + let env = mock_env(); + let res = execute(deps.as_mut(), env.clone(), info, msg).unwrap(); + assert_eq!(0, res.messages.len()); + + let query_msg = QueryMsg::GetQuotas { + channel_id: "channel".to_string(), + }; + + let res = query(deps.as_ref(), mock_env(), query_msg.clone()).unwrap(); + + let value: Vec = from_binary(&res).unwrap(); + verify_query_response( + &value[0], + "daily", + (3, 5), + 1600, + 0, + 0, + env.block.time.plus_seconds(1600), + ); + + assert_eq!(value.len(), 1); + + // Add another channel + let msg = ExecuteMsg::AddChannel { + channel_id: "channel2".to_string(), + quotas: vec![QuotaMsg { + name: "daily".to_string(), + duration: 1600, + send_recv: (3, 5), + }], + }; + let info = mock_info(IBC_ADDR, &vec![]); + + let env = mock_env(); + execute(deps.as_mut(), env.clone(), info, msg).unwrap(); + + // remove the first one + let msg = ExecuteMsg::RemoveChannel { + channel_id: "channel".to_string(), + }; + + let info = mock_info(IBC_ADDR, &vec![]); + let env = mock_env(); + execute(deps.as_mut(), env.clone(), info, msg).unwrap(); + + // The channel is not there anymore + let err = query(deps.as_ref(), mock_env(), query_msg.clone()).unwrap_err(); + assert!(matches!(err, StdError::NotFound { .. })); + + // The second channel is still there + let query_msg = QueryMsg::GetQuotas { + channel_id: "channel2".to_string(), + }; + let res = query(deps.as_ref(), mock_env(), query_msg.clone()).unwrap(); + let value: Vec = from_binary(&res).unwrap(); + assert_eq!(value.len(), 1); + verify_query_response( + &value[0], + "daily", + (3, 5), + 1600, + 0, + 0, + env.block.time.plus_seconds(1600), + ); + + // Channels are overriden if they share a name + let msg = ExecuteMsg::AddChannel { + channel_id: "channel2".to_string(), + quotas: vec![QuotaMsg { + name: "different".to_string(), + duration: 5000, + send_recv: (50, 30), + }], + }; + let info = mock_info(IBC_ADDR, &vec![]); + + let env = mock_env(); + execute(deps.as_mut(), env.clone(), info, msg).unwrap(); + + let query_msg = QueryMsg::GetQuotas { + channel_id: "channel2".to_string(), + }; + let res = query(deps.as_ref(), mock_env(), query_msg.clone()).unwrap(); + let value: Vec = from_binary(&res).unwrap(); + assert_eq!(value.len(), 1); + println!("{value:?}"); + verify_query_response( + &value[0], + "different", + (50, 30), + 5000, + 0, + 0, + env.block.time.plus_seconds(5000), + ); + } +} diff --git a/x/ibc-rate-limit/contracts/rate-limiter/src/msg.rs b/x/ibc-rate-limit/contracts/rate-limiter/src/msg.rs new file mode 100644 index 00000000000..35f2812db57 --- /dev/null +++ b/x/ibc-rate-limit/contracts/rate-limiter/src/msg.rs @@ -0,0 +1,69 @@ +use cosmwasm_std::Addr; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema)] +pub struct Channel { + pub name: String, + pub quotas: Vec, +} + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema)] +pub struct QuotaMsg { + pub name: String, + pub duration: u64, + pub send_recv: (u32, u32), +} + +impl QuotaMsg { + pub fn new(name: &str, seconds: u64, send_percentage: u32, recv_percentage: u32) -> Self { + QuotaMsg { + name: name.to_string(), + duration: seconds, + send_recv: (send_percentage, recv_percentage), + } + } +} + +/// Initialize the contract with the address of the IBC module and any existing channels. +/// Only the ibc module is allowed to execute actions on this contract +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema)] +pub struct InstantiateMsg { + pub gov_module: Addr, + pub ibc_module: Addr, + pub channels: Vec, +} + +/// The caller (IBC module) is responsibble for correctly calculating the funds +/// being sent through the channel +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema)] +#[serde(rename_all = "snake_case")] +pub enum ExecuteMsg { + SendPacket { + channel_id: String, + channel_value: u128, + funds: u128, + }, + RecvPacket { + channel_id: String, + channel_value: u128, + funds: u128, + }, + AddChannel { + channel_id: String, + quotas: Vec, + }, + RemoveChannel { + channel_id: String, + }, + ResetChannelQuota { + channel_id: String, + quota_id: String, + }, +} + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema)] +#[serde(rename_all = "snake_case")] +pub enum QueryMsg { + GetQuotas { channel_id: String }, +} diff --git a/x/ibc-rate-limit/contracts/rate-limiter/src/state.rs b/x/ibc-rate-limit/contracts/rate-limiter/src/state.rs new file mode 100644 index 00000000000..adaa323aad6 --- /dev/null +++ b/x/ibc-rate-limit/contracts/rate-limiter/src/state.rs @@ -0,0 +1,149 @@ +use cosmwasm_std::{Addr, Timestamp}; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; +use std::cmp; + +use cw_storage_plus::{Item, Map}; + +use crate::msg::QuotaMsg; + +pub const RESET_TIME_DAILY: u64 = 60 * 60 * 24; +pub const RESET_TIME_WEEKLY: u64 = 60 * 60 * 24 * 7; +pub const RESET_TIME_MONTHLY: u64 = 60 * 60 * 24 * 30; + +#[derive(Debug, Clone)] +pub enum FlowType { + In, + Out, +} + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema, Copy)] +pub struct Flow { + pub inflow: u128, + pub outflow: u128, + pub period_end: Timestamp, +} + +impl Flow { + pub fn new( + inflow: impl Into, + outflow: impl Into, + now: Timestamp, + duration: u64, + ) -> Self { + Self { + inflow: inflow.into(), + outflow: outflow.into(), + period_end: now.plus_seconds(duration), + } + } + + pub fn balance(&self) -> u128 { + self.inflow.abs_diff(self.outflow) + } + + pub fn is_expired(&self, now: Timestamp) -> bool { + self.period_end < now + } + + // Mutating methods + pub fn expire(&mut self, now: Timestamp, duration: u64) { + self.inflow = 0; + self.outflow = 0; + self.period_end = now.plus_seconds(duration); + } + + pub fn add_flow(&mut self, direction: FlowType, value: u128) { + match direction { + FlowType::In => self.inflow = self.inflow.saturating_add(value), + FlowType::Out => self.outflow = self.outflow.saturating_add(value), + } + } +} + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema)] +pub struct Quota { + pub name: String, + pub max_percentage_send: u32, + pub max_percentage_recv: u32, + pub duration: u64, +} + +impl Quota { + /// Calculates the max capacity based on the total value of the channel + pub fn capacity_at(&self, total_value: &u128, direction: &FlowType) -> u128 { + let max_percentage = match direction { + FlowType::In => self.max_percentage_recv, + FlowType::Out => self.max_percentage_send, + }; + total_value * (max_percentage as u128) / 100_u128 + } +} + +impl From<&QuotaMsg> for Quota { + fn from(msg: &QuotaMsg) -> Self { + let send_recv = ( + cmp::min(msg.send_recv.0, 100), + cmp::min(msg.send_recv.1, 100), + ); + Quota { + name: msg.name.clone(), + max_percentage_send: send_recv.0, + max_percentage_recv: send_recv.1, + duration: msg.duration, + } + } +} + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema)] +pub struct ChannelFlow { + pub quota: Quota, + pub flow: Flow, +} + +/// Only this module can manage the contract +pub const GOVMODULE: Item = Item::new("gov_module"); +/// Only this module can execute transfers +pub const IBCMODULE: Item = Item::new("ibc_module"); +// For simplicity, the map keys (ibc channel) refers to the "host" channel on the +// osmosis side. This means that on PacketSend it will refer to the source +// channel while on PacketRecv it refers to the destination channel. +// +// It is the responsibility of the go module to pass the appropriate channel +// when sending the messages +pub const CHANNEL_FLOWS: Map<&str, Vec> = Map::new("flow"); + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn flow() { + let epoch = Timestamp::from_seconds(0); + let mut flow = Flow::new(0_u32, 0_u32, epoch, RESET_TIME_WEEKLY); + + assert!(!flow.is_expired(epoch)); + assert!(!flow.is_expired(epoch.plus_seconds(RESET_TIME_DAILY))); + assert!(!flow.is_expired(epoch.plus_seconds(RESET_TIME_WEEKLY))); + assert!(flow.is_expired(epoch.plus_seconds(RESET_TIME_WEEKLY).plus_nanos(1))); + + assert_eq!(flow.balance(), 0_u128); + flow.add_flow(FlowType::In, 5); + assert_eq!(flow.balance(), 5_u128); + flow.add_flow(FlowType::Out, 2); + assert_eq!(flow.balance(), 3_u128); + // Adding flow doesn't affect expiration + assert!(!flow.is_expired(epoch.plus_seconds(RESET_TIME_DAILY))); + + flow.expire(epoch.plus_seconds(RESET_TIME_WEEKLY), RESET_TIME_WEEKLY); + assert_eq!(flow.balance(), 0_u128); + assert_eq!(flow.inflow, 0_u128); + assert_eq!(flow.outflow, 0_u128); + assert_eq!(flow.period_end, epoch.plus_seconds(RESET_TIME_WEEKLY * 2)); + + // Expiration has moved + assert!(!flow.is_expired(epoch.plus_seconds(RESET_TIME_WEEKLY).plus_nanos(1))); + assert!(!flow.is_expired(epoch.plus_seconds(RESET_TIME_WEEKLY * 2))); + assert!(flow.is_expired(epoch.plus_seconds(RESET_TIME_WEEKLY * 2).plus_nanos(1))); + } +}