orml_traits/
rate_limit.rs

1use frame_support::Parameter;
2use parity_scale_codec::Encode;
3use sp_runtime::{traits::Member, RuntimeDebug};
4
5#[derive(PartialEq, Eq, RuntimeDebug)]
6pub enum RateLimiterError {
7	ExceedLimit,
8}
9
10/// Rate Limiter
11pub trait RateLimiter {
12	/// The type for the rate limiter.
13	type RateLimiterId: Parameter + Member + Copy;
14
15	/// Check whether the rate limiter of can be bypassed according to the
16	/// `key`.
17	fn is_whitelist(limiter_id: Self::RateLimiterId, key: impl Encode) -> bool;
18
19	/// Check whether the `value` can be consumed under the limit of
20	/// `limit_key`.
21	fn can_consume(
22		limiter_id: Self::RateLimiterId,
23		limit_key: impl Encode,
24		value: u128,
25	) -> Result<(), RateLimiterError>;
26
27	/// The handler function to consume quota.
28	fn consume(limiter_id: Self::RateLimiterId, limit_key: impl Encode, value: u128);
29
30	/// Try consume quota.
31	fn try_consume(
32		limiter_id: Self::RateLimiterId,
33		limit_key: impl Encode + Clone,
34		value: u128,
35		whitelist_check: Option<impl Encode>,
36	) -> Result<(), RateLimiterError> {
37		let need_consume = match whitelist_check {
38			Some(whitelist_key) => !Self::is_whitelist(limiter_id, whitelist_key),
39			None => true,
40		};
41
42		if need_consume {
43			Self::can_consume(limiter_id, limit_key.clone(), value)?;
44			Self::consume(limiter_id, limit_key, value);
45		}
46
47		Ok(())
48	}
49}
50
51impl RateLimiter for () {
52	type RateLimiterId = ();
53
54	fn is_whitelist(_: Self::RateLimiterId, _: impl Encode) -> bool {
55		true
56	}
57
58	fn can_consume(_: Self::RateLimiterId, _: impl Encode, _: u128) -> Result<(), RateLimiterError> {
59		Ok(())
60	}
61
62	fn consume(_: Self::RateLimiterId, _: impl Encode, _: u128) {}
63}