Provide ability to get JSON representation of policy AST

closes #265

Signed-off-by: Anand Krishnamoorthi <[email protected]>

Cargo.toml

@@ -27,6 +27,7 @@ doctest = false
 default = ["full-opa", "arc"]
 
 arc = ["scientific/arc"]
+ast = []
 base64 = ["dep:data-encoding"]
 base64url = ["dep:data-encoding"]
 coverage = []

bindings/ffi/Cargo.toml

@@ -13,7 +13,8 @@ regorus = { path  = "../..", default-features = false }
 serde_json = "1.0.113"
 
 [features]
-default = ["std", "coverage", "regorus/arc", "regorus/full-opa"]
+default = ["ast", "std", "coverage", "regorus/arc", "regorus/full-opa"]
+ast = ["regorus/ast"]
 std = ["regorus/std"]
 coverage = ["regorus/coverage"]
 custom_allocator = []

bindings/ffi/src/lib.rs

@@ -411,6 +411,23 @@ pub extern "C" fn regorus_engine_take_prints(engine: *mut RegorusEngine) -> Rego
     }
 }
 
+/// Get AST of policies.
+///
+/// See https://docs.rs/regorus/latest/regorus/coverage/struct.Engine.html#method.get_ast_as_json
+#[no_mangle]
+#[cfg(feature = "ast")]
+pub extern "C" fn regorus_engine_get_ast_as_json(engine: *mut RegorusEngine) -> RegorusResult {
+    let output = || -> Result<String> { to_ref(&engine)?.engine.get_ast_as_json()? }();
+    match output {
+        Ok(out) => RegorusResult {
+            status: RegorusStatus::RegorusStatusOk,
+            output: to_c_str(out),
+            error_message: std::ptr::null_mut(),
+        },
+        Err(e) => to_regorus_result(Err(e)),
+    }
+}
+
 #[cfg(feature = "custom_allocator")]
 extern "C" {
     fn regorus_aligned_alloc(alignment: usize, size: usize) -> *mut u8;

bindings/java/Cargo.toml

@@ -11,7 +11,9 @@ keywords = ["interpreter", "opa", "policy-as-code", "rego"]
 crate-type = ["cdylib"]
 
 [features]
-default = ["regorus/std", "regorus/full-opa"]
+default = ["ast", "coverage", "regorus/std", "regorus/full-opa"]
+coverage = ["regorus/coverage"]
+ast = ["regorus/ast"]
 
 [dependencies]
 anyhow = "1.0"

bindings/java/src/lib.rs

@@ -205,6 +205,7 @@ pub extern "system" fn Java_com_microsoft_regorus_Engine_nativeEvalRule(
 }
 
 #[no_mangle]
+#[cfg(feature = "coverage")]
 pub extern "system" fn Java_com_microsoft_regorus_Engine_nativeSetEnableCoverage(
     env: JNIEnv,
     _class: JClass,
@@ -219,6 +220,7 @@ pub extern "system" fn Java_com_microsoft_regorus_Engine_nativeSetEnableCoverage
 }
 
 #[no_mangle]
+#[cfg(feature = "coverage")]
 pub extern "system" fn Java_com_microsoft_regorus_Engine_nativeGetCoverageReport(
     env: JNIEnv,
     _class: JClass,
@@ -238,6 +240,7 @@ pub extern "system" fn Java_com_microsoft_regorus_Engine_nativeGetCoverageReport
 }
 
 #[no_mangle]
+#[cfg(feature = "coverage")]
 pub extern "system" fn Java_com_microsoft_regorus_Engine_nativeGetCoverageReportPretty(
     env: JNIEnv,
     _class: JClass,
@@ -257,6 +260,7 @@ pub extern "system" fn Java_com_microsoft_regorus_Engine_nativeGetCoverageReport
 }
 
 #[no_mangle]
+#[cfg(feature = "coverage")]
 pub extern "system" fn Java_com_microsoft_regorus_Engine_nativeClearCoverageData(
     env: JNIEnv,
     _class: JClass,
@@ -302,6 +306,26 @@ pub extern "system" fn Java_com_microsoft_regorus_Engine_nativeTakePrints(
     }
 }
 
+#[no_mangle]
+#[cfg(feature = "ast")]
+pub extern "system" fn Java_com_microsoft_regorus_Engine_getAstAsJson(
+    env: JNIEnv,
+    _class: JClass,
+    engine_ptr: jlong,
+) -> jstring {
+    let res = throw_err(env, |env| {
+        let engine = unsafe { &mut *(engine_ptr as *mut Engine) };
+        let ast = engine.get_ast_as_json()?;
+        let output = env.new_string(&ast)?;
+        Ok(output.into_raw())
+    });
+
+    match res {
+        Ok(val) => val,
+        Err(_) => JObject::null().into_raw(),
+    }
+}
+
 #[no_mangle]
 pub extern "system" fn Java_com_microsoft_regorus_Engine_nativeDestroyEngine(
     _env: JNIEnv,

bindings/python/Cargo.toml

@@ -12,7 +12,9 @@ keywords = ["interpreter", "opa", "policy-as-code", "rego"]
 crate-type = ["cdylib"]
 
 [features]
-default = ["regorus/std", "regorus/full-opa"]
+default = ["ast", "coverage", "regorus/std", "regorus/full-opa"]
+ast = ["regorus/ast"]
+coverage = ["regorus/coverage"]
 
 [dependencies]
 anyhow = "1.0"

bindings/python/src/lib.rs

@@ -312,19 +312,22 @@ impl Engine {
 
     /// Get coverage report as json.
     ///
+    #[cfg(feature = "coverage")]
     pub fn get_coverage_report_as_json(&self) -> Result<String> {
         let report = self.engine.get_coverage_report()?;
         serde_json::to_string_pretty(&report).map_err(|e| anyhow!("{e}"))
     }
 
     /// Get coverage report as pretty printable string.
     ///
+    #[cfg(feature = "coverage")]
     pub fn get_coverage_report_pretty(&self) -> Result<String> {
         self.engine.get_coverage_report()?.to_string_pretty()
     }
 
     /// Clear coverage data.
     ///
+    #[cfg(feature = "coverage")]
     pub fn clear_coverage_data(&mut self) {
         self.engine.clear_coverage_data();
     }
@@ -350,6 +353,13 @@ impl Engine {
             engine: self.engine.clone(),
         }
     }
+
+    /// Get AST of policies.
+    ///
+    #[cfg(feature = "ast")]
+    pub fn get_ast_as_json(&self) -> Result<String> {
+        self.engine.get_ast_as_json()
+    }
 }
 
 #[pymodule]

bindings/ruby/ext/regorusrb/Cargo.toml

@@ -10,10 +10,12 @@ crate-type = ["cdylib"]
 path = "src/lib.rs"
 
 [features]
-default = ["regorus/std", "regorus/full-opa"]
+default = ["ast", "coverage", "regorus/std", "regorus/full-opa"]
+ast = ["regorus/ast"]
+coverage = ["regorus/coverage"]
 
 [dependencies]
 magnus = { version = "0.6.4" }
-regorus = { git = "https://github.com/microsoft/regorus", default-features = false, features = ["arc"] }
+regorus = { path = "../../../..", default-features = false, features = ["arc"] }
 serde_json = "1.0.117"
 serde_magnus = "0.8.1"

bindings/ruby/ext/regorusrb/src/lib.rs

@@ -192,11 +192,13 @@ impl Engine {
         Ok(self.engine.borrow_mut().eval_deny_query(query, false))
     }
 
+    #[cfg(feature = "coverage")]
     fn set_enable_coverage(&self, enable: bool) -> Result<(), Error> {
         self.engine.borrow_mut().set_enable_coverage(enable);
         Ok(())
     }
 
+    #[cfg(feature = "coverage")]
     fn get_coverage_report_as_json(&self) -> Result<String, Error> {
         let report = self
             .engine
@@ -217,6 +219,7 @@ impl Engine {
         })
     }
 
+    #[cfg(feature = "coverage")]
     fn get_coverage_report_pretty(&self) -> Result<String, Error> {
         let report = self
             .engine
@@ -237,6 +240,7 @@ impl Engine {
         })
     }
 
+    #[cfg(feature = "coverage")]
     fn clear_coverage_data(&self) -> Result<(), Error> {
         self.engine.borrow_mut().clear_coverage_data();
         Ok(())
@@ -256,6 +260,14 @@ impl Engine {
             )
         })
     }
+
+    #[cfg(feature = "ast")]
+    fn get_ast_as_json(&self) -> Result<String, Error> {
+        self.engine
+            .borrow()
+            .get_ast_as_json()
+            .map_err(|e| Error::new(runtime_error(), format!("Failed to get ast: {e}")))
+    }
 }
 
 #[magnus::init]

bindings/wasm/Cargo.toml

@@ -11,7 +11,9 @@ keywords = ["interpreter", "opa", "policy-as-code", "rego"]
 crate-type = ["cdylib"]
 
 [features]
-default = ["regorus/std", "regorus/full-opa"]
+default = ["ast", "coverage", "regorus/std", "regorus/full-opa"]
+ast = ["regorus/ast"]
+coverage = ["regorus/coverage"]
 
 [dependencies]
 regorus = { path  = "../..", default-features = false, features = ["arc"] }

bindings/wasm/src/lib.rs

@@ -131,13 +131,15 @@ impl Engine {
     ///
     /// See https://docs.rs/regorus/latest/regorus/struct.Engine.html#method.set_enable_coverage
     /// * `b`: Whether to enable gathering coverage or not.
+    #[cfg(feature = "coverage")]
     pub fn setEnableCoverage(&mut self, enable: bool) {
         self.engine.set_enable_coverage(enable)
     }
 
     /// Get the coverage report as json.
     ///
     /// See https://docs.rs/regorus/latest/regorus/struct.Engine.html#method.get_coverage_report
+    #[cfg(feature = "coverage")]
     pub fn getCoverageReport(&self) -> Result<String, JsValue> {
         let report = self
             .engine
@@ -149,20 +151,30 @@ impl Engine {
     /// Clear gathered coverage data.
     ///
     /// See https://docs.rs/regorus/latest/regorus/struct.Engine.html#method.clear_coverage_data
+    #[cfg(feature = "coverage")]
     pub fn clearCoverageData(&mut self) {
         self.engine.clear_coverage_data()
     }
 
     /// Get ANSI color coded coverage report.
     ///
     /// See https://docs.rs/regorus/latest/regorus/coverage/struct.Report.html#method.to_string_pretty
+    #[cfg(feature = "coverage")]
     pub fn getCoverageReportPretty(&self) -> Result<String, JsValue> {
         let report = self
             .engine
             .get_coverage_report()
             .map_err(error_to_jsvalue)?;
         report.to_string_pretty().map_err(error_to_jsvalue)
     }
+
+    /// Get AST of policies.
+    ///
+    /// See https://docs.rs/regorus/latest/regorus/struct.Engine.html#method.get_ast_as_json
+    #[cfg(feature = "ast")]
+    pub fn getAstAsJson(&self) -> Result<String, JsValue> {
+        self.engine.get_ast_as_json().map_err(error_to_jsvalue)
+    }
 }
 
 #[cfg(test)]

examples/regorus.rs

@@ -170,8 +170,40 @@ fn rego_parse(file: String) -> Result<()> {
     Ok(())
 }
 
+#[allow(unused_variables)]
+fn rego_ast(file: String) -> Result<()> {
+    #[cfg(feature = "ast")]
+    {
+        // Create engine.
+        let mut engine = regorus::Engine::new();
+
+        // Create source.
+        #[cfg(feature = "std")]
+        engine.add_policy_from_file(file)?;
+
+        #[cfg(not(feature = "std"))]
+        engine.add_policy(file.clone(), read_file(&file)?)?;
+
+        let ast = engine.get_ast_as_json()?;
+
+        println!("{ast}");
+        Ok(())
+    }
+
+    #[cfg(not(feature = "ast"))]
+    {
+        bail!("`ast` feature must be enabled");
+    }
+}
+
 #[derive(clap::Subcommand)]
 enum RegorusCommand {
+    /// Parse a Rego policy and dump AST.
+    Ast {
+        /// Rego policy file.
+        file: String,
+    },
+
     /// Evaluate a Rego Query.
     Eval {
         /// Directories containing Rego files.
@@ -254,5 +286,6 @@ fn main() -> Result<()> {
         ),
         RegorusCommand::Lex { file, verbose } => rego_lex(file, verbose),
         RegorusCommand::Parse { file } => rego_parse(file),
+        RegorusCommand::Ast { file } => rego_ast(file),
     }
 }