sentinel_modsec/variables/
collection.rs

1//! Collection trait and implementations.
2
3use regex::Regex;
4use std::collections::HashMap;
5
6/// A collection of key-value pairs.
7pub trait Collection: Send + Sync {
8    /// Get all key-value pairs.
9    fn all(&self) -> Vec<(&str, &str)>;
10
11    /// Get values by key.
12    fn get(&self, key: &str) -> Option<Vec<&str>>;
13
14    /// Get values matching a regex pattern.
15    fn get_regex(&self, pattern: &Regex) -> Vec<(&str, &str)>;
16
17    /// Count items.
18    fn count(&self) -> usize;
19
20    /// Count items with specific key.
21    fn count_key(&self, key: &str) -> usize;
22}
23
24/// A mutable collection (TX, SESSION, etc.).
25pub trait MutableCollection: Collection {
26    /// Set a value.
27    fn set(&mut self, key: String, value: String);
28
29    /// Delete a key.
30    fn delete(&mut self, key: &str);
31
32    /// Increment a numeric value.
33    fn increment(&mut self, key: &str, amount: i64);
34
35    /// Decrement a numeric value.
36    fn decrement(&mut self, key: &str, amount: i64);
37}
38
39/// Simple hash map based collection.
40#[derive(Debug, Clone, Default)]
41pub struct HashMapCollection {
42    data: HashMap<String, Vec<String>>,
43}
44
45impl HashMapCollection {
46    /// Create a new empty collection.
47    pub fn new() -> Self {
48        Self::default()
49    }
50
51    /// Add a value to a key.
52    pub fn add(&mut self, key: String, value: String) {
53        self.data.entry(key).or_default().push(value);
54    }
55
56    /// Clear all data.
57    pub fn clear(&mut self) {
58        self.data.clear();
59    }
60}
61
62impl Collection for HashMapCollection {
63    fn all(&self) -> Vec<(&str, &str)> {
64        self.data
65            .iter()
66            .flat_map(|(k, vs)| vs.iter().map(move |v| (k.as_str(), v.as_str())))
67            .collect()
68    }
69
70    fn get(&self, key: &str) -> Option<Vec<&str>> {
71        self.data
72            .get(key)
73            .map(|vs| vs.iter().map(|s| s.as_str()).collect())
74    }
75
76    fn get_regex(&self, pattern: &Regex) -> Vec<(&str, &str)> {
77        self.data
78            .iter()
79            .filter(|(k, _)| pattern.is_match(k))
80            .flat_map(|(k, vs)| vs.iter().map(move |v| (k.as_str(), v.as_str())))
81            .collect()
82    }
83
84    fn count(&self) -> usize {
85        self.data.values().map(|v| v.len()).sum()
86    }
87
88    fn count_key(&self, key: &str) -> usize {
89        self.data.get(key).map(|v| v.len()).unwrap_or(0)
90    }
91}
92
93impl MutableCollection for HashMapCollection {
94    fn set(&mut self, key: String, value: String) {
95        self.data.insert(key, vec![value]);
96    }
97
98    fn delete(&mut self, key: &str) {
99        self.data.remove(key);
100    }
101
102    fn increment(&mut self, key: &str, amount: i64) {
103        let current: i64 = self
104            .data
105            .get(key)
106            .and_then(|v| v.first())
107            .and_then(|s| s.parse().ok())
108            .unwrap_or(0);
109        self.set(key.to_string(), (current + amount).to_string());
110    }
111
112    fn decrement(&mut self, key: &str, amount: i64) {
113        self.increment(key, -amount);
114    }
115}
sentinel_modsec/variables/collection.rs

sentinel_modsec/variables/
collection.rs
