From fb188db8035c39876d795282f9335df020a8b1d9 Mon Sep 17 00:00:00 2001 From: Anand Krishnamoorthi Date: Sun, 12 Nov 2023 21:35:53 -0800 Subject: [PATCH] Allow with modifier for builtin and user functions Signed-off-by: Anand Krishnamoorthi --- src/builtins/mod.rs | 3 +- src/builtins/time.rs | 30 +++++ src/interpreter.rs | 126 +++++++++++++----- src/parser.rs | 48 +++---- src/scheduler.rs | 34 +++-- src/utils.rs | 36 +++-- src/value.rs | 5 + .../cases/builtins/arrays/reverse.yaml | 2 +- tests/opa/mod.rs | 10 +- 9 files changed, 216 insertions(+), 78 deletions(-) create mode 100644 src/builtins/time.rs diff --git a/src/builtins/mod.rs b/src/builtins/mod.rs index 95a1ef64..9589f014 100644 --- a/src/builtins/mod.rs +++ b/src/builtins/mod.rs @@ -13,6 +13,7 @@ pub mod numbers; mod objects; pub mod sets; mod strings; +mod time; mod tracing; pub mod types; mod utils; @@ -51,7 +52,7 @@ lazy_static! { encoding::register(&mut m); //token_signing::register(&mut m); //token_verification::register(&mut m); - //time::register(&mut m); + time::register(&mut m); //cryptography::register(&mut m); //graphs::register(&mut m); //graphql::register(&mut m); diff --git a/src/builtins/time.rs b/src/builtins/time.rs new file mode 100644 index 00000000..18b8f925 --- /dev/null +++ b/src/builtins/time.rs @@ -0,0 +1,30 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use crate::ast::Expr; +use crate::builtins; +use crate::builtins::utils::ensure_args_count; +use crate::lexer::Span; +use crate::value::Value; + +use std::collections::HashMap; +use std::time::SystemTime; + +use anyhow::{bail, Result}; + +pub fn register(m: &mut HashMap<&'static str, builtins::BuiltinFcn>) { + m.insert("time.now_ns", (now_ns, 0)); +} + +fn now_ns(span: &Span, params: &[Expr], args: &[Value]) -> Result { + let name = "time.now_ns"; + ensure_args_count(span, name, params, args, 0)?; + + let now = SystemTime::now(); + let elapsed = match now.duration_since(SystemTime::UNIX_EPOCH) { + Ok(e) => e, + Err(e) => bail!(span.error(format!("could not fetch elapsed time. {e}").as_str())), + }; + let nanos = elapsed.as_nanos(); + Ok(Value::from_u128(nanos)) +} diff --git a/src/interpreter.rs b/src/interpreter.rs index c00e4726..8ecc270f 100644 --- a/src/interpreter.rs +++ b/src/interpreter.rs @@ -2,7 +2,7 @@ // Licensed under the MIT License. use crate::ast::*; -use crate::builtins; +use crate::builtins::{self, BuiltinFcn}; use crate::lexer::Span; use crate::parser::Parser; use crate::scheduler::*; @@ -27,6 +27,7 @@ pub struct Interpreter<'source> { data: Value, init_data: Value, with_document: Value, + with_functions: BTreeMap, scopes: Vec, // TODO: handle recursive calls where same expr could have different values. loop_var_values: BTreeMap, Value>, @@ -96,6 +97,7 @@ impl<'source> Interpreter<'source> { data: Value::new_object(), init_data: Value::new_object(), with_document, + with_functions: BTreeMap::new(), scopes: vec![Scope::new()], contexts: vec![], loop_var_values: BTreeMap::new(), @@ -842,7 +844,11 @@ impl<'source> Interpreter<'source> { span, fcn, params, - get_extra_arg(expr, &self.functions), + get_extra_arg( + expr, + Some(self.current_module_path.as_str()), + &self.functions, + ), true, )?, _ => self.eval_expr(expr)?, @@ -872,7 +878,11 @@ impl<'source> Interpreter<'source> { span, fcn, params, - get_extra_arg(expr, &self.functions), + get_extra_arg( + expr, + Some(self.current_module_path.as_str()), + &self.functions, + ), false, )?, _ => self.eval_expr(expr)?, @@ -947,19 +957,43 @@ impl<'source> Interpreter<'source> { stmt: &'source LiteralStmt, stmts: &[&'source LiteralStmt], ) -> Result { + let mut skip_exec = false; 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(); + let with_functions = self.with_functions.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(); + let target = path.join("."); + + let value = match self.eval_expr(&wm.r#as) { + Ok(Value::Undefined) => { + if let Ok(fcn_path) = get_path_string(&wm.r#as, None) { + let span = wm.r#as.span(); + if self.lookup_function_by_name(&fcn_path).is_some() { + let mut fcn_path = get_path_string(&wm.r#as, None)?; + if !fcn_path.starts_with("data.") { + fcn_path = self.current_module_path.clone() + "." + &fcn_path; + } + self.with_functions.insert(target, fcn_path); + continue; + } else if self.lookup_builtin(span, &fcn_path).is_ok() { + self.with_functions.insert(target, fcn_path); + continue; + } + } + skip_exec = true; + continue; + } + Ok(v) => v, + Err(e) => return Err(e), + }; if path[0] == "input" || path[0] == "data" { *Self::make_or_get_value_mut(&mut self.with_document, &path[..])? = value; @@ -968,28 +1002,38 @@ impl<'source> Interpreter<'source> { } 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) + (with_document, input, data, processed, with_functions) } else { ( Value::Undefined, Value::Undefined, Value::Undefined, BTreeSet::new(), + BTreeMap::new(), ) }; - let r = self.eval_stmt_impl(stmt, stmts); + let r = if !skip_exec { + self.eval_stmt_impl(stmt, stmts) + } else { + Ok(false) + }; // Restore state. if saved_state.0 != Value::Undefined { - (self.with_document, self.input, self.data, self.processed) = saved_state; + ( + self.with_document, + self.input, + self.data, + self.processed, + self.with_functions, + ) = saved_state; } r @@ -1443,24 +1487,22 @@ impl<'source> Interpreter<'source> { } } - fn lookup_function(&self, fcn: &'source Expr) -> Result<&Vec<&'source Rule>> { - let mut path = Self::get_path_string(fcn, None)?; + fn lookup_function_by_name(&self, path: &str) -> Option<&Vec<&'source Rule>> { + let mut path = path.to_owned(); if !path.starts_with("data.") { path = self.current_module_path.clone() + "." + &path; } match self.functions.get(&path) { - Some((r, _)) => Ok(r), - _ => { - bail!(fcn.span().error("function not found")) - } + Some((f, _)) => Some(f), + _ => None, } } fn eval_builtin_call( &mut self, span: &'source Span, - name: String, + name: &str, builtin: builtins::BuiltinFcn, params: &'source [Expr], ) -> Result { @@ -1474,7 +1516,7 @@ impl<'source> Interpreter<'source> { } } - let cache = builtins::must_cache(name.as_str()); + let cache = builtins::must_cache(name); if let Some(name) = &cache { if let Some(v) = self.builtins_cache.get(&(name, args.clone())) { return Ok(v.clone()); @@ -1496,30 +1538,47 @@ impl<'source> Interpreter<'source> { Ok(v) } + fn lookup_builtin(&self, span: &'source Span, path: &str) -> Result> { + Ok(if let Some(builtin) = builtins::BUILTINS.get(path) { + Some(builtin) + } else if let Some(builtin) = builtins::DEPRECATED.get(path) { + if !self.allow_deprecated { + bail!(span.error(format!("{path} is deprecated").as_str())) + } + Some(builtin) + } else { + None + }) + } + fn eval_call_impl( &mut self, span: &'source Span, fcn: &'source Expr, params: &'source [Expr], ) -> Result { - let fcns_rules = match self.lookup_function(fcn) { - Ok(r) => r, + let fcn_path = match get_path_string(fcn, None) { + Ok(p) => p, + _ => bail!(span.error("invalid function expression")), + }; + + let orig_fcn_path = fcn_path; + let fcn_path = match self.with_functions.remove(&orig_fcn_path) { + Some(p) => p, + _ => orig_fcn_path.clone(), + }; + + let fcns_rules = match self.lookup_function_by_name(&fcn_path) { + Some(r) => r, _ => { // Look up builtin function. - // TODO: handle with modifier - if let Ok(path) = Self::get_path_string(fcn, None) { - if let Some(builtin) = builtins::BUILTINS.get(path.as_str()) { - return self.eval_builtin_call(span, path, *builtin, params); - } - if let Some(builtin) = builtins::DEPRECATED.get(path.as_str()) { - if self.allow_deprecated { - return self.eval_builtin_call(span, path, *builtin, params); - } else { - bail!(span.error(format!("{path} is deprecated").as_str())) - } + if let Ok(Some(builtin)) = self.lookup_builtin(span, &fcn_path) { + let r = self.eval_builtin_call(span, &fcn_path.clone(), *builtin, params); + if orig_fcn_path != fcn_path { + self.with_functions.insert(orig_fcn_path, fcn_path); } + return r; } - return Err(span .source .error(span.line, span.col, "could not find function")); @@ -1629,6 +1688,11 @@ impl<'source> Interpreter<'source> { "functions must not produce multiple outputs for same inputs", )); } + + if orig_fcn_path != fcn_path { + self.with_functions.insert(orig_fcn_path, fcn_path); + } + Ok(results[0].clone()) } diff --git a/src/parser.rs b/src/parser.rs index 6534db03..54a043dc 100644 --- a/src/parser.rs +++ b/src/parser.rs @@ -522,19 +522,16 @@ impl<'source> Parser<'source> { } "(" if possible_fcn => { self.next_token()?; - if *self.tok.1.text() == ")" { - return Err(self - .tok - .1 - .error("at least one argument required for function calls")); - } - let mut args = vec![self.parse_in_expr()?]; - while *self.tok.1.text() == "," { - self.next_token()?; - match *self.tok.1.text() { - ")" => break, - "" if self.tok.0 == TokenKind::Eof => break, - _ => args.push(self.parse_in_expr()?), + let mut args = vec![]; + if *self.tok.1.text() != ")" { + args.push(self.parse_in_expr()?); + while *self.tok.1.text() == "," { + self.next_token()?; + match *self.tok.1.text() { + ")" => break, + "" if self.tok.0 == TokenKind::Eof => break, + _ => args.push(self.parse_in_expr()?), + } } } self.expect(")", "while parsing call expr")?; @@ -925,19 +922,19 @@ impl<'source> Parser<'source> { pub fn parse_query(&mut self, mut span: Span, end_delim: &str) -> Result { let state = self.clone(); - let _is_definite_query = matches!(*self.tok.1.text(), "some" | "every"); + let is_definite_query = matches!(*self.tok.1.text(), "some" | "every"); // TODO: empty query? let mut literals = vec![]; let stmt = match self.parse_literal_stmt() { Ok(stmt) => stmt, - Err(e) if _is_definite_query => return Err(e), - _ => { + Err(e) if is_definite_query => return Err(e), + Err(_) => { // There was error parsing the first literal // Restore the state and return. *self = state; - return Err(anyhow!("encountered , when expecting {}", end_delim)); + bail!(span.error(format!("expecting {end_delim}").as_str())); } }; @@ -1170,13 +1167,16 @@ impl<'source> Parser<'source> { "(" => { self.check_rule_ref(&rule_ref)?; self.next_token()?; - let mut args = vec![self.parse_term()?]; - while *self.tok.1.text() == "," { - self.next_token()?; - match *self.tok.1.text() { - ")" => break, - "" if self.tok.0 == TokenKind::Eof => break, - _ => args.push(self.parse_term()?), + let mut args = vec![]; + if *self.tok.1.text() != ")" { + args.push(self.parse_term()?); + while *self.tok.1.text() == "," { + self.next_token()?; + match *self.tok.1.text() { + ")" => break, + "" if self.tok.0 == TokenKind::Eof => break, + _ => args.push(self.parse_term()?), + } } } self.expect(")", "while parsing function rule args")?; diff --git a/src/scheduler.rs b/src/scheduler.rs index b8c7715d..93b66aa9 100644 --- a/src/scheduler.rs +++ b/src/scheduler.rs @@ -222,8 +222,6 @@ fn traverse<'a>(expr: &'a Expr, f: &mut dyn FnMut(&'a Expr) -> Result) -> ArrayCompr { .. } | SetCompr { .. } | ObjectCompr { .. } => (), Call { params, .. } => { - // TODO: is traversing function needed? - // traverse(fcn, f)?; for p in params { traverse(p, f)?; } @@ -372,6 +370,7 @@ pub struct Analyzer<'a> { scopes: Vec>, order: BTreeMap, Vec>, functions: FunctionTable<'a>, + current_module_path: String, } #[derive(Clone)] @@ -394,6 +393,7 @@ impl<'a> Analyzer<'a> { scopes: vec![], order: BTreeMap::new(), functions: FunctionTable::new(), + current_module_path: String::default(), } } @@ -453,7 +453,7 @@ impl<'a> Analyzer<'a> { Some(s) => s, _ => bail!("internal error: package scope missing"), }; - + self.current_module_path = path; self.scopes.push(scope.clone()); for r in &m.policy { self.analyze_rule(r)?; @@ -597,6 +597,7 @@ impl<'a> Analyzer<'a> { scope: &mut Scope<'a>, first_use: &mut BTreeMap<&'a str, Span>, definitions: &mut Vec>, + return_arg: &Option>, ) -> Result<(Vec<&'a str>, Vec<&'a Expr>)> { let mut used_vars = vec![]; let mut comprs = vec![]; @@ -608,8 +609,8 @@ impl<'a> Analyzer<'a> { { used_vars.push(name); first_use.entry(name).or_insert(v.clone()); - } else if !scope.inputs.contains(name) { - bail!(v.error(format!("Use of undefined variable `{name}` is unsafe").as_str())); + } else if !scope.inputs.contains(name) && Some(Ref::make(e)) != *return_arg { + bail!(v.error(format!("use of undefined variable `{name}` is unsafe").as_str())); } Ok(false) } @@ -623,6 +624,7 @@ impl<'a> Analyzer<'a> { scope, first_use, definitions, + return_arg, )?; definitions.push(Definition { var, @@ -764,6 +766,7 @@ impl<'a> Analyzer<'a> { scope, first_use, definitions, + &None, )?; self.process_comprs(&comprs[..], scope, first_use, &mut used_vars)?; let check_first_use = *op == AssignOp::ColEq; @@ -789,6 +792,7 @@ impl<'a> Analyzer<'a> { scope, first_use, definitions, + &None, )?; let check_first_use = false; self.process_comprs(&comprs[..], scope, first_use, &mut used_vars)?; @@ -825,8 +829,13 @@ impl<'a> Analyzer<'a> { self.process_assign_expr(op, lhs, rhs, scope, first_use, definitions) } _ => { - let (mut used_vars, comprs) = - Self::gather_used_vars_comprs_index_vars(expr, scope, first_use, definitions)?; + let (mut used_vars, comprs) = Self::gather_used_vars_comprs_index_vars( + expr, + scope, + first_use, + definitions, + &None, + )?; self.process_comprs(&comprs[..], scope, first_use, &mut used_vars)?; definitions.push(Definition { var: "", used_vars }); Ok(()) @@ -922,6 +931,7 @@ impl<'a> Analyzer<'a> { &mut scope, &mut first_use, &mut col_definitions, + &None, )?; definitions.append(&mut col_definitions); @@ -948,6 +958,7 @@ impl<'a> Analyzer<'a> { &mut scope, &mut first_use, &mut definitions, + &None, )?; if !definitions.is_empty() { bail!("internal error: non empty definitions"); @@ -964,12 +975,18 @@ impl<'a> Analyzer<'a> { // TODO: vars in compr } Literal::Expr { expr, .. } => { - if let Some(Expr::Var(return_arg)) = get_extra_arg(expr, &self.functions) { + let extra_arg = get_extra_arg( + expr, + Some(self.current_module_path.as_str()), + &self.functions, + ); + if let Some(ra @ Expr::Var(return_arg)) = &extra_arg { let (mut used_vars, comprs) = Self::gather_used_vars_comprs_index_vars( expr, &mut scope, &mut first_use, &mut definitions, + &Some(Ref::make(ra)), )?; let var = if *return_arg.text() != "_" { // The var in the return argument slot would have been processed as @@ -1010,6 +1027,7 @@ impl<'a> Analyzer<'a> { &mut scope, &mut first_use, &mut definitions, + &None, )?; self.process_comprs(&comprs[..], &mut scope, &mut first_use, &mut uv)?; definitions.push(Definition { diff --git a/src/utils.rs b/src/utils.rs index c9532ea3..e2524807 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -39,25 +39,41 @@ pub fn get_path_string(refr: &Expr, document: Option<&str>) -> Result { pub type FunctionTable<'a> = BTreeMap, u8)>; -pub fn get_extra_arg<'a>(expr: &'a Expr, functions: &FunctionTable) -> Option<&'a Expr> { +fn get_extra_arg_impl<'a>( + expr: &'a Expr, + module: Option<&str>, + functions: &FunctionTable, +) -> Result> { if let Expr::Call { fcn, params, .. } = expr { - if let Ok(path) = get_path_string(fcn, None) { - let n_args = if let Some((_, n_args)) = functions.get(&path) { - *n_args - } else if let Some((_, n_args)) = BUILTINS.get(path.as_str()) { + let full_path = get_path_string(fcn, module)?; + let n_args = if let Some((_, n_args)) = functions.get(&full_path) { + *n_args + } else { + let path = get_path_string(fcn, None)?; + if let Some((_, n_args)) = BUILTINS.get(path.as_str()) { *n_args } else if let Some((_, n_args)) = DEPRECATED.get(path.as_str()) { *n_args } else { - return None; - }; - if n_args as usize == params.len() - 1 { - return params.last(); + return Ok(None); } + }; + if (n_args as usize) + 1 == params.len() { + return Ok(params.last()); } } + Ok(None) +} - None +pub fn get_extra_arg<'a>( + expr: &'a Expr, + module: Option<&str>, + functions: &FunctionTable, +) -> Option<&'a Expr> { + match get_extra_arg_impl(expr, module, functions) { + Ok(a) => a, + _ => None, + } } pub fn gather_functions<'a>(modules: &[&'a Module]) -> Result> { diff --git a/src/value.rs b/src/value.rs index 21aebfcf..de2d7a7e 100644 --- a/src/value.rs +++ b/src/value.rs @@ -190,6 +190,11 @@ impl Value { Value::Number(Number(OrderedFloat(v))) } + pub fn from_u128(v: u128) -> Value { + // TODO: fix precision loss + Value::Number(Number(OrderedFloat(v as f64))) + } + pub fn from_array(a: Vec) -> Value { Value::Array(Rc::new(a)) } diff --git a/tests/interpreter/cases/builtins/arrays/reverse.yaml b/tests/interpreter/cases/builtins/arrays/reverse.yaml index 98be4df2..910d960a 100644 --- a/tests/interpreter/cases/builtins/arrays/reverse.yaml +++ b/tests/interpreter/cases/builtins/arrays/reverse.yaml @@ -39,7 +39,7 @@ cases: package test x = array.reverse() query: data.test - error: "at least one argument required for function calls" + error: "expects 1 argument" - note: reverse-more-args data: {} diff --git a/tests/opa/mod.rs b/tests/opa/mod.rs index c4e516fa..0852bf00 100644 --- a/tests/opa/mod.rs +++ b/tests/opa/mod.rs @@ -72,7 +72,7 @@ fn run_opa_tests() -> Result<()> { return Ok(()); } }; - dbg!(&opa_tests_dir); + let tests_path = Path::new(&opa_tests_dir); let mut status = BTreeMap::::new(); let mut n = 0; @@ -152,14 +152,18 @@ fn run_opa_tests() -> Result<()> { } println!("TESTSUITE STATUS"); - println!(" {:30} {:4} {:4}", "FOLDER", "PASS", "FAIL"); + println!(" {:40} {:4} {:4}", "FOLDER", "PASS", "FAIL"); + let (mut npass, mut nfail) = (0, 0); for (dir, (pass, fail)) in status { if fail == 0 { println!("\x1b[32m {dir:40}: {pass:4} {fail:4}\x1b[0m"); } else { println!("\x1b[31m {dir:40}: {pass:4} {fail:4}\x1b[0m"); } + npass += pass; + nfail += fail; } - + println!(); + println!("\x1b[31m {:40}: {npass:4} {nfail:4}\x1b[0m", "TOTAL"); Ok(()) }