kata-containers/regorus
1
0
Fork 0
mirror of https://github.com/microsoft/regorus.git synced 2026-08-05 02:16:11 +00:00
Code Issues Packages Projects Releases Wiki Activity

Ability to add custom builtin functions (#132)

Browse Source 
Signed-off-by: Anand Krishnamoorthi <anakrish@microsoft.com>
This commit is contained in:
Anand Krishnamoorthi
2024-02-07 15:15:55 -08:00
committed by GitHub
parent 5717f9c249
commit a95a9d21b3
5 changed files with 308 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

90
src/engine.rs
View File

@@ -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<Value> | {
/// // params is mut and therefore individual values can be removed from it and modified.
/// // The number of parameters (1) has already been validated.
///
/// match &params[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<Value>| {
/// 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<dyn Extension>,
) -> Result<()> {
self.interpreter.add_extension(path, nargs, extension)
}
}

123
src/interpreter.rs
View File

@@ -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<String, Ref<Expr>>,
extensions: HashMap<String, (u8, Box<dyn Extension>)>,
}
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<Ref<Rule>> = 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<Value> = Vec::new();
@@ -3016,43 +3040,7 @@ impl Interpreter {
Ok(())
}
pub fn eval_rule(&mut self, module: &Ref<Module>, rule: &Ref<Rule>) -> 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<Module>, rule: &Ref<Rule>) -> 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<Module>, rule: &Ref<Rule>) -> 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<dyn Extension>,
) -> Result<()> {
if let std::collections::hash_map::Entry::Vacant(v) = self.extensions.entry(path) {
v.insert((nargs, extension));
Ok(())
} else {
bail!("extension already added");
}
}
}

31
src/lib.rs
View File

@@ -260,6 +260,37 @@ pub struct QueryResults {
pub result: Vec<QueryResult>,
}
/// A user defined builtin function implementation.
///
/// It is not necessary to implement this trait directly.
pub trait Extension: FnMut(Vec<Value>) -> anyhow::Result<Value> {
/// 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<dyn 'a + Extension>
where
Self: 'a;
}
/// Automatically make matching closures a valid [`Extension`].
impl<F> Extension for F
where
F: FnMut(Vec<Value>) -> anyhow::Result<Value> + Clone,
{
fn clone_box<'a>(&self) -> Box<dyn 'a + Extension>
where
Self: 'a,
{
Box::new(self.clone())
}
}
/// Implement clone for a boxed extension using [`Extension::clone_box`].
impl<'a> Clone for Box<dyn 'a + Extension> {
fn clone(&self) -> Self {
(**self).clone_box()
}
}
/// Items in `unstable` are likely to change.
#[doc(hidden)]
pub mod unstable {

103
tests/engine/mod.rs Normal file
View File

@@ -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<Value>) -> Result<Value> {
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<Value>| {
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(())
}

1
tests/mod.rs
View File

@@ -1,6 +1,7 @@
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.
mod engine;
mod lexer;
mod parser;
mod value;