use thiserror::Error;
use crate::{CounterKey, FixedWindowPolicy, SubjectKey};
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct Check<'a> {
policy: &'a FixedWindowPolicy,
subject: SubjectKey,
cost: u64,
}
impl<'a> Check<'a> {
pub const fn new(policy: &'a FixedWindowPolicy, subject: SubjectKey) -> Self {
Self {
policy,
subject,
cost: 1,
}
}
pub fn with_cost(
policy: &'a FixedWindowPolicy,
subject: SubjectKey,
cost: u64,
) -> Result<Self, CheckError> {
validate_cost(policy, cost)?;
Ok(Self {
policy,
subject,
cost,
})
}
pub fn try_with_cost(mut self, cost: u64) -> Result<Self, CheckError> {
validate_cost(self.policy, cost)?;
self.cost = cost;
Ok(self)
}
pub const fn policy(&self) -> &'a FixedWindowPolicy {
self.policy
}
pub const fn subject(&self) -> SubjectKey {
self.subject
}
pub const fn counter_key(&self) -> CounterKey {
CounterKey::new(self.policy.fingerprint(), self.subject)
}
pub const fn cost(&self) -> u64 {
self.cost
}
}
fn validate_cost(policy: &FixedWindowPolicy, cost: u64) -> Result<(), CheckError> {
if cost == 0 {
return Err(CheckError::ZeroCost);
}
if cost > policy.limit() {
return Err(CheckError::CostExceedsLimit {
cost,
limit: policy.limit(),
});
}
Ok(())
}
#[derive(Clone, Copy, Debug, Error, Eq, PartialEq)]
pub enum CheckError {
#[error("check cost must be greater than zero")]
ZeroCost,
#[error("check cost ({cost}) exceeds the policy limit ({limit})")]
CostExceedsLimit {
cost: u64,
limit: u64,
},
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use super::{Check, CheckError};
use crate::{FixedWindowPolicy, PolicyId, ScopeId, SubjectKey};
fn policy() -> FixedWindowPolicy {
FixedWindowPolicy::new(
PolicyId::new("auth.login").unwrap(),
ScopeId::new("client").unwrap(),
8,
Duration::from_secs(60),
)
.unwrap()
}
#[test]
fn defaults_to_one_unit_of_cost() {
let policy = policy();
let subject = SubjectKey::from_digest([1; 32]);
let check = Check::new(&policy, subject);
assert_eq!(check.policy(), &policy);
assert_eq!(check.subject(), subject);
assert_eq!(check.cost(), 1);
}
#[test]
fn accepts_cost_up_to_and_including_the_limit() {
let policy = policy();
let subject = SubjectKey::from_digest([2; 32]);
assert_eq!(Check::with_cost(&policy, subject, 3).unwrap().cost(), 3);
assert_eq!(
Check::new(&policy, subject)
.try_with_cost(policy.limit())
.unwrap()
.cost(),
policy.limit()
);
}
#[test]
fn rejects_zero_cost() {
let policy = policy();
assert_eq!(
Check::with_cost(&policy, SubjectKey::from_digest([3; 32]), 0),
Err(CheckError::ZeroCost)
);
}
#[test]
fn rejects_cost_above_policy_limit() {
let policy = policy();
assert_eq!(
Check::with_cost(&policy, SubjectKey::from_digest([4; 32]), 9),
Err(CheckError::CostExceedsLimit { cost: 9, limit: 8 })
);
}
}