Skip to main content

runlimit_core/
check.rs

1use thiserror::Error;
2
3use crate::{CounterKey, FixedWindowPolicy, RateLimitPolicy, SubjectKey};
4
5/// One proposed quota charge against a policy and opaque subject.
6///
7/// A newly constructed check has cost 1. Custom costs are validated against
8/// the referenced policy so a backend never receives a zero-cost or
9/// intrinsically impossible check.
10#[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    /// Constructs a check with the default cost of 1.
27    pub const fn new(policy: &'a P, subject: SubjectKey) -> Self {
28        Self {
29            policy,
30            subject,
31            cost: 1,
32        }
33    }
34
35    /// Validates and constructs a check with a custom cost.
36    ///
37    /// # Errors
38    ///
39    /// Returns an error when `cost` is zero or exceeds the policy capacity.
40    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    /// Returns this check with a validated custom cost.
50    ///
51    /// # Errors
52    ///
53    /// Returns an error when `cost` is zero or exceeds the policy capacity.
54    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    /// Returns the policy evaluated by this check.
61    pub const fn policy(&self) -> &'a P {
62        self.policy
63    }
64
65    /// Returns the opaque subject key evaluated by this check.
66    pub const fn subject(&self) -> SubjectKey {
67        self.subject
68    }
69
70    /// Returns the complete logical identity of the stored counter.
71    pub fn counter_key(&self) -> CounterKey {
72        CounterKey::new(self.policy.fingerprint(), self.subject)
73    }
74
75    /// Returns the nonzero quota cost.
76    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/// An invalid check cost.
95#[derive(Clone, Copy, Debug, Error, Eq, PartialEq)]
96pub enum CheckError {
97    /// The requested cost was zero.
98    #[error("check cost must be greater than zero")]
99    ZeroCost,
100    /// The requested cost exceeded the referenced policy's capacity.
101    #[error("check cost ({cost}) exceeds the policy capacity ({capacity})")]
102    CostExceedsCapacity {
103        /// Requested cost.
104        cost: u64,
105        /// Maximum cost accepted by the policy.
106        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}