diff --git a/src/engine.rs b/src/engine.rs index 8877d63..f1d64c6 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -8,7 +8,7 @@ use crate::parser::*; use crate::scheduler::*; use crate::utils::gather_functions; use crate::value::*; -use crate::QueryResults; +use crate::{Extension, QueryResults}; use std::convert::AsRef; use std::path::Path; @@ -348,4 +348,92 @@ impl Engine { self.interpreter.create_rule_prefixes()?; Ok(self.interpreter.get_data_mut().clone()) } + + /// Add a custom builtin (extension). + /// + /// * `path`: The fully qualified path of the builtin. + /// * `nargs`: The number of arguments the builtin takes. + /// * `extension`: The [`Extension`] instance. + /// + /// ```rust + /// # use regorus::*; + /// # use anyhow::{bail, Result}; + /// # fn main() -> Result<()> { + /// let mut engine = Engine::new(); + /// + /// // Policy uses `do_magic` custom builtin. + /// engine.add_policy( + /// "test.rego".to_string(), + /// r#"package test + /// x = do_magic(1) + /// "#.to_string(), + /// )?; + /// + /// // Evaluating fails since `do_magic` is not defined. + /// assert!(engine.eval_query("data.test.x".to_string(), false).is_err()); + /// + /// // Add extension to implement `do_magic`. The extension can be stateful. + /// let mut magic = 8; + /// engine.add_extension("do_magic".to_string(), 1 , Box::new(move | mut params: Vec | { + /// // params is mut and therefore individual values can be removed from it and modified. + /// // The number of parameters (1) has already been validated. + /// + /// match ¶ms[0].as_i64() { + /// Ok(i) => { + /// // Compute value + /// let v = *i + magic; + /// // Update extension state. + /// magic += 1; + /// Ok(Value::from(v)) + /// } + /// // Extensions can raise errors. Regorus will add location information to + /// // the error. + /// _ => bail!("do_magic expects i64 value") + /// } + /// }))?; + /// + /// // Evaluation will now succeed. + /// let r = engine.eval_query("data.test.x".to_string(), false)?; + /// assert_eq!(r.result[0].expressions[0].value.as_i64()?, 9); + /// + /// // Cloning the engine will also clone the extension. + /// let mut engine1 = engine.clone(); + /// + /// // Evaluating again will return a different value since the extension is stateful. + /// let r = engine.eval_query("data.test.x".to_string(), false)?; + /// assert_eq!(r.result[0].expressions[0].value.as_i64()?, 10); + /// + /// // The second engine has a clone of the extension. + /// let r = engine1.eval_query("data.test.x".to_string(), false)?; + /// assert_eq!(r.result[0].expressions[0].value.as_i64()?, 10); + /// + /// // Once added, the extension cannot be replaced or removed. + /// assert!(engine.add_extension("do_magic".to_string(), 1, Box::new(|_:Vec| { + /// Ok(Value::Undefined) + /// })).is_err()); + /// + /// // Extensions don't support out-parameter syntax. + /// engine.add_policy( + /// "policy.rego".to_string(), + /// r#"package invalid + /// x = y { + /// # y = do_magic(2) + /// do_magic(2, y) # y is supplied as an out parameter. + /// } + /// "#.to_string() + /// )?; + /// + /// // Evaluation fails since y is not defined. + /// assert!(engine.eval_query("data.invalid.y".to_string(), false).is_err()); + /// # Ok(()) + /// # } + /// ``` + pub fn add_extension( + &mut self, + path: String, + nargs: u8, + extension: Box, + ) -> Result<()> { + self.interpreter.add_extension(path, nargs, extension) + } } diff --git a/src/interpreter.rs b/src/interpreter.rs index 0f69e68..f23aa3d 100644 --- a/src/interpreter.rs +++ b/src/interpreter.rs @@ -9,7 +9,7 @@ use crate::parser::Parser; use crate::scheduler::*; use crate::utils::*; use crate::value::*; -use crate::{Expression, Location, QueryResult, QueryResults}; +use crate::{Expression, Extension, Location, QueryResult, QueryResults}; use anyhow::{anyhow, bail, Result}; use log::info; @@ -65,6 +65,7 @@ pub struct Interpreter { allow_deprecated: bool, strict_builtin_errors: bool, imports: BTreeMap>, + extensions: HashMap)>, } impl Default for Interpreter { @@ -175,6 +176,7 @@ impl Interpreter { allow_deprecated: true, strict_builtin_errors: true, imports: BTreeMap::default(), + extensions: HashMap::new(), } } @@ -2101,6 +2103,7 @@ impl Interpreter { } let orig_fcn_path = fcn_path; + let mut with_functions_saved = None; let fcn_path = match self.with_functions.get(&orig_fcn_path) { Some(FunctionModifier::Function(p)) => { @@ -2121,6 +2124,7 @@ impl Interpreter { _ => orig_fcn_path.clone(), }; + let mut extension = None; let empty: Vec> = vec![]; let (fcns_rules, fcn_module) = match self.lookup_function_by_name(&fcn_path) { Some((fcns, m)) => (fcns, Some(m.clone())), @@ -2134,6 +2138,11 @@ impl Interpreter { // process default functions later. (&empty, self.module.clone()) } + // Look up extension. + else if let Some(ext) = self.extensions.get_mut(&fcn_path) { + extension = Some(ext); + (&empty, None) + } // Look up builtin function. else if let Some(builtin) = self.lookup_builtin(span, &fcn_path)? { let r = self.eval_builtin_call(span, &fcn_path.clone(), *builtin, params); @@ -2153,6 +2162,21 @@ impl Interpreter { return Ok(Value::Undefined); } + if let Some((nargs, ext)) = extension { + if param_values.len() != *nargs as usize { + bail!(span.error("incorrect number of parameters supplied to extension")); + } + let r = ext(param_values); + // Restore with_functions. + if let Some(with_functions) = with_functions_saved { + self.with_functions = with_functions; + } + match r { + Ok(v) => return Ok(v), + Err(e) => bail!(span.error(&format!("{e}"))), + } + } + let fcns = fcns_rules.clone(); let mut results: Vec = Vec::new(); @@ -3016,43 +3040,7 @@ impl Interpreter { Ok(()) } - pub fn eval_rule(&mut self, module: &Ref, rule: &Ref) -> Result<()> { - // Skip reprocessing rule - if self.processed.contains(rule) { - return Ok(()); - } - - // Skip default rules - if let Rule::Default { .. } = rule.as_ref() { - return Ok(()); - } - - self.active_rules.push(rule.clone()); - if self.active_rules.iter().filter(|&r| r == rule).count() == 2 { - let mut msg = String::default(); - for r in &self.active_rules { - let refr = Self::get_rule_refr(r); - let span = refr.span(); - msg += span - .source - .message(span.line, span.col, "depends on", "") - .as_str(); - } - msg += "cyclic evaluation"; - let refr = Self::get_rule_refr(rule); - let span = refr.span(); - return Err(span.source.error( - span.line, - span.col, - format!("recursion detected when evaluating rule:{msg}").as_str(), - )); - } - - // Back up local variables of current function and empty - // the local variables of callee function. - let scopes = std::mem::take(&mut self.scopes); - let prev_module = self.set_current_module(Some(module.clone()))?; - + fn eval_rule_impl(&mut self, module: &Ref, rule: &Ref) -> Result<()> { match rule.as_ref() { Rule::Spec { span, @@ -3133,10 +3121,53 @@ impl Interpreter { } _ => bail!("internal error: unexpected"), } + Ok(()) + } + + pub fn eval_rule(&mut self, module: &Ref, rule: &Ref) -> Result<()> { + // Skip reprocessing rule + if self.processed.contains(rule) { + return Ok(()); + } + + // Skip default rules + if let Rule::Default { .. } = rule.as_ref() { + return Ok(()); + } + + self.active_rules.push(rule.clone()); + if self.active_rules.iter().filter(|&r| r == rule).count() == 2 { + let mut msg = String::default(); + for r in &self.active_rules { + let refr = Self::get_rule_refr(r); + let span = refr.span(); + msg += span + .source + .message(span.line, span.col, "depends on", "") + .as_str(); + } + msg += "cyclic evaluation"; + self.active_rules.pop(); + let refr = Self::get_rule_refr(rule); + let span = refr.span(); + return Err(span.source.error( + span.line, + span.col, + format!("recursion detected when evaluating rule:{msg}").as_str(), + )); + } + + // Back up local variables of current function and empty + // the local variables of callee function. + let scopes = std::mem::take(&mut self.scopes); + let prev_module = self.set_current_module(Some(module.clone()))?; + + let res = self.eval_rule_impl(module, rule); + self.set_current_module(prev_module)?; self.scopes = scopes; match self.active_rules.pop() { - Some(ref r) if r == rule => Ok(()), + Some(ref r) if r == rule => res, _ => bail!("internal error: current rule not active"), } } @@ -3409,4 +3440,18 @@ impl Interpreter { } Ok(()) } + + pub fn add_extension( + &mut self, + path: String, + nargs: u8, + extension: Box, + ) -> Result<()> { + if let std::collections::hash_map::Entry::Vacant(v) = self.extensions.entry(path) { + v.insert((nargs, extension)); + Ok(()) + } else { + bail!("extension already added"); + } + } } diff --git a/src/lib.rs b/src/lib.rs index 03824a7..2808219 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -260,6 +260,37 @@ pub struct QueryResults { pub result: Vec, } +/// A user defined builtin function implementation. +/// +/// It is not necessary to implement this trait directly. +pub trait Extension: FnMut(Vec) -> anyhow::Result { + /// Fn, FnMut etc are not sized and cannot be cloned in their boxed form. + /// clone_box exists to overcome that. + fn clone_box<'a>(&self) -> Box + where + Self: 'a; +} + +/// Automatically make matching closures a valid [`Extension`]. +impl Extension for F +where + F: FnMut(Vec) -> anyhow::Result + Clone, +{ + fn clone_box<'a>(&self) -> Box + where + Self: 'a, + { + Box::new(self.clone()) + } +} + +/// Implement clone for a boxed extension using [`Extension::clone_box`]. +impl<'a> Clone for Box { + fn clone(&self) -> Self { + (**self).clone_box() + } +} + /// Items in `unstable` are likely to change. #[doc(hidden)] pub mod unstable { diff --git a/tests/engine/mod.rs b/tests/engine/mod.rs new file mode 100644 index 0000000..7b8fa2b --- /dev/null +++ b/tests/engine/mod.rs @@ -0,0 +1,103 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use anyhow::{bail, Result}; +use regorus::*; + +#[test] +fn extension() -> Result<()> { + fn repeat(mut params: Vec) -> Result { + match params.remove(0) { + Value::String(s) => { + let s = s.as_ref().to_owned(); + Ok(Value::from(s.clone() + &s)) + } + _ => bail!("param must be string"), + } + } + let mut engine = Engine::new(); + engine.add_policy( + "test.rego".to_string(), + r#"package test + x = repeat("hello") + "# + .to_string(), + )?; + + // Raises error since repeat is not defined. + assert!(engine.eval_query("data.test.x".to_string(), false).is_err()); + + // Register extension. + engine.add_extension("repeat".to_string(), 1, Box::new(repeat))?; + + // Adding extension twice is error. + assert!(engine + .add_extension( + "repeat".to_string(), + 1, + Box::new(|_| { Ok(Value::Undefined) }) + ) + .is_err()); + + let r = engine.eval_query("data.test.x".to_string(), false)?; + assert_eq!( + r.result[0].expressions[0].value.as_string()?.as_ref(), + "hellohello" + ); + + Ok(()) +} + +#[test] +fn extension_with_state() -> Result<()> { + #[derive(Clone)] + struct Gen { + n: i64, + } + + let mut engine = Engine::new(); + engine.add_policy( + "test.rego".to_string(), + r#"package test + x = gen() + "# + .to_string(), + )?; + + let mut g = Box::new(Gen { n: 5 }); + engine.add_extension( + "gen".to_string(), + 0, + Box::new(move |_: Vec| { + let v = Value::from(g.n); + g.n += 1; + Ok(v) + }), + )?; + + // First eval. + let r = engine.eval_query("data.test.x".to_string(), false)?; + assert_eq!(r.result[0].expressions[0].value.as_i64()?, 5); + + // Second eval will produce a new value since for each query, the + // internal evaluation state of the interpreter is cleared. + // This might change in the future. + let r = engine.eval_query("data.test.x".to_string(), false)?; + assert_eq!(r.result[0].expressions[0].value.as_i64()?, 6); + + // Clone the engine. + // This should also clone the stateful extension. + let mut engine1 = engine.clone(); + + // Both the engines should produce the same value. + let r = engine.eval_query("data.test.x".to_string(), false)?; + let r1 = engine1.eval_query("data.test.x".to_string(), false)?; + assert_eq!( + r.result[0].expressions[0].value, + r1.result[0].expressions[0].value + ); + + assert_eq!(r.result[0].expressions[0].value.as_i64()?, 7); + + Ok(()) +} diff --git a/tests/mod.rs b/tests/mod.rs index aa08f98..4782498 100644 --- a/tests/mod.rs +++ b/tests/mod.rs @@ -1,6 +1,7 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. +mod engine; mod lexer; mod parser; mod value;