From 545a93cc8b55bbc0e6bfdd07406f883444ac1422 Mon Sep 17 00:00:00 2001 From: Anand Krishnamoorthi Date: Wed, 7 Aug 2024 11:20:15 -0700 Subject: [PATCH] fix: Merge data to init document 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 | 16 ++++++---- src/tests/interpreter/mod.rs | 29 +++++++++++++++++ tests/kata/main.rs | 61 +++++++++++++++++++++++++++++++++++- 4 files changed, 135 insertions(+), 14 deletions(-) diff --git a/src/engine.rs b/src/engine.rs index 3a31b703..9550fe69 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 f4ae621e..956612c3 100644 --- a/src/interpreter.rs +++ b/src/interpreter.rs @@ -216,16 +216,20 @@ impl Interpreter { self.modules = modules.to_vec(); } - pub fn set_init_data(&mut self, init_data: Value) { - self.init_data = init_data; + pub fn get_data_mut(&mut self) -> &mut Value { + &mut self.data } - pub fn set_data(&mut self, data: Value) { - self.data = data; + pub fn set_init_data(&mut self, data: Value) { + self.init_data = data; } - pub fn get_data_mut(&mut self) -> &mut Value { - &mut self.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) { diff --git a/src/tests/interpreter/mod.rs b/src/tests/interpreter/mod.rs index f9869956..4d71fa27 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 7b6f3b30..9b962352 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() }