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() }