kora_lib/usage_limit/
config.rs

1use serde::{Deserialize, Serialize};
2use solana_sdk::pubkey::Pubkey;
3use std::str::FromStr;
4use utoipa::ToSchema;
5
6use crate::{constant::DEFAULT_USAGE_LIMIT_FALLBACK_IF_UNAVAILABLE, error::KoraError};
7
8use super::rules::{InstructionRule, TransactionRule, UsageRule};
9
10/// Unified usage limit configuration
11#[derive(Clone, Serialize, Deserialize, ToSchema)]
12pub struct UsageLimitConfig {
13    /// Enable per-wallet usage limiting
14    pub enabled: bool,
15    /// Cache URL for shared usage limiting across multiple Kora instances (e.g., "redis://localhost:6379")
16    pub cache_url: Option<String>,
17    /// Fallback behavior when cache is unavailable - if true, allow transactions; if false, deny
18    pub fallback_if_unavailable: bool,
19    /// Usage limit rules - can be transaction-level or instruction-level
20    #[serde(default)]
21    pub rules: Vec<UsageLimitRuleConfig>,
22}
23
24impl Default for UsageLimitConfig {
25    fn default() -> Self {
26        Self {
27            enabled: false,
28            cache_url: None,
29            fallback_if_unavailable: DEFAULT_USAGE_LIMIT_FALLBACK_IF_UNAVAILABLE,
30            rules: vec![],
31        }
32    }
33}
34
35impl UsageLimitConfig {
36    /// Convert config rules to usage rule enums
37    pub fn build_rules(&self) -> Result<Vec<UsageRule>, KoraError> {
38        self.rules.iter().map(|r| r.build()).collect()
39    }
40}
41
42/// Configuration for a single usage limit rule (TOML-serializable)
43///
44/// Use `type` field to specify the rule type:
45/// - `type = "transaction"` - Counts all transactions
46/// - `type = "instruction"` - Counts specific instruction types
47///
48/// Example TOML:
49/// ```toml
50/// [[kora.usage_limit.rules]]
51/// type = "transaction"
52/// max = 100
53/// window_seconds = 3600
54///
55/// [[kora.usage_limit.rules]]
56/// type = "instruction"
57/// program = "TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA"
58/// instruction = "Transfer"
59/// max = 10
60/// ```
61#[derive(Clone, Serialize, Deserialize, ToSchema)]
62#[serde(tag = "type", rename_all = "lowercase")]
63pub enum UsageLimitRuleConfig {
64    /// Transaction-level limit - counts all transactions for a wallet
65    Transaction {
66        /// Maximum transactions allowed
67        max: u64,
68        /// Time window in seconds (None = lifetime)
69        #[serde(default)]
70        window_seconds: Option<u64>,
71    },
72    /// Instruction-level limit - counts specific instruction types
73    Instruction {
74        /// Program ID (e.g., "TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA")
75        program: String,
76        /// Instruction type (e.g., "Transfer", "Burn")
77        instruction: String,
78        /// Maximum allowed
79        max: u64,
80        /// Time window in seconds (None = lifetime)
81        #[serde(default)]
82        window_seconds: Option<u64>,
83    },
84}
85
86impl UsageLimitRuleConfig {
87    /// Build a usage rule enum from this config
88    pub fn build(&self) -> Result<UsageRule, KoraError> {
89        match self {
90            UsageLimitRuleConfig::Transaction { max, window_seconds } => {
91                Ok(UsageRule::Transaction(TransactionRule::new(*max, *window_seconds)))
92            }
93            UsageLimitRuleConfig::Instruction { program, instruction, max, window_seconds } => {
94                let program_pubkey = Pubkey::from_str(program).map_err(|e| {
95                    KoraError::InternalServerError(format!(
96                        "Invalid program in usage limit rule '{}': {}",
97                        program, e
98                    ))
99                })?;
100                Ok(UsageRule::Instruction(InstructionRule::new(
101                    program_pubkey,
102                    instruction.clone(),
103                    *max,
104                    *window_seconds,
105                )))
106            }
107        }
108    }
109}
110
111#[cfg(test)]
112mod tests {
113    use super::*;
114    use crate::tests::toml_mock::ConfigBuilder;
115
116    #[test]
117    fn test_usage_limit_config_parsing() {
118        let config = ConfigBuilder::new()
119            .with_usage_limit_config(true, Some("redis://localhost:6379"), false)
120            .with_usage_limit_rules(vec![
121                // Transaction lifetime rule
122                (None, None, 100, None),
123            ])
124            .build_config()
125            .unwrap();
126
127        assert!(config.kora.usage_limit.enabled);
128        assert_eq!(config.kora.usage_limit.cache_url, Some("redis://localhost:6379".to_string()));
129        assert!(!config.kora.usage_limit.fallback_if_unavailable);
130        assert_eq!(config.kora.usage_limit.rules.len(), 1);
131        match &config.kora.usage_limit.rules[0] {
132            UsageLimitRuleConfig::Transaction { max, window_seconds } => {
133                assert_eq!(*max, 100);
134                assert_eq!(*window_seconds, None);
135            }
136            _ => panic!("Expected Transaction rule"),
137        }
138    }
139
140    #[test]
141    fn test_usage_limit_config_default() {
142        let config = ConfigBuilder::new().build_config().unwrap();
143
144        assert!(!config.kora.usage_limit.enabled);
145        assert_eq!(config.kora.usage_limit.cache_url, None);
146        assert!(config.kora.usage_limit.rules.is_empty());
147        assert_eq!(
148            config.kora.usage_limit.fallback_if_unavailable,
149            DEFAULT_USAGE_LIMIT_FALLBACK_IF_UNAVAILABLE
150        );
151    }
152
153    #[test]
154    fn test_usage_limit_time_bucket_rule() {
155        let config = ConfigBuilder::new()
156            .with_usage_limit_config(true, Some("redis://localhost:6379"), false)
157            .with_usage_limit_rules(vec![
158                // Transaction time-bucket rule: 50 per hour
159                (None, None, 50, Some(3600)),
160            ])
161            .build_config()
162            .unwrap();
163
164        assert!(config.kora.usage_limit.enabled);
165        assert_eq!(config.kora.usage_limit.rules.len(), 1);
166        match &config.kora.usage_limit.rules[0] {
167            UsageLimitRuleConfig::Transaction { max, window_seconds } => {
168                assert_eq!(*max, 50);
169                assert_eq!(*window_seconds, Some(3600));
170            }
171            _ => panic!("Expected Transaction rule"),
172        }
173    }
174
175    #[test]
176    fn test_usage_limit_instruction_rules() {
177        let config = ConfigBuilder::new()
178            .with_usage_limit_config(true, Some("redis://localhost:6379"), false)
179            .with_usage_limit_rules(vec![
180                // Time-windowed instruction rule
181                (
182                    Some("TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA"),
183                    Some("Transfer"),
184                    10,
185                    Some(86400),
186                ),
187                // Lifetime instruction rule (no window)
188                (
189                    Some("ATokenGPvbdGVxr1b2hvZbsiqW5xWH25efTNsLJA8knL"),
190                    Some("CreateIdempotent"),
191                    3,
192                    None,
193                ),
194            ])
195            .build_config()
196            .unwrap();
197
198        assert!(config.kora.usage_limit.enabled);
199        assert_eq!(config.kora.usage_limit.rules.len(), 2);
200
201        // Check time-windowed rule
202        match &config.kora.usage_limit.rules[0] {
203            UsageLimitRuleConfig::Instruction { program, instruction, max, window_seconds } => {
204                assert_eq!(program, "TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA");
205                assert_eq!(instruction, "Transfer");
206                assert_eq!(*max, 10);
207                assert_eq!(*window_seconds, Some(86400));
208            }
209            _ => panic!("Expected Instruction rule"),
210        }
211
212        // Check lifetime rule
213        match &config.kora.usage_limit.rules[1] {
214            UsageLimitRuleConfig::Instruction { program, instruction, max, window_seconds } => {
215                assert_eq!(program, "ATokenGPvbdGVxr1b2hvZbsiqW5xWH25efTNsLJA8knL");
216                assert_eq!(instruction, "CreateIdempotent");
217                assert_eq!(*max, 3);
218                assert_eq!(*window_seconds, None);
219            }
220            _ => panic!("Expected Instruction rule"),
221        }
222    }
223
224    #[test]
225    fn test_transaction_rule_config_parse() {
226        let config = ConfigBuilder::new()
227            .with_usage_limit_config(true, Some("redis://localhost:6379"), false)
228            .with_usage_limit_rules(vec![
229                // Transaction rule with window
230                (None, None, 100, Some(3600)),
231            ])
232            .build_config()
233            .unwrap();
234
235        match &config.kora.usage_limit.rules[0] {
236            UsageLimitRuleConfig::Transaction { max, window_seconds } => {
237                assert_eq!(*max, 100);
238                assert_eq!(*window_seconds, Some(3600));
239            }
240            _ => panic!("Expected Transaction rule"),
241        }
242    }
243
244    #[test]
245    fn test_transaction_rule_config_lifetime() {
246        let config = ConfigBuilder::new()
247            .with_usage_limit_config(true, Some("redis://localhost:6379"), false)
248            .with_usage_limit_rules(vec![
249                // Transaction lifetime rule
250                (None, None, 50, None),
251            ])
252            .build_config()
253            .unwrap();
254
255        match &config.kora.usage_limit.rules[0] {
256            UsageLimitRuleConfig::Transaction { max, window_seconds } => {
257                assert_eq!(*max, 50);
258                assert_eq!(*window_seconds, None);
259            }
260            _ => panic!("Expected Transaction rule"),
261        }
262    }
263
264    #[test]
265    fn test_instruction_rule_config_parse() {
266        let config = ConfigBuilder::new()
267            .with_usage_limit_config(true, Some("redis://localhost:6379"), false)
268            .with_usage_limit_rules(vec![
269                // Instruction rule with window
270                (
271                    Some("TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA"),
272                    Some("Transfer"),
273                    10,
274                    Some(86400),
275                ),
276            ])
277            .build_config()
278            .unwrap();
279
280        match &config.kora.usage_limit.rules[0] {
281            UsageLimitRuleConfig::Instruction { program, instruction, max, window_seconds } => {
282                assert_eq!(program, "TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA");
283                assert_eq!(instruction, "Transfer");
284                assert_eq!(*max, 10);
285                assert_eq!(*window_seconds, Some(86400));
286            }
287            _ => panic!("Expected Instruction rule"),
288        }
289    }
290
291    #[test]
292    fn test_instruction_rule_config_lifetime() {
293        let config = ConfigBuilder::new()
294            .with_usage_limit_config(true, Some("redis://localhost:6379"), false)
295            .with_usage_limit_rules(vec![
296                // Instruction lifetime rule
297                (Some("TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA"), Some("Burn"), 5, None),
298            ])
299            .build_config()
300            .unwrap();
301
302        match &config.kora.usage_limit.rules[0] {
303            UsageLimitRuleConfig::Instruction { program, instruction, max, window_seconds } => {
304                assert_eq!(program, "TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA");
305                assert_eq!(instruction, "Burn");
306                assert_eq!(*max, 5);
307                assert_eq!(*window_seconds, None);
308            }
309            _ => panic!("Expected Instruction rule"),
310        }
311    }
312
313    #[test]
314    fn test_build_transaction_rule() {
315        let config = UsageLimitRuleConfig::Transaction { max: 100, window_seconds: Some(3600) };
316        let rule = config.build().unwrap();
317
318        assert_eq!(rule.rule_type(), "transaction");
319        assert_eq!(rule.max(), 100);
320        assert_eq!(rule.window_seconds(), Some(3600));
321    }
322
323    #[test]
324    fn test_build_instruction_rule() {
325        let config = UsageLimitRuleConfig::Instruction {
326            program: "TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA".to_string(),
327            instruction: "Transfer".to_string(),
328            max: 10,
329            window_seconds: None,
330        };
331        let rule = config.build().unwrap();
332
333        assert_eq!(rule.rule_type(), "instruction");
334        assert_eq!(rule.max(), 10);
335        assert_eq!(rule.window_seconds(), None);
336    }
337
338    #[test]
339    fn test_build_instruction_rule_invalid_program() {
340        let config = UsageLimitRuleConfig::Instruction {
341            program: "invalid".to_string(),
342            instruction: "Transfer".to_string(),
343            max: 10,
344            window_seconds: None,
345        };
346        assert!(config.build().is_err());
347    }
348
349    #[test]
350    fn test_usage_limit_config_build_rules() {
351        let config = UsageLimitConfig {
352            enabled: true,
353            cache_url: None,
354            fallback_if_unavailable: true,
355            rules: vec![
356                UsageLimitRuleConfig::Transaction { max: 100, window_seconds: None },
357                UsageLimitRuleConfig::Instruction {
358                    program: "TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA".to_string(),
359                    instruction: "Transfer".to_string(),
360                    max: 10,
361                    window_seconds: Some(86400),
362                },
363            ],
364        };
365
366        let rules = config.build_rules().unwrap();
367        assert_eq!(rules.len(), 2);
368        assert_eq!(rules[0].rule_type(), "transaction");
369        assert_eq!(rules[1].rule_type(), "instruction");
370    }
371}
kora_lib/usage_limit/config.rs

kora_lib/usage_limit/
config.rs
