From c53d0023471ff3cf1997fe1af262c8290e6aa66b Mon Sep 17 00:00:00 2001 From: Anand Krishnamoorthi <35780660+anakrish@users.noreply.github.com> Date: Wed, 1 Nov 2023 21:57:06 -0700 Subject: [PATCH] Improvements (#33) 1. Skip recording undefined variables 2. Parse `in` correctly if it is not imported. 3. base64.decode 4. Handle `with` modifier for qualified data and input. Signed-off-by: Anand Krishnamoorthi --- Cargo.toml | 1 + src/builtins/encoding.rs | 28 ++++++ src/builtins/mod.rs | 3 +- src/interpreter.rs | 139 +++++++++++++++++---------- src/parser.rs | 6 +- src/scheduler.rs | 4 +- tests/interpreter/mod.rs | 50 ++++------ tests/parser/cases/some/some.in.yaml | 1 + 8 files changed, 149 insertions(+), 83 deletions(-) create mode 100644 src/builtins/encoding.rs diff --git a/Cargo.toml b/Cargo.toml index e612bce..cfb6137 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,6 +14,7 @@ log = "0.4.17" env_logger="0.10.0" lazy_static = "1.4.0" rand = "0.8.5" +data-encoding = "2.4.0" [dev-dependencies] clap = { version = "4.4.7", features = ["derive"] } diff --git a/src/builtins/encoding.rs b/src/builtins/encoding.rs new file mode 100644 index 0000000..938f470 --- /dev/null +++ b/src/builtins/encoding.rs @@ -0,0 +1,28 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use crate::ast::Expr; +use crate::builtins; +use crate::builtins::utils::{ensure_args_count, ensure_string}; +use crate::lexer::Span; +use crate::value::Value; + +use std::collections::HashMap; + +use anyhow::Result; +use data_encoding::BASE64; + +pub fn register(m: &mut HashMap<&'static str, builtins::BuiltinFcn>) { + m.insert("base64.decode", (base64_decode, 1)); +} + +fn base64_decode(span: &Span, params: &[Expr], args: &[Value]) -> Result { + let name = "base64.decode"; + ensure_args_count(span, name, params, args, 1)?; + + let encoded_str = ensure_string(name, ¶ms[0], &args[0])?; + let decoded_bytes = BASE64.decode(encoded_str.as_bytes())?; + Ok(Value::String( + String::from_utf8_lossy(&decoded_bytes).to_string(), + )) +} diff --git a/src/builtins/mod.rs b/src/builtins/mod.rs index b040aad..cf50e6b 100644 --- a/src/builtins/mod.rs +++ b/src/builtins/mod.rs @@ -7,6 +7,7 @@ mod bitwise; pub mod comparison; mod conversions; mod debugging; +mod encoding; pub mod numbers; mod objects; pub mod sets; @@ -44,7 +45,7 @@ lazy_static! { conversions::register(&mut m); //units::register(&mut m); types::register(&mut m); - //encoding::register(&mut m); + encoding::register(&mut m); //token_signing::register(&mut m); //token_verification::register(&mut m); //time::register(&mut m); diff --git a/src/interpreter.rs b/src/interpreter.rs index 03e8620..60bbd90 100644 --- a/src/interpreter.rs +++ b/src/interpreter.rs @@ -26,6 +26,7 @@ pub struct Interpreter<'source> { input: Value, data: Value, init_data: Value, + with_document: Value, scopes: Vec, // TODO: handle recursive calls where same expr could have different values. loop_var_values: BTreeMap<&'source Expr<'source>, Value>, @@ -81,6 +82,9 @@ struct LoopExpr<'source> { impl<'source> Interpreter<'source> { pub fn new(modules: Vec<&'source Module<'source>>) -> Result> { + let mut with_document = Value::new_object(); + *Self::make_or_get_value_mut(&mut with_document, &["data"])? = Value::new_object(); + *Self::make_or_get_value_mut(&mut with_document, &["input"])? = Value::new_object(); Ok(Interpreter { modules, module: None, @@ -90,6 +94,7 @@ impl<'source> Interpreter<'source> { input: Value::new_object(), data: Value::new_object(), init_data: Value::new_object(), + with_document, scopes: vec![Scope::new()], contexts: vec![], loop_var_values: BTreeMap::new(), @@ -241,6 +246,9 @@ impl<'source> Interpreter<'source> { match ident { "_" => true, _ => match self.lookup_local_var(ident) { + // Vars declared using `some v` can be loop vars. + // They are initialized to undefined. + Some(Value::Undefined) => true, // If ident is a local var (in current or parent scopes), // then it is not a loop var. Some(_) => false, @@ -478,6 +486,11 @@ impl<'source> Interpreter<'source> { } }; + // Omit recording undefined values. + if value == Value::Undefined { + return Ok(Value::Bool(false)); + } + self.add_variable_or(name)?; // TODO: optimize this @@ -803,44 +816,12 @@ impl<'source> Interpreter<'source> { Value::from_map(expr) } - fn eval_stmt( + fn eval_stmt_impl( &mut self, stmt: &'source LiteralStmt<'source>, stmts: &[&'source LiteralStmt<'source>], ) -> Result { - let mut to_restore = vec![]; - for wm in &stmt.with_mods { - // Evaluate value and ref - let value = self.eval_expr(&wm.r#as)?; - let path = Parser::get_path_ref_components(&wm.refr)?; - let mut path: Vec<&str> = path.iter().map(|s| s.text()).collect(); - - // TODO: multiple modules and qualified path - if path.len() > 2 && format!("{}.{}", path[0], path[1]) == self.current_module_path { - path = path[1..].to_vec(); - } - - // Set new values in modifications table - let mut saved = false; - for (i, _) in path.iter().enumerate() { - let vref = Self::make_or_get_value_mut(&mut self.data, &path[0..i])?; - if vref == &Value::Undefined { - to_restore.push((path[0..i].to_vec(), vref.clone())); - saved = false; - break; - } - } - - // TODO: input - let vref = Self::make_or_get_value_mut(&mut self.data, &path[..])?; - if !saved { - to_restore.push((path, vref.clone())); - } - - *vref = value; - } - - let r = Ok(match &stmt.literal { + Ok(match &stmt.literal { Literal::Expr { span, expr, .. } => { let value = match expr { Expr::Call { span, fcn, params } => self.eval_call( @@ -944,16 +925,59 @@ impl<'source> Interpreter<'source> { } self.eval_every(span, key, value, domain, query)? } - }); + }) + } - for (path, value) in to_restore.into_iter().rev() { - if value == Value::Undefined { - unimplemented!("handle undefined restore"); - } else { - let vref = Self::make_or_get_value_mut(&mut self.data, &path[..])?; - *vref = value; + fn eval_stmt( + &mut self, + stmt: &'source LiteralStmt<'source>, + stmts: &[&'source LiteralStmt<'source>], + ) -> Result { + let saved_state = if !stmt.with_mods.is_empty() { + // Save state; + let with_document = self.with_document.clone(); + let input = self.input.clone(); + let data = self.data.clone(); + let processed = self.processed.clone(); + + // Apply with modifiers. + for wm in &stmt.with_mods { + // Evaluate value and ref + let value = self.eval_expr(&wm.r#as)?; + let path = Parser::get_path_ref_components(&wm.refr)?; + let path: Vec<&str> = path.iter().map(|s| s.text()).collect(); + + if path[0] == "input" || path[0] == "data" { + *Self::make_or_get_value_mut(&mut self.with_document, &path[..])? = value; + } /* else if path.len() == 1 { + // TODO: handle var in current module. + } else { + // TODO: error about input, data + } */ + // TODO: functions } + + self.data = self.with_document["data"].clone(); + self.input = self.with_document["input"].clone(); + self.processed.clear(); + + (with_document, input, data, processed) + } else { + ( + Value::Undefined, + Value::Undefined, + Value::Undefined, + BTreeSet::new(), + ) + }; + + let r = self.eval_stmt_impl(stmt, stmts); + + // Restore state. + if saved_state.0 != Value::Undefined { + (self.with_document, self.input, self.data, self.processed) = saved_state; } + r } @@ -989,8 +1013,8 @@ impl<'source> Interpreter<'source> { if let Some(idx) = self.lookup_local_var(loop_expr.index) { if loop_expr_value[&idx] != Value::Undefined { result = self.eval_stmts_in_loop(stmts, &loops[1..])? || result; + return Ok(result); } - return Ok(result); } // Save the current scope and restore it after evaluating the statements so @@ -1672,6 +1696,14 @@ impl<'source> Interpreter<'source> { // Ensure that rules are evaluated if name == "data" { + let v = Self::get_value_chained(self.data.clone(), fields); + + // If the rule has already been evaluated or specified via a with modifier, + // use that value. + if v != Value::Undefined { + return Ok(v); + } + // Evaluate rule corresponding to longest matching path. for i in (1..fields.len() + 1).rev() { let path = "data.".to_owned() + &fields[0..i].join("."); @@ -1680,15 +1712,9 @@ impl<'source> Interpreter<'source> { break; } } + Ok(Self::get_value_chained(self.data.clone(), fields)) } else if !self.modules.is_empty() { - // Add module prefix and ensure that any matching rule is evaluated. - let module_path = - Self::get_path_string(&self.current_module()?.package.refr, Some("data"))?; - let path = module_path + "." + name; - - self.ensure_rule_evaluated(path)?; - let mut path: Vec<&str> = Parser::get_path_ref_components(&self.module.unwrap().package.refr)? .iter() @@ -1696,6 +1722,21 @@ impl<'source> Interpreter<'source> { .collect(); path.push(name); + let v = Self::get_value_chained(self.data.clone(), &path); + + // If the rule has already been evaluated or specified via a with modifier, + // use that value. + if v != Value::Undefined { + return Ok(v); + } + + // Add module prefix and ensure that any matching rule is evaluated. + let module_path = + Self::get_path_string(&self.current_module()?.package.refr, Some("data"))?; + let rule_path = module_path + "." + name; + + self.ensure_rule_evaluated(rule_path)?; + let value = Self::get_value_chained(self.data.clone(), &path[..]); Ok(Self::get_value_chained(value, fields)) } else { diff --git a/src/parser.rs b/src/parser.rs index 0952d66..0b57c56 100644 --- a/src/parser.rs +++ b/src/parser.rs @@ -714,7 +714,7 @@ impl<'source> Parser<'source> { let start = self.tok.1.start; let mut expr = self.parse_bool_expr()?; - while self.tok.1.text() == "in" { + while self.tok.1.text() == "in" && self.future_keywords.get("in").is_some() { expr = self.parse_membership_tail(start, expr, None)?; } @@ -731,7 +731,7 @@ impl<'source> Parser<'source> { expr = self.parse_membership_tail(start, expr, Some(value))?; } - while self.tok.1.text() == "in" { + while self.tok.1.text() == "in" && self.is_imported_future_keyword("in") { expr = self.parse_membership_tail(start, expr, None)?; } @@ -834,7 +834,7 @@ impl<'source> Parser<'source> { vars.push(span); } - if self.tok.1.text() != "in" || self.future_keywords.get("in").is_none() { + if self.tok.1.text() != "in" || !self.is_imported_future_keyword("in") { if self.tok.1.text() == "in" { self.warn_future_keyword(); } diff --git a/src/scheduler.rs b/src/scheduler.rs index 2712f0b..826958e 100644 --- a/src/scheduler.rs +++ b/src/scheduler.rs @@ -913,8 +913,10 @@ impl<'a> Analyzer<'a> { collection, &mut scope, &mut first_use, - &mut col_definitions, // TODO: handle these definitions + &mut col_definitions, )?; + definitions.append(&mut col_definitions); + self.process_comprs( &col_comprs[..], &mut scope, diff --git a/tests/interpreter/mod.rs b/tests/interpreter/mod.rs index c330a32..5d765be 100644 --- a/tests/interpreter/mod.rs +++ b/tests/interpreter/mod.rs @@ -181,18 +181,14 @@ pub fn check_output(computed_results: &[Value], expected_results: &[Value]) -> R Ok(()) } -fn query_results_to_value(query_results: QueryResults) -> Result { +fn push_query_results(query_results: QueryResults, results: &mut Vec) { if let Some(query_result) = query_results.result.last() { if !query_result.bindings.is_empty_object() { - return Ok(query_result.bindings.clone()); - } else { - return match query_result.expressions.last() { - Some(v) => Ok(v["value"].clone()), - _ => bail!("no expressions in query results"), - }; + results.push(query_result.bindings.clone()); + } else if let Some(v) = query_result.expressions.last() { + results.push(v["value"].clone()); } } - bail!("query result incomplete") } pub fn eval_file_first_rule( @@ -268,22 +264,20 @@ pub fn eval_file_first_rule( } // Now eval the query. - results.push(query_results_to_value(interpreter.eval_user_query( - &query_node, - &query_schedule, - enable_tracing, - )?)?); + push_query_results( + interpreter.eval_user_query(&query_node, &query_schedule, enable_tracing)?, + &mut results, + ); } } else { // it no input is defined then one evaluation of all modules is performed interpreter.eval(&data_opt, &None, enable_tracing, Some(schedule))?; - // Now eval the query. - results.push(query_results_to_value(interpreter.eval_user_query( - &query_node, - &query_schedule, - enable_tracing, - )?)?); + // Now eval the query + push_query_results( + interpreter.eval_user_query(&query_node, &query_schedule, enable_tracing)?, + &mut results, + ); } Ok(results) @@ -359,22 +353,20 @@ pub fn eval_file( interpreter.eval_modules(&Some(input), enable_tracing)?; // Now eval the query. - results.push(query_results_to_value(interpreter.eval_user_query( - &query_node, - &query_schedule, - enable_tracing, - )?)?); + push_query_results( + interpreter.eval_user_query(&query_node, &query_schedule, enable_tracing)?, + &mut results, + ); } } else { // it no input is defined then one evaluation of all modules is performed interpreter.eval(&data_opt, &None, enable_tracing, Some(schedule))?; // Now eval the query. - results.push(query_results_to_value(interpreter.eval_user_query( - &query_node, - &query_schedule, - enable_tracing, - )?)?); + push_query_results( + interpreter.eval_user_query(&query_node, &query_schedule, enable_tracing)?, + &mut results, + ); } Ok(results) diff --git a/tests/parser/cases/some/some.in.yaml b/tests/parser/cases/some/some.in.yaml index 2fc717a..ef47190 100644 --- a/tests/parser/cases/some/some.in.yaml +++ b/tests/parser/cases/some/some.in.yaml @@ -134,6 +134,7 @@ cases: - note: same-line rego: | package test + import future.keywords.in x = y{ some a b in {4, 5} [1, 2, 3][a] == 3