Skip to main content

kora_lib/usage_limit/rules/
transaction.rs

1use super::super::limiter::LimiterContext;
2
3const TX_KEY_PREFIX: &str = "kora:tx";
4
5/// Rule that limits the total number of transactions per user
6///
7/// Supports both lifetime limits (never resets) and time-windowed limits (resets periodically).
8#[derive(Debug)]
9pub struct TransactionRule {
10    max: u64,
11    window_seconds: Option<u64>,
12}
13
14impl TransactionRule {
15    pub fn new(max: u64, window_seconds: Option<u64>) -> Self {
16        Self { max, window_seconds }
17    }
18
19    pub fn storage_key(&self, user_id: &str, timestamp: u64) -> String {
20        let base = format!("{TX_KEY_PREFIX}:{user_id}");
21        match self.window_seconds {
22            Some(window) if window > 0 => format!("{base}:{}", timestamp / window),
23            _ => base,
24        }
25    }
26
27    /// How many units to increment for this transaction (always 1)
28    pub fn count_increment(&self, _ctx: &mut LimiterContext<'_>) -> u64 {
29        1
30    }
31
32    /// Maximum allowed count within the window (or lifetime)
33    pub fn max(&self) -> u64 {
34        self.max
35    }
36
37    /// Time window in seconds
38    pub fn window_seconds(&self) -> Option<u64> {
39        self.window_seconds
40    }
41
42    pub fn description(&self) -> String {
43        let window = self.window_seconds.map_or("lifetime".to_string(), |w| format!("per {w}s"));
44        format!("transaction ({window})")
45    }
46}
47
48#[cfg(test)]
49mod tests {
50    use super::*;
51    use crate::tests::transaction_mock::create_mock_resolved_transaction;
52
53    #[test]
54    fn test_transaction_rule_lifetime_key() {
55        let rule = TransactionRule::new(100, None);
56        let user_id = "test-user-123";
57
58        let key = rule.storage_key(user_id, 1000000);
59        assert_eq!(key, format!("kora:tx:{}", user_id));
60    }
61
62    #[test]
63    fn test_transaction_rule_windowed_key() {
64        let rule = TransactionRule::new(100, Some(3600));
65        let user_id = "test-user-456";
66
67        let key1 = rule.storage_key(user_id, 3600);
68        let key2 = rule.storage_key(user_id, 7199);
69        let key3 = rule.storage_key(user_id, 7200);
70
71        assert_eq!(key1, format!("kora:tx:{}:1", user_id));
72        assert_eq!(key2, format!("kora:tx:{}:1", user_id));
73        assert_eq!(key3, format!("kora:tx:{}:2", user_id));
74    }
75
76    #[test]
77    fn test_transaction_rule_count_increment() {
78        let rule = TransactionRule::new(100, None);
79        let tx = create_mock_resolved_transaction();
80        let user_id = "test-user-789".to_string();
81        let mut tx = tx;
82        let mut ctx =
83            LimiterContext { transaction: &mut tx, user_id, kora_signer: None, timestamp: 1000000 };
84
85        assert_eq!(rule.count_increment(&mut ctx), 1);
86    }
87
88    #[test]
89    fn test_transaction_rule_description() {
90        let lifetime = TransactionRule::new(100, None);
91        assert_eq!(lifetime.description(), "transaction (lifetime)");
92
93        let windowed = TransactionRule::new(100, Some(3600));
94        assert_eq!(windowed.description(), "transaction (per 3600s)");
95    }
96}