feat: OPA v0.68.0. Engine::set_rego_v1 (#305)

Provide ability in the engine to treat subsequently loaded policies as rego.v1. Signed-off-by: Anand Krishnamoorthi <[email protected]>
microsoft · Aug 30, 2024 · af50714 · af50714
1 parent edd0ccc
commit af50714
Show file tree

Hide file tree

Showing 6 changed files with 300 additions and 129 deletions.
diff --git a/README.md b/README.md
@@ -10,7 +10,7 @@ Regorus is also
   - *cross-platform* - Written in platform-agnostic Rust.
   - *no_std compatible* - Regorus can be used in `no_std` environments too. Most of the builtins are supported.
   - *current* - We strive to keep Regorus up to date with latest OPA release. Regorus supports `import rego.v1`.
-  - *compliant* - Regorus is mostly compliant with the latest [OPA release v0.67.0](https://github.com/open-policy-agent/opa/releases/tag/v0.67.0). See [OPA Conformance](#opa-conformance) for details. Note that while we behaviorally produce the same results, we don't yet support all the builtins.
+  - *compliant* - Regorus is mostly compliant with the latest [OPA release v0.68.0](https://github.com/open-policy-agent/opa/releases/tag/v0.68.0). See [OPA Conformance](#opa-conformance) for details. Note that while we behaviorally produce the same results, we don't yet support all the builtins.
   - *extensible* - Extend the Rego language by implementing custom stateful builtins in Rust.
     See [add_extension](https://github.com/microsoft/regorus/blob/fc68bf9c8bea36427dae9401a7d1f6ada771f7ab/src/engine.rs#L352).
     Support for extensibility using other languages coming soon.
@@ -99,7 +99,7 @@ $ cargo build -r --example regorus --no-default-features; strip target/release/e
 -rwxr-xr-x  1 anand  staff   1.9M May 11 22:04 target/release/examples/regorus*
 ```
 
-Regorus passes the [OPA v0.67.0 test-suite](https://www.openpolicyagent.org/docs/latest/ir/#test-suite) barring a few
+Regorus passes the [OPA v0.68.0 test-suite](https://www.openpolicyagent.org/docs/latest/ir/#test-suite) barring a few
 builtins. See [OPA Conformance](#opa-conformance) below.
 
 ## Bindings
@@ -276,7 +276,7 @@ Benchmark 1: opa eval -b tests/aci -d tests/aci/data.json -i tests/aci/input.jso
 ```
 ## OPA Conformance
 
-Regorus has been verified to be compliant with [OPA v0.67.0](https://github.com/open-policy-agent/opa/releases/tag/v0.67.0)
+Regorus has been verified to be compliant with [OPA v0.68.0](https://github.com/open-policy-agent/opa/releases/tag/v0.68.0)
 using a [test driver](https://github.com/microsoft/regorus/blob/main/tests/opa.rs) that loads and runs the OPA testsuite using Regorus, and verifies that expected outputs are produced.
 
 The test driver can be invoked by running:

diff --git a/examples/regorus.rs b/examples/regorus.rs
@@ -33,6 +33,7 @@ fn add_policy_from_file(engine: &mut regorus::Engine, path: String) -> Result<St
     engine.add_policy(path.clone(), read_file(&path)?)
 }
 
+#[allow(clippy::too_many_arguments)]
 fn rego_eval(
     bundles: &[String],
     files: &[String],
@@ -41,6 +42,7 @@ fn rego_eval(
     enable_tracing: bool,
     non_strict: bool,
     #[cfg(feature = "coverage")] coverage: bool,
+    v1: bool,
 ) -> Result<()> {
     // Create engine.
     let mut engine = regorus::Engine::new();
@@ -50,6 +52,8 @@ fn rego_eval(
     #[cfg(feature = "coverage")]
     engine.set_enable_coverage(coverage);
 
+    engine.set_rego_v1(v1);
+
     // Load files from given bundles.
     for dir in bundles.iter() {
         let entries =
@@ -233,6 +237,10 @@ enum RegorusCommand {
         #[cfg(feature = "coverage")]
         #[arg(long, short)]
         coverage: bool,
+
+        /// Turn on rego.v1
+        #[arg(long)]
+        v1: bool,
     },
 
     /// Tokenize a Rego policy.
@@ -274,6 +282,7 @@ fn main() -> Result<()> {
             non_strict,
             #[cfg(feature = "coverage")]
             coverage,
+            v1,
         } => rego_eval(
             &bundles,
             &data,
@@ -283,6 +292,7 @@ fn main() -> Result<()> {
             non_strict,
             #[cfg(feature = "coverage")]
             coverage,
+            v1,
         ),
         RegorusCommand::Lex { file, verbose } => rego_lex(file, verbose),
         RegorusCommand::Parse { file } => rego_parse(file),

diff --git a/src/engine.rs b/src/engine.rs
@@ -20,6 +20,7 @@ pub struct Engine {
     modules: Vec<Ref<Module>>,
     interpreter: Interpreter,
     prepared: bool,
+    rego_v1: bool,
 }
 
 /// Create a default engine.
@@ -36,9 +37,35 @@ impl Engine {
             modules: vec![],
             interpreter: Interpreter::new(),
             prepared: false,
+            rego_v1: false,
         }
     }
 
+    /// Turn rego.v1 on/off for subsequently added policies.
+    ///
+    /// Explicit import rego.v1 is not needed if set.
+    ///
+    /// ```
+    /// # use regorus::*;
+    /// # fn main() -> anyhow::Result<()> {
+    /// let mut engine = Engine::new();
+    ///
+    /// engine.set_rego_v1(true);
+    /// engine.add_policy(
+    ///    "test.rego".to_string(),
+    ///    r#"
+    ///    package test
+    ///    allow if true # if keyword is automatically imported
+    ///    "#.to_string())?;
+    ///
+    /// # Ok(())
+    /// # }
+    /// ```
+    ///
+    pub fn set_rego_v1(&mut self, rego_v1: bool) {
+        self.rego_v1 = rego_v1;
+    }
+
     /// Add a policy.
     ///
     /// The policy file will be parsed and converted to AST representation.
@@ -67,7 +94,7 @@ impl Engine {
     ///
     pub fn add_policy(&mut self, path: String, rego: String) -> Result<String> {
         let source = Source::from_contents(path, rego)?;
-        let mut parser = Parser::new(&source)?;
+        let mut parser = self.make_parser(&source)?;
         let module = Ref::new(parser.parse()?);
         self.modules.push(module.clone());
         // if policies change, interpreter needs to be prepared again
@@ -98,7 +125,7 @@ impl Engine {
     #[cfg_attr(docsrs, doc(cfg(feature = "std")))]
     pub fn add_policy_from_file<P: AsRef<std::path::Path>>(&mut self, path: P) -> Result<String> {
         let source = Source::from_file(path)?;
-        let mut parser = Parser::new(&source)?;
+        let mut parser = self.make_parser(&source)?;
         let module = Ref::new(parser.parse()?);
         self.modules.push(module.clone());
         // if policies change, interpreter needs to be prepared again
@@ -428,7 +455,7 @@ impl Engine {
 
         // Parse the query.
         let query_source = Source::from_contents("<query.rego>".to_string(), query)?;
-        let mut parser = Parser::new(&query_source)?;
+        let mut parser = self.make_parser(&query_source)?;
         let query_node = parser.parse_user_query()?;
         if query_node.span.text() == "data" {
             self.eval_modules(enable_tracing)?;
@@ -545,7 +572,7 @@ impl Engine {
 
         // Parse the query.
         let query_source = Source::from_contents("<query.rego>".to_string(), query)?;
-        let mut parser = Parser::new(&query_source)?;
+        let mut parser = self.make_parser(&query_source)?;
         let query_node = parser.parse_user_query()?;
         let query_schedule = Analyzer::new().analyze_query_snippet(&self.modules, &query_node)?;
         self.interpreter.eval_user_query(
@@ -870,4 +897,12 @@ impl Engine {
 
         serde_json::to_string_pretty(&ast).map_err(anyhow::Error::msg)
     }
+
+    fn make_parser<'a>(&self, source: &'a Source) -> Result<Parser<'a>> {
+        let mut parser = Parser::new(source)?;
+        if self.rego_v1 {
+            parser.enable_rego_v1()?;
+        }
+        Ok(parser)
+    }
 }
diff --git a/src/parser.rs b/src/parser.rs
@@ -40,6 +40,18 @@ impl<'source> Parser<'source> {
         })
     }
 
+    pub fn enable_rego_v1(&mut self) -> Result<()> {
+        self.turn_on_rego_v1(self.tok.1.clone())
+    }
+
+    fn turn_on_rego_v1(&mut self, span: Span) -> Result<()> {
+        self.rego_v1 = true;
+        for kw in FUTURE_KEYWORDS {
+            self.set_future_keyword(kw, &span)?;
+        }
+        Ok(())
+    }
+
     pub fn token_text(&self) -> &str {
         match self.tok.0 {
             TokenKind::Symbol | TokenKind::Number | TokenKind::Ident | TokenKind::Eof => {
@@ -1648,10 +1660,7 @@ impl<'source> Parser<'source> {
 
             let is_future_kw =
                 if comps.len() == 2 && comps[0].text() == "rego" && comps[1].text() == "v1" {
-                    self.rego_v1 = true;
-                    for kw in FUTURE_KEYWORDS {
-                        self.set_future_keyword(kw, &span)?;
-                    }
+                    self.turn_on_rego_v1(span.clone())?;
                     true
                 } else {
                     self.handle_import_future_keywords(&comps)?