Parse Limit's variables as exp and resolve

Signed-off-by: Alex Snaps <[email protected]>
Kuadrant · Nov 29, 2024 · 1652ba8 · 1652ba8
1 parent d60e5f7
commit 1652ba8
Show file tree

Hide file tree

Showing 15 changed files with 304 additions and 141 deletions.
diff --git a/limitador-server/src/envoy_rls/server.rs b/limitador-server/src/envoy_rls/server.rs
@@ -254,7 +254,7 @@ mod tests {
             1,
             60,
             vec!["req_method == 'GET'".try_into().expect("failed parsing!")],
-            vec!["app_id"],
+            vec!["app_id".try_into().expect("failed parsing!")],
         );
 
         let limiter = RateLimiter::new(10_000);
@@ -399,7 +399,7 @@ mod tests {
                 10,
                 60,
                 vec!["x == '1'".try_into().expect("failed parsing!")],
-                vec!["z"],
+                vec!["z".try_into().expect("failed parsing!")],
             ),
             Limit::new(
                 namespace,
@@ -409,7 +409,7 @@ mod tests {
                     "x == '1'".try_into().expect("failed parsing!"),
                     "y == '2'".try_into().expect("failed parsing!"),
                 ],
-                vec!["z"],
+                vec!["z".try_into().expect("failed parsing!")],
             ),
         ]
         .into_iter()
@@ -479,7 +479,7 @@ mod tests {
             10,
             60,
             vec!["x == '1'".try_into().expect("failed parsing!")],
-            vec!["y"],
+            vec!["y".try_into().expect("failed parsing!")],
         );
 
         let limiter = RateLimiter::new(10_000);
@@ -554,7 +554,7 @@ mod tests {
             1,
             60,
             vec!["x == '1'".try_into().expect("failed parsing!")],
-            vec!["y"],
+            vec!["y".try_into().expect("failed parsing!")],
         );
 
         let limiter = RateLimiter::new(10_000);

diff --git a/limitador-server/src/http_api/request_types.rs b/limitador-server/src/http_api/request_types.rs
@@ -1,5 +1,5 @@
 use limitador::counter::Counter as LimitadorCounter;
-use limitador::limit::{Limit as LimitadorLimit, ParseError, Predicate};
+use limitador::limit::{Expression, Limit as LimitadorLimit, ParseError, Predicate};
 use paperclip::actix::Apiv2Schema;
 use serde::{Deserialize, Serialize};
 use std::collections::{BTreeMap, HashMap};
@@ -46,6 +46,8 @@ impl TryFrom<Limit> for LimitadorLimit {
     fn try_from(limit: Limit) -> Result<Self, Self::Error> {
         let conditions: Result<Vec<Predicate>, ParseError> =
             limit.conditions.into_iter().map(|p| p.try_into()).collect();
+        let variables: Result<Vec<Expression>, ParseError> =
+            limit.variables.into_iter().map(|v| v.try_into()).collect();
 
         let mut limitador_limit = if let Some(id) = limit.id {
             Self::with_id(
@@ -54,15 +56,15 @@ impl TryFrom<Limit> for LimitadorLimit {
                 limit.max_value,
                 limit.seconds,
                 conditions?,
-                limit.variables,
+                variables?,
             )
         } else {
             Self::new(
                 limit.namespace,
                 limit.max_value,
                 limit.seconds,
                 conditions?,
-                limit.variables,
+                variables?,
             )
         };
 

diff --git a/limitador-server/src/http_api/server.rs b/limitador-server/src/http_api/server.rs
@@ -554,7 +554,7 @@ mod tests {
             max,
             60,
             vec!["req_method == 'GET'".try_into().expect("failed parsing!")],
-            vec!["app_id"],
+            vec!["app_id".try_into().expect("failed parsing!")],
         );
 
         match &limiter {

diff --git a/limitador/benches/bench.rs b/limitador/benches/bench.rs
@@ -540,7 +540,7 @@ fn generate_test_limits(scenario: &TestScenario) -> (Vec<Limit>, Vec<TestCallPar
     let mut variables = vec![];
     for idx_var in 0..scenario.n_vars_per_limit {
         let var_name = format!("var_{idx_var}");
-        variables.push(var_name.clone());
+        variables.push(var_name.clone().try_into().expect("failed parsing!"));
         test_values.insert(var_name, "1".into());
     }
 

diff --git a/limitador/src/counter.rs b/limitador/src/counter.rs
@@ -1,4 +1,5 @@
 use crate::limit::{Limit, Namespace};
+use crate::LimitadorResult;
 use serde::{Deserialize, Serialize};
 use std::collections::{BTreeMap, HashMap};
 use std::hash::{Hash, Hasher};
@@ -16,19 +17,37 @@ pub struct Counter {
 }
 
 impl Counter {
-    pub fn new<L: Into<Arc<Limit>>>(limit: L, set_variables: HashMap<String, String>) -> Self {
-        // TODO: check that all the variables defined in the limit are set.
+    pub fn new<L: Into<Arc<Limit>>>(
+        limit: L,
+        set_variables: HashMap<String, String>,
+    ) -> LimitadorResult<Self> {
+        let limit = limit.into();
+        let mut vars = set_variables;
+        vars.retain(|var, _| limit.has_variable(var));
 
+        let variables = limit.resolve_variables(vars)?;
+        Ok(Self {
+            limit,
+            set_variables: variables,
+            remaining: None,
+            expires_in: None,
+        })
+    }
+
+    pub(super) fn resolved_vars<L: Into<Arc<Limit>>>(
+        limit: L,
+        set_variables: HashMap<String, String>,
+    ) -> LimitadorResult<Self> {
         let limit = limit.into();
         let mut vars = set_variables;
         vars.retain(|var, _| limit.has_variable(var));
 
-        Self {
+        Ok(Self {
             limit,
             set_variables: vars.into_iter().collect(),
             remaining: None,
             expires_in: None,
-        }
+        })
     }
 
     #[cfg(any(feature = "redis_storage", feature = "disk_storage"))]
@@ -120,3 +139,28 @@ impl PartialEq for Counter {
         self.limit == other.limit && self.set_variables == other.set_variables
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::counter::Counter;
+    use crate::limit::Limit;
+    use std::collections::HashMap;
+
+    #[test]
+    fn resolves_variables() {
+        let limit = Limit::new(
+            "",
+            10,
+            60,
+            Vec::default(),
+            ["int(x) * 3".try_into().expect("failed parsing!")],
+        );
+        let key = "x".to_string();
+        let counter = Counter::new(limit, HashMap::from([(key.clone(), "14".to_string())]))
+            .expect("failed creating counter");
+        assert_eq!(
+            counter.set_variables.get(&key),
+            Some("42".to_string()).as_ref()
+        );
+    }
+}
diff --git a/limitador/src/lib.rs b/limitador/src/lib.rs
@@ -55,7 +55,7 @@
 //!      10,
 //!      60,
 //!      vec!["req_method == 'GET'".try_into().expect("failed parsing!")],
-//!      vec!["user_id"],
+//!      vec!["user_id".try_into().expect("failed parsing!")],
 //! );
 //! ```
 //!
@@ -72,7 +72,7 @@
 //!      10,
 //!      60,
 //!      vec!["req_method == 'GET'".try_into().expect("failed parsing!")],
-//!      vec!["user_id"],
+//!      vec!["user_id".try_into().expect("failed parsing!")],
 //! );
 //! let mut rate_limiter = RateLimiter::new(1000);
 //!
@@ -104,7 +104,7 @@
 //!      2,
 //!      60,
 //!      vec!["req_method == 'GET'".try_into().expect("failed parsing!")],
-//!      vec!["user_id"],
+//!      vec!["user_id".try_into().expect("failed parsing!")],
 //! );
 //! rate_limiter.add_limit(limit);
 //!
@@ -168,7 +168,7 @@
 //!      10,
 //!      60,
 //!      vec!["req_method == 'GET'".try_into().expect("failed parsing!")],
-//!      vec!["user_id"],
+//!      vec!["user_id".try_into().expect("failed parsing!")],
 //! );
 //!
 //! async {
@@ -481,13 +481,11 @@ impl RateLimiter {
     ) -> LimitadorResult<Vec<Counter>> {
         let limits = self.storage.get_limits(namespace);
 
-        let counters = limits
+        limits
             .iter()
             .filter(|lim| lim.applies(values))
             .map(|lim| Counter::new(Arc::clone(lim), values.clone()))
-            .collect();
-
-        Ok(counters)
+            .collect()
     }
 }
 
@@ -660,13 +658,11 @@ impl AsyncRateLimiter {
     ) -> LimitadorResult<Vec<Counter>> {
         let limits = self.storage.get_limits(namespace);
 
-        let counters = limits
+        limits
             .iter()
             .filter(|lim| lim.applies(values))
             .map(|lim| Counter::new(Arc::clone(lim), values.clone()))
-            .collect();
-
-        Ok(counters)
+            .collect()
     }
 }
 
@@ -693,7 +689,7 @@ fn classify_limits_by_namespace(
 
 #[cfg(test)]
 mod test {
-    use crate::limit::Limit;
+    use crate::limit::{Expression, Limit};
     use crate::RateLimiter;
     use std::collections::HashMap;
 
@@ -702,7 +698,7 @@ mod test {
         let rl = RateLimiter::new(100);
         let namespace = "foo";
 
-        let l = Limit::new(namespace, 42, 100, vec![], Vec::<String>::default());
+        let l = Limit::new(namespace, 42, 100, vec![], Vec::<Expression>::default());
         rl.add_limit(l.clone());
         let limits = rl.get_limits(&namespace.into());
         assert_eq!(limits.len(), 1);