1use std::collections::{HashMap, hash_map::Entry};
2
3use thiserror::Error;
4
5use crate::{Check, QuotaMode, RateLimitPolicy};
6
7pub fn validate_batch<P: RateLimitPolicy>(
18 checks: &[Check<'_, P>],
19 maximum: usize,
20) -> Result<(), BatchError> {
21 if checks.len() > maximum {
22 return Err(BatchError::BatchTooLarge {
23 actual: checks.len(),
24 maximum,
25 });
26 }
27
28 let first_mode = checks.first().map(|check| check.policy().quota_mode());
29 let mut first_indices = HashMap::with_capacity(checks.len());
30 for (duplicate_index, check) in checks.iter().enumerate() {
31 if let Some(first) = first_mode
32 && check.policy().quota_mode() != first
33 {
34 return Err(BatchError::MixedQuotaModes {
35 first,
36 index: duplicate_index,
37 actual: check.policy().quota_mode(),
38 });
39 }
40 match first_indices.entry(check.counter_key()) {
41 Entry::Vacant(entry) => {
42 entry.insert(duplicate_index);
43 }
44 Entry::Occupied(entry) => {
45 return Err(BatchError::DuplicateKey {
46 first_index: *entry.get(),
47 duplicate_index,
48 });
49 }
50 }
51 }
52
53 Ok(())
54}
55
56#[derive(Clone, Copy, Debug, Error, Eq, PartialEq)]
58pub enum BatchError {
59 #[error("batch check at index {duplicate_index} duplicates the counter at index {first_index}")]
61 DuplicateKey {
62 first_index: usize,
64 duplicate_index: usize,
66 },
67 #[error("batch has {actual} checks but the configured maximum is {maximum}")]
69 BatchTooLarge {
70 actual: usize,
72 maximum: usize,
74 },
75 #[error(
77 "batch policy at index {index} uses quota mode {actual:?}, which differs from the first policy's {first:?} mode"
78 )]
79 MixedQuotaModes {
80 first: QuotaMode,
82 index: usize,
84 actual: QuotaMode,
86 },
87}
88
89#[cfg(test)]
90mod tests {
91 use std::time::Duration;
92
93 use super::{BatchError, validate_batch};
94 use crate::{Check, FixedWindowPolicy, PolicyId, QuotaMode, ScopeId, SubjectKey};
95
96 fn policy(id: &str) -> FixedWindowPolicy {
97 FixedWindowPolicy::new(
98 PolicyId::new(id).unwrap(),
99 ScopeId::new("client").unwrap(),
100 10,
101 Duration::from_secs(60),
102 )
103 .unwrap()
104 }
105
106 fn subject(byte: u8) -> SubjectKey {
107 SubjectKey::from_digest([byte; 32])
108 }
109
110 #[test]
111 fn accepts_empty_and_distinct_batches() {
112 let first_policy = policy("auth.alpha");
113 let second_policy = policy("auth.beta");
114 let checks = [
115 Check::new(&first_policy, subject(1)),
116 Check::new(&second_policy, subject(1)),
117 Check::new(&first_policy, subject(2)),
118 ];
119
120 assert_eq!(validate_batch::<FixedWindowPolicy>(&[], 0), Ok(()));
121 assert_eq!(validate_batch(&checks, checks.len()), Ok(()));
122 }
123
124 #[test]
125 fn reports_the_first_duplicate_in_caller_order() {
126 let alpha = policy("auth.alpha");
127 let beta = policy("auth.beta");
128 let checks = [
129 Check::new(&beta, subject(1)),
130 Check::new(&beta, subject(1)),
131 Check::new(&alpha, subject(2)),
132 Check::new(&alpha, subject(2)),
133 ];
134
135 assert_eq!(
136 validate_batch(&checks, checks.len()),
137 Err(BatchError::DuplicateKey {
138 first_index: 0,
139 duplicate_index: 1,
140 })
141 );
142 }
143
144 #[test]
145 fn batch_size_error_takes_precedence_over_duplicate_detection() {
146 let policy = policy("auth.alpha");
147 let checks = [
148 Check::new(&policy, subject(1)),
149 Check::new(&policy, subject(1)),
150 ];
151
152 assert_eq!(
153 validate_batch(&checks, 1),
154 Err(BatchError::BatchTooLarge {
155 actual: 2,
156 maximum: 1,
157 })
158 );
159 }
160
161 #[test]
162 fn rejects_mixed_enforcement_modes_before_duplicate_detection() {
163 let enforced = policy("auth.alpha");
164 let shadow = policy("auth.beta").with_quota_mode(QuotaMode::Shadow);
165 let checks = [
166 Check::new(&enforced, subject(1)),
167 Check::new(&shadow, subject(2)),
168 Check::new(&enforced, subject(1)),
169 ];
170
171 assert_eq!(
172 validate_batch(&checks, checks.len()),
173 Err(BatchError::MixedQuotaModes {
174 first: QuotaMode::Enforce,
175 index: 1,
176 actual: QuotaMode::Shadow,
177 })
178 );
179 }
180}