Skip to main content

runlimit_core/
batch.rs

1use std::collections::{HashMap, hash_map::Entry};
2
3use thiserror::Error;
4
5use crate::{Check, QuotaMode, RateLimitPolicy};
6
7/// Validates backend-independent structural requirements for an atomic batch.
8///
9/// Duplicate counters are reported in caller order: `duplicate_index` is the
10/// earliest repeated input, and `first_index` is that counter's first input.
11///
12/// # Errors
13///
14/// Returns [`BatchError::BatchTooLarge`] before inspecting keys when the batch
15/// exceeds `maximum`. Otherwise returns [`BatchError::DuplicateKey`] for the
16/// first repeated logical counter.
17pub 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/// A backend-independent invalid atomic batch.
57#[derive(Clone, Copy, Debug, Error, Eq, PartialEq)]
58pub enum BatchError {
59    /// The batch included the same logical counter more than once.
60    #[error("batch check at index {duplicate_index} duplicates the counter at index {first_index}")]
61    DuplicateKey {
62        /// Index of the first occurrence in caller order.
63        first_index: usize,
64        /// Index of the repeated occurrence in caller order.
65        duplicate_index: usize,
66    },
67    /// The batch contained more checks than the backend's configured maximum.
68    #[error("batch has {actual} checks but the configured maximum is {maximum}")]
69    BatchTooLarge {
70        /// Submitted check count.
71        actual: usize,
72        /// Configured maximum.
73        maximum: usize,
74    },
75    /// The batch mixed enforced and shadow quota policies.
76    #[error(
77        "batch policy at index {index} uses quota mode {actual:?}, which differs from the first policy's {first:?} mode"
78    )]
79    MixedQuotaModes {
80        /// Quota mode of the first input policy.
81        first: QuotaMode,
82        /// Index of the first input with a different mode.
83        index: usize,
84        /// Quota mode at `index`.
85        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}