From ef45382c5cd525676fe302293f3432927fddc8f6 Mon Sep 17 00:00:00 2001 From: Alex Snaps Date: Tue, 3 Dec 2024 05:34:43 -0500 Subject: [PATCH] Treat CEL NoSuchKey as ignored false in pred, and miss in keys Signed-off-by: Alex Snaps --- limitador/src/counter.rs | 19 ++++--- limitador/src/lib.rs | 12 +++- limitador/src/limit.rs | 13 +++-- limitador/src/limit/cel.rs | 56 ++++++++++++++----- limitador/src/storage/disk/rocksdb_storage.rs | 3 +- limitador/src/storage/distributed/mod.rs | 6 +- limitador/src/storage/in_memory.rs | 10 +++- limitador/src/storage/keys.rs | 23 +++++--- limitador/src/storage/redis/counters_cache.rs | 1 + limitador/src/storage/redis/redis_cached.rs | 9 ++- 10 files changed, 107 insertions(+), 45 deletions(-) diff --git a/limitador/src/counter.rs b/limitador/src/counter.rs index 7dfebe0b..788bf8a7 100644 --- a/limitador/src/counter.rs +++ b/limitador/src/counter.rs @@ -20,18 +20,21 @@ impl Counter { pub fn new>>( limit: L, set_variables: HashMap, - ) -> LimitadorResult { + ) -> LimitadorResult> { 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, - }) + match variables { + None => Ok(None), + Some(variables) => Ok(Some(Self { + limit, + set_variables: variables, + remaining: None, + expires_in: None, + })), + } } pub(super) fn resolved_vars>>( @@ -162,7 +165,7 @@ mod tests { ) .expect("failed creating counter"); assert_eq!( - counter.set_variables.get(var), + counter.unwrap().set_variables.get(var), Some("13".to_string()).as_ref() ); } diff --git a/limitador/src/lib.rs b/limitador/src/lib.rs index 72507ea9..d5016a18 100644 --- a/limitador/src/lib.rs +++ b/limitador/src/lib.rs @@ -483,7 +483,11 @@ impl RateLimiter { limits .iter() .filter(|lim| lim.applies(&ctx)) - .map(|lim| Counter::new(Arc::clone(lim), values.clone())) + .filter_map(|lim| match Counter::new(Arc::clone(lim), values.clone()) { + Ok(None) => None, + Ok(Some(c)) => Some(Ok(c)), + Err(e) => Some(Err(e)), + }) .collect() } } @@ -660,7 +664,11 @@ impl AsyncRateLimiter { limits .iter() .filter(|lim| lim.applies(&ctx)) - .map(|lim| Counter::new(Arc::clone(lim), values.clone())) + .filter_map(|lim| match Counter::new(Arc::clone(lim), values.clone()) { + Ok(None) => None, + Ok(Some(c)) => Some(Ok(c)), + Err(e) => Some(Err(e)), + }) .collect() } } diff --git a/limitador/src/limit.rs b/limitador/src/limit.rs index 5aeaa046..6653cb07 100644 --- a/limitador/src/limit.rs +++ b/limitador/src/limit.rs @@ -134,15 +134,19 @@ impl Limit { pub fn resolve_variables( &self, vars: HashMap, - ) -> Result, EvaluationError> { + ) -> Result>, EvaluationError> { let ctx = Context::new(String::default(), &vars); let mut map = BTreeMap::new(); for variable in &self.variables { let name = variable.source().into(); - let value = variable.eval(&ctx)?; - map.insert(name, value); + match variable.eval(&ctx)? { + None => return Ok(None), + Some(value) => { + map.insert(name, value); + } + } } - Ok(map) + Ok(Some(map)) } #[cfg(feature = "disk_storage")] @@ -464,6 +468,7 @@ mod tests { assert_eq!( Counter::new(limit, map) .expect("failed") + .unwrap() .set_variables() .get("bar.endsWith('baz')"), Some(&"true".to_string()) diff --git a/limitador/src/limit/cel.rs b/limitador/src/limit/cel.rs index 8a4561cf..4674692a 100644 --- a/limitador/src/limit/cel.rs +++ b/limitador/src/limit/cel.rs @@ -159,19 +159,25 @@ impl Expression { } } - pub fn eval(&self, ctx: &Context) -> Result { - match self.resolve(ctx)? { - Value::Int(i) => Ok(i.to_string()), - Value::UInt(i) => Ok(i.to_string()), - Value::Float(f) => Ok(f.to_string()), - Value::String(s) => Ok(s.to_string()), - Value::Null => Ok("null".to_owned()), - Value::Bool(b) => Ok(b.to_string()), - val => Err(err_on_value(val)), + pub fn eval(&self, ctx: &Context) -> Result, EvaluationError> { + let result = self.resolve(ctx); + match result { + Ok(value) => match value { + Value::Int(i) => Ok(i.to_string()), + Value::UInt(i) => Ok(i.to_string()), + Value::Float(f) => Ok(f.to_string()), + Value::String(s) => Ok(s.to_string()), + Value::Null => Ok("null".to_owned()), + Value::Bool(b) => Ok(b.to_string()), + val => Err(err_on_value(val)), + } + .map(Some), + Err(ExecutionError::NoSuchKey(_)) => Ok(None), + Err(err) => Err(err.into()), } } - pub fn resolve(&self, ctx: &Context) -> Result { + pub(super) fn resolve(&self, ctx: &Context) -> Result { Value::resolve(&self.expression, &ctx.ctx) } @@ -287,9 +293,14 @@ impl Predicate { { return Ok(false); } - match self.expression.resolve(ctx)? { - Value::Bool(b) => Ok(b), - v => Err(err_on_value(v)), + + match self.expression.resolve(ctx) { + Ok(value) => match value { + Value::Bool(b) => Ok(b), + v => Err(err_on_value(v)), + }, + Err(ExecutionError::NoSuchKey(_)) => Ok(false), + Err(err) => Err(err.into()), } } } @@ -345,12 +356,12 @@ impl From for String { #[cfg(test)] mod tests { use super::{Context, Expression, Predicate}; - use std::collections::HashSet; + use std::collections::{HashMap, HashSet}; #[test] fn expression() { let exp = Expression::parse("100").expect("failed to parse"); - assert_eq!(exp.eval(&ctx()), Ok(String::from("100"))); + assert_eq!(exp.eval(&ctx()), Ok(Some(String::from("100")))); } #[test] @@ -377,6 +388,21 @@ mod tests { assert_eq!(pred.test(&ctx()), Ok(true)); } + #[test] + fn predicate_no_var() { + let pred = Predicate::parse("not_there == 42").expect("failed to parse"); + assert_eq!(pred.test(&ctx()), Ok(false)); + } + + #[test] + fn predicate_no_key() { + let pred = Predicate::parse("there.not == 42").expect("failed to parse"); + assert_eq!( + pred.test(&(&HashMap::from([("there".to_string(), String::default())])).into()), + Ok(false) + ); + } + #[test] fn unexpected_value_predicate() { let pred = Predicate::parse("42").expect("failed to parse"); diff --git a/limitador/src/storage/disk/rocksdb_storage.rs b/limitador/src/storage/disk/rocksdb_storage.rs index 05d04f04..356aaaaf 100644 --- a/limitador/src/storage/disk/rocksdb_storage.rs +++ b/limitador/src/storage/disk/rocksdb_storage.rs @@ -253,7 +253,8 @@ mod tests { limit, HashMap::from([("app_id".to_string(), "foo".to_string())]), ) - .unwrap(); + .unwrap() + .expect("must have a counter"); let tmp = TempDir::new().expect("We should have a dir!"); { diff --git a/limitador/src/storage/distributed/mod.rs b/limitador/src/storage/distributed/mod.rs index 81d9d6ed..aee13fa2 100644 --- a/limitador/src/storage/distributed/mod.rs +++ b/limitador/src/storage/distributed/mod.rs @@ -48,7 +48,8 @@ impl CounterStorage for CrInMemoryStorage { limits.entry(key.clone()).or_insert(Arc::new(CounterEntry { key, counter: Counter::new(limit.clone(), HashMap::default()) - .expect("counter creation can't fail! no vars to resolve!"), + .expect("counter creation can't fail! no vars to resolve!") + .expect("must have a counter"), value: CrCounterValue::new( self.identifier.clone(), limit.max_value(), @@ -343,6 +344,7 @@ fn encode_limit_to_key(limit: &Limit) -> Vec { .map(|k| (k, "".to_string())) .collect(), ) - .expect("counter creation can't fail! faked vars!"); + .expect("counter creation can't fail! faked vars!") + .expect("must have a counter"); key_for_counter_v2(&counter) } diff --git a/limitador/src/storage/in_memory.rs b/limitador/src/storage/in_memory.rs index ac5e2c07..3b273d9c 100644 --- a/limitador/src/storage/in_memory.rs +++ b/limitador/src/storage/in_memory.rs @@ -212,7 +212,9 @@ impl InMemoryStorage { if limit.namespace() == namespace { res.insert( // todo fixme - Counter::new(limit.clone(), HashMap::default()).unwrap(), + Counter::new(limit.clone(), HashMap::default()) + .unwrap() + .unwrap(), counter.clone(), ); } @@ -271,12 +273,14 @@ mod tests { limit_1, HashMap::from([("app_id".to_string(), "foo".to_string())]), ) - .expect("counter creation failed!"); + .expect("counter creation failed!") + .expect("Should have a counter"); let counter_2 = Counter::new( limit_2, HashMap::from([("app_id".to_string(), "foo".to_string())]), ) - .expect("counter creation failed!"); + .expect("counter creation failed!") + .expect("Should have a counter"); storage.update_counter(&counter_1, 1).unwrap(); storage.update_counter(&counter_2, 1).unwrap(); diff --git a/limitador/src/storage/keys.rs b/limitador/src/storage/keys.rs index 9d4e1d83..58506328 100644 --- a/limitador/src/storage/keys.rs +++ b/limitador/src/storage/keys.rs @@ -157,7 +157,8 @@ mod tests { limit.clone(), HashMap::from([("app_id".to_string(), "foo".to_string())]), ) - .expect("counter creation failed!"); + .expect("counter creation failed!") + .expect("must have a counter"); let raw = key_for_counter(&counter); assert_eq!(counter, partial_counter_from_counter_key(&raw)); } @@ -176,7 +177,8 @@ mod tests { limit.clone(), HashMap::from([("app_id".to_string(), "foo".to_string())]), ) - .expect("counter creation failed!"); + .expect("counter creation failed!") + .expect("must have a counter"); let mut other = counter.clone(); other.set_remaining(123); other.set_expires_in(Duration::from_millis(456)); @@ -387,7 +389,9 @@ pub mod bin { vars.insert("role".to_string(), "admin".to_string()); vars.insert("app_id".to_string(), "123".to_string()); vars.insert("wat".to_string(), "dunno".to_string()); - let counter = Counter::new(limit.clone(), vars).expect("counter creation failed!"); + let counter = Counter::new(limit.clone(), vars) + .expect("counter creation failed!") + .expect("must have a counter"); let raw = key_for_counter(&counter); let key_back: CounterKey = @@ -408,7 +412,9 @@ pub mod bin { ); let mut variables = HashMap::default(); variables.insert("app_id".to_string(), "123".to_string()); - let counter = Counter::new(limit.clone(), variables).expect("counter creation failed!"); + let counter = Counter::new(limit.clone(), variables) + .expect("counter creation failed!") + .expect("must have a counter"); let raw = key_for_counter(&counter); assert_eq!(counter, partial_counter_from_counter_key(&raw)); } @@ -427,7 +433,8 @@ pub mod bin { limit, HashMap::from([("app_id".to_string(), "foo".to_string())]), ) - .expect("counter creation failed!"); + .expect("counter creation failed!") + .expect("must have a counter"); let serialized_counter = key_for_counter(&counter); let prefix = prefix_for_namespace(namespace); @@ -457,14 +464,16 @@ pub mod bin { limit_with_id, HashMap::from([("app_id".to_string(), "foo".to_string())]), ) - .expect("counter creation failed!"); + .expect("counter creation failed!") + .expect("must have a counter"); let serialized_with_id_counter = key_for_counter(&counter_with_id); let counter_without_id = Counter::new( limit_without_id, HashMap::from([("app_id".to_string(), "foo".to_string())]), ) - .expect("counter creation failed!"); + .expect("counter creation failed!") + .expect("must have a counter"); let serialized_without_id_counter = key_for_counter(&counter_without_id); // the original key_for_counter continues to encode kinda big diff --git a/limitador/src/storage/redis/counters_cache.rs b/limitador/src/storage/redis/counters_cache.rs index 2645996d..03cdc900 100644 --- a/limitador/src/storage/redis/counters_cache.rs +++ b/limitador/src/storage/redis/counters_cache.rs @@ -687,5 +687,6 @@ mod tests { values, ) .expect("failed creating counter") + .expect("Should have a counter") } } diff --git a/limitador/src/storage/redis/redis_cached.rs b/limitador/src/storage/redis/redis_cached.rs index b9b12ca7..70aaac75 100644 --- a/limitador/src/storage/redis/redis_cached.rs +++ b/limitador/src/storage/redis/redis_cached.rs @@ -449,7 +449,8 @@ mod tests { ), HashMap::from([("app_id".to_string(), "foo".to_string())]), ) - .expect("counter creation failed!"); + .expect("counter creation failed!") + .expect("must have a counter"); let arc = Arc::new(CachedCounterValue::from_authority( &counter, @@ -512,7 +513,8 @@ mod tests { ), HashMap::from([("app_id".to_string(), "foo".to_string())]), ) - .expect("counter creation failed!"); + .expect("counter creation failed!") + .expect("must have a counter"); let mock_response = Value::Array(vec![ Value::Int(8), @@ -572,7 +574,8 @@ mod tests { ), HashMap::from([("app_id".to_string(), "foo".to_string())]), ) - .expect("counter creation failed!"); + .expect("counter creation failed!") + .expect("must have a counter"); let error: RedisError = io::Error::new(io::ErrorKind::TimedOut, "That was long!").into(); assert!(error.is_timeout());