From ef549a6528b6b0fb14b33cb6af6add4f923d3970 Mon Sep 17 00:00:00 2001 From: Anand Krishnamoorthi <35780660+anakrish@users.noreply.github.com> Date: Wed, 7 Aug 2024 11:39:23 -0700 Subject: [PATCH] fix: Merge data to init document (#293) Init document is the aggregated data documen that the user has specified using multiple `add_data` calls. Each query evaluation starts of by initializing the current data to the init document. Previously `add_data` was incorrectly added to the current document, causing the added data to be lost if the addition happened after query evaluation. With this fix, scenarios where data addition may be interspersed with query evaluation calls are supported. Also provide a get_data method to obtain the (init) data document. Signed-off-by: Anand Krishnamoorthi --- src/engine.rs | 43 ++++++++++++++++++++----- src/interpreter.rs | 20 +++++++----- src/tests/interpreter/mod.rs | 29 +++++++++++++++++ tests/kata/main.rs | 61 +++++++++++++++++++++++++++++++++++- 4 files changed, 137 insertions(+), 16 deletions(-) diff --git a/src/engine.rs b/src/engine.rs index 3a31b70..9550fe6 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -239,7 +239,7 @@ impl Engine { /// # } /// ``` pub fn clear_data(&mut self) { - self.interpreter.set_data(Value::new_object()); + self.interpreter.set_init_data(Value::new_object()); self.prepared = false; } @@ -276,7 +276,40 @@ impl Engine { bail!("data must be object"); } self.prepared = false; - self.interpreter.get_data_mut().merge(data) + self.interpreter.get_init_data_mut().merge(data) + } + + /// Get the data document. + /// + /// The returned value is the data document that has been constructed using + /// one or more calls to [`Engine::add_data`]. The values of policy rules are + /// not included in the returned document. + /// + /// + /// ``` + /// # use regorus::*; + /// # fn main() -> anyhow::Result<()> { + /// let mut engine = Engine::new(); + /// + /// // If not set, data document is empty. + /// assert_eq!(engine.get_data(), Value::new_object()); + /// + /// // Merge { "x" : 1, "y" : {} } + /// assert!(engine.add_data(Value::from_json_str(r#"{ "x" : 1, "y" : {}}"#)?).is_ok()); + /// + /// // Merge { "z" : 2 } + /// assert!(engine.add_data(Value::from_json_str(r#"{ "z" : 2 }"#)?).is_ok()); + /// + /// let data = engine.get_data(); + /// assert_eq!(data["x"], Value::from(1)); + /// assert_eq!(data["y"], Value::new_object()); + /// assert_eq!(data["z"], Value::from(2)); + /// + /// # Ok(()) + /// # } + /// ``` + pub fn get_data(&self) -> Value { + self.interpreter.get_init_data().clone() } pub fn add_data_json(&mut self, data_json: &str) -> Result<()> { @@ -537,11 +570,7 @@ impl Engine { self.interpreter.set_modules(&self.modules); self.interpreter.clear_builtins_cache(); - // when the interpreter is prepared the initial data is saved - // the data will be reset to init_data each time clean_internal_evaluation_state is called - let init_data = self.interpreter.get_data_mut().clone(); - self.interpreter.set_init_data(init_data); - + // clean_internal_evaluation_state will set data to an efficient clont of use supplied init_data // Initialize the with-document with initial data values. // with-modifiers will be applied to this document. self.interpreter.init_with_document()?; diff --git a/src/interpreter.rs b/src/interpreter.rs index f4ae621..956612c 100644 --- a/src/interpreter.rs +++ b/src/interpreter.rs @@ -216,18 +216,22 @@ impl Interpreter { self.modules = modules.to_vec(); } - pub fn set_init_data(&mut self, init_data: Value) { - self.init_data = init_data; - } - - pub fn set_data(&mut self, data: Value) { - self.data = data; - } - pub fn get_data_mut(&mut self) -> &mut Value { &mut self.data } + pub fn set_init_data(&mut self, data: Value) { + self.init_data = data; + } + + pub fn get_init_data(&self) -> &Value { + &self.init_data + } + + pub fn get_init_data_mut(&mut self) -> &mut Value { + &mut self.init_data + } + pub fn set_traces(&mut self, enable_tracing: bool) { self.traces = match enable_tracing { true => Some(vec![]), diff --git a/src/tests/interpreter/mod.rs b/src/tests/interpreter/mod.rs index f986995..4d71fa2 100644 --- a/src/tests/interpreter/mod.rs +++ b/src/tests/interpreter/mod.rs @@ -422,3 +422,32 @@ fn one_yaml() -> Result<()> { fn run(path: &str) { yaml_test(path).unwrap() } + +#[test] +fn test_get_data() -> Result<()> { + let mut engine = Engine::new(); + + // Merge { "x" : 1, "y" : {} } + engine.add_data(Value::from_json_str(r#"{ "x" : 1, "y" : {}}"#)?)?; + + // Merge { "z" : 2 } + engine.add_data(Value::from_json_str(r#"{ "z" : 2 }"#)?)?; + + // Add a policy + engine.add_policy("policy.rego".to_string(), "package a".to_string())?; + + // Evaluate virtual data document. The virtual document includes all rules as well. + let v_data = engine.eval_query("data".to_string(), false)?.result[0].expressions[0] + .value + .clone(); + // There must be an empty package. + assert_eq!(v_data["a"], Value::new_object()); + + // Get the data document. + let data = engine.get_data(); + + // There must NOT be any value of `a`. + assert_eq!(data["a"], Value::Undefined); + + Ok(()) +} diff --git a/tests/kata/main.rs b/tests/kata/main.rs index 7b6f3b3..9b96235 100644 --- a/tests/kata/main.rs +++ b/tests/kata/main.rs @@ -172,6 +172,64 @@ struct Cli { generate: bool, } +fn stateful_policy_test() -> Result<()> { + // Create an engine for evaluating Rego policies. + let mut engine = regorus::Engine::new(); + + let policy = String::from( + r#" + package example + import rego.v1 + + default allow := false + + allow if { + print("data.allowed_actions = ", data.allowed_actions) + input.action in data.allowed_actions["user1"] + print("This rule should be allowed") + } + "#, + ); + + // Add policy to the engine. + engine.add_policy(String::from("policy.rego"), policy)?; + + // Evaluate first input. Expect to evaluate to false, since state is not set + engine.set_input(regorus::Value::from_json_str( + r#"{ + "action": "write" + }"#, + )?); + + let r = engine.eval_bool_query(String::from("data.example.allow"), false)?; + println!("Received result: {:?}", r); + assert_eq!(r, false); + + // Add data to engine. Set state + engine.add_data(regorus::Value::from_json_str( + r#"{ + "allowed_actions": { + "user1" : ["read", "write"] + }}"#, + )?)?; + + // Evaluate second input. Expect to evaluate to true, since state has been set now + engine.set_input(regorus::Value::from_json_str( + r#"{ + "action": "write" + }"#, + )?); + + let r = engine.eval_bool_query(String::from("data.example.allow"), false)?; + println!("Received result: {:?}", r); + assert_eq!( + r, true, + "expect result to be true since rule evaluates to true after state has been updated, per rego logs" + ); + + Ok(()) +} + fn main() -> Result<()> { let cli = Cli::parse(); run_kata_tests( @@ -179,5 +237,6 @@ fn main() -> Result<()> { &cli.name, cli.coverage, cli.generate, - ) + )?; + stateful_policy_test() }