1use thiserror::Error;
2
3use crate::{CounterKey, FixedWindowPolicy, RateLimitPolicy, SubjectKey};
4
5#[derive(Debug, Eq, PartialEq)]
11pub struct Check<'a, P: RateLimitPolicy + ?Sized = FixedWindowPolicy> {
12 policy: &'a P,
13 subject: SubjectKey,
14 cost: u64,
15}
16
17impl<P: RateLimitPolicy + ?Sized> Copy for Check<'_, P> {}
18
19impl<P: RateLimitPolicy + ?Sized> Clone for Check<'_, P> {
20 fn clone(&self) -> Self {
21 *self
22 }
23}
24
25impl<'a, P: RateLimitPolicy + ?Sized> Check<'a, P> {
26 pub const fn new(policy: &'a P, subject: SubjectKey) -> Self {
28 Self {
29 policy,
30 subject,
31 cost: 1,
32 }
33 }
34
35 pub fn with_cost(policy: &'a P, subject: SubjectKey, cost: u64) -> Result<Self, CheckError> {
41 validate_cost(policy, cost)?;
42 Ok(Self {
43 policy,
44 subject,
45 cost,
46 })
47 }
48
49 pub fn try_with_cost(mut self, cost: u64) -> Result<Self, CheckError> {
55 validate_cost(self.policy, cost)?;
56 self.cost = cost;
57 Ok(self)
58 }
59
60 pub const fn policy(&self) -> &'a P {
62 self.policy
63 }
64
65 pub const fn subject(&self) -> SubjectKey {
67 self.subject
68 }
69
70 pub fn counter_key(&self) -> CounterKey {
72 CounterKey::new(self.policy.fingerprint(), self.subject)
73 }
74
75 pub const fn cost(&self) -> u64 {
77 self.cost
78 }
79}
80
81fn validate_cost<P: RateLimitPolicy + ?Sized>(policy: &P, cost: u64) -> Result<(), CheckError> {
82 if cost == 0 {
83 return Err(CheckError::ZeroCost);
84 }
85 if cost > policy.capacity() {
86 return Err(CheckError::CostExceedsCapacity {
87 cost,
88 capacity: policy.capacity(),
89 });
90 }
91 Ok(())
92}
93
94#[derive(Clone, Copy, Debug, Error, Eq, PartialEq)]
96pub enum CheckError {
97 #[error("check cost must be greater than zero")]
99 ZeroCost,
100 #[error("check cost ({cost}) exceeds the policy capacity ({capacity})")]
102 CostExceedsCapacity {
103 cost: u64,
105 capacity: u64,
107 },
108}
109
110#[cfg(test)]
111mod tests {
112 use std::time::Duration;
113
114 use super::{Check, CheckError};
115 use crate::{FixedWindowPolicy, PolicyId, ScopeId, SubjectKey};
116
117 fn policy() -> FixedWindowPolicy {
118 FixedWindowPolicy::new(
119 PolicyId::new("auth.login").unwrap(),
120 ScopeId::new("client").unwrap(),
121 8,
122 Duration::from_secs(60),
123 )
124 .unwrap()
125 }
126
127 #[test]
128 fn defaults_to_one_unit_of_cost() {
129 let policy = policy();
130 let subject = SubjectKey::from_digest([1; 32]);
131 let check = Check::new(&policy, subject);
132
133 assert_eq!(check.policy(), &policy);
134 assert_eq!(check.subject(), subject);
135 assert_eq!(check.cost(), 1);
136 }
137
138 #[test]
139 fn accepts_cost_up_to_and_including_the_limit() {
140 let policy = policy();
141 let subject = SubjectKey::from_digest([2; 32]);
142
143 assert_eq!(Check::with_cost(&policy, subject, 3).unwrap().cost(), 3);
144 assert_eq!(
145 Check::new(&policy, subject)
146 .try_with_cost(policy.limit())
147 .unwrap()
148 .cost(),
149 policy.limit()
150 );
151 }
152
153 #[test]
154 fn rejects_zero_cost() {
155 let policy = policy();
156
157 assert_eq!(
158 Check::with_cost(&policy, SubjectKey::from_digest([3; 32]), 0),
159 Err(CheckError::ZeroCost)
160 );
161 }
162
163 #[test]
164 fn rejects_cost_above_policy_limit() {
165 let policy = policy();
166
167 assert_eq!(
168 Check::with_cost(&policy, SubjectKey::from_digest([4; 32]), 9),
169 Err(CheckError::CostExceedsCapacity {
170 cost: 9,
171 capacity: 8
172 })
173 );
174 }
175}