Initial implementation of policy coverage (#146)

Signed-off-by: Anand Krishnamoorthi <anakrish@microsoft.com>

.github/workflows/pr.yml

@@ -23,27 +23,29 @@ jobs:
     - name: Format Check
       run: cargo fmt --check
     - name: Build
-      run: cargo build --verbose
-    - name: Build Tests
-      run: cargo build --all-targets --verbose
-    - name: Clippy
-      run: cargo clippy --all-targets --no-deps -- -Dwarnings
+      run: cargo build -r --verbose
+    - name: Doc Tests
+      run: cargo test -r --doc
     - name: Run tests
       run: cargo test -r --verbose
-    - name: Build (MUSL)
-      run: cargo build --verbose --all-targets --target x86_64-unknown-linux-musl
-    - name: Run tests (MUSL)
-      run: cargo test -r --verbose --target x86_64-unknown-linux-musl
     - name: Run tests (ACI)
       run: cargo test -r --test aci
     - name: Run tests (OPA Conformance)
       run: >-
         cargo test -r --test opa --features opa-testutil  -- $(tr '\n' ' ' < tests/opa.passing)
-
-    # - name: Install wasm-pack
-    #   run: curl https://rustwasm.github.io/wasm-pack/installer/init.sh -sSf | sh
-
-    # - name: Run wasm binding tests
-    #   run: |
-    #     cd bindings/wasm
-    #     wasm-pack test --node -r
+    - name: Run tests (coverage)
+      run: cargo test -r --verbose --features "coverage"
+    - name: Run tests (ACI coverage)
+      run: cargo test -r --test aci --features "coverage"
+    - name: Run tests (OPA Conformance Coverage)
+      run: >-
+        cargo test -r --test opa --features "opa-testutil,coverage"  -- $(tr '\n' ' ' < tests/opa.passing)
+    - name: Build (MUSL)
+      run: cargo build --verbose --all-targets --target x86_64-unknown-linux-musl
+    - name: Run tests (MUSL)
+      run: cargo test -r --verbose --target x86_64-unknown-linux-musl
+    - name: Run tests (MUSL ACI)
+      run: cargo test -r --test aci --features "coverage" --target x86_64-unknown-linux-musl
+    - name: Run tests (MUSL OPA Conformance Coverage)
+      run: >-
+        cargo test -r --test opa --features "opa-testutil,coverage" --target x86_64-unknown-linux-musl  -- $(tr '\n' ' ' < tests/opa.passing)

Cargo.toml

@@ -18,12 +18,16 @@ keywords = ["interpreter", "opa", "policy-as-code", "rego"]
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
+[lib]
+doctest = false
+
 [features]
 default = ["full-opa", "arc"]
 
 arc = ["scientific/arc"]
 base64 = ["dep:data-encoding"]
 base64url = ["dep:data-encoding"]
+coverage = []
 crypto = ["dep:constant_time_eq", "dep:hmac", "dep:hex", "dep:md-5", "dep:sha1", "dep:sha2"]
 deprecated = []
 hex = ["dep:data-encoding"]
@@ -97,6 +101,7 @@ jsonwebtoken = { version = "9.2.0", optional = true }
 itertools = "0.12.1"
 
 [dev-dependencies]
+cfg-if = "1.0.0"
 clap = { version = "4.4.7", features = ["derive"] }
 colored-diff = "0.2.3"
 serde_yaml = "0.9.16"

examples/regorus.rs

@@ -10,6 +10,7 @@ fn rego_eval(
     query: String,
     enable_tracing: bool,
     non_strict: bool,
+    #[cfg(feature = "coverage")] coverage: bool,
 ) -> Result<()> {
     // Create engine.
     let mut engine = regorus::Engine::new();
@@ -70,6 +71,28 @@ fn rego_eval(
     let results = engine.eval_query(query, enable_tracing)?;
     println!("{}", serde_json::to_string_pretty(&results)?);
 
+    #[cfg(feature = "coverage")]
+    if coverage {
+        println!("\n\nCOVERAGE REPORT");
+        // Fetch coverage report.
+        let report = engine.get_coverage_report()?;
+        for file in report.files.into_iter() {
+            if file.uncovered.is_empty() {
+                println!("{} has full coverage", file.path);
+                continue;
+            }
+
+            println!("{}:", file.path);
+            for (line, code) in file.code.split('\n').enumerate() {
+                if file.uncovered.contains(&(line as u32 + 1)) {
+                    println!("\x1b[31m {line:4}  {code}\x1b[0m");
+                } else {
+                    println!(" {line:4}  {code}");
+                }
+            }
+        }
+    }
+
     Ok(())
 }
 
@@ -140,6 +163,11 @@ enum RegorusCommand {
         // Non strict execution
         #[arg(long, short)]
         non_strict: bool,
+
+        // Display coverage information
+        #[cfg(feature = "coverage")]
+        #[arg(long, short)]
+        coverage: bool,
     },
 
     /// Tokenize a Rego policy.
@@ -183,7 +211,18 @@ fn main() -> Result<()> {
             query,
             trace,
             non_strict,
-        } => rego_eval(&bundles, &data, input, query, trace, non_strict),
+            #[cfg(feature = "coverage")]
+            coverage,
+        } => rego_eval(
+            &bundles,
+            &data,
+            input,
+            query,
+            trace,
+            non_strict,
+            #[cfg(feature = "coverage")]
+            coverage,
+        ),
         RegorusCommand::Lex { file, verbose } => rego_lex(file, verbose),
         RegorusCommand::Parse { file } => rego_parse(file),
     }

scripts/pre-push

@@ -9,10 +9,18 @@ if [ -f Cargo.toml ]; then
    dir=$(dirname "${BASH_SOURCE[0]}")
    "$dir/pre-commit"
 
+   # Ensure that the public API works
+   cargo test -r --doc
+   
    # Ensure that all tests pass
    cargo test -r
    cargo test -r --test aci
 
    # Ensure that OPA conformance tests don't regress.
    cargo test -r --features opa-testutil --test opa -- $(tr '\n' ' ' < tests/opa.passing)
+
+   # Run the same tests using "coverage" build
+   cargo test -r --features "coverage"
+   cargo test -r --test aci --features "coverage"
+   cargo test -r --features "coverage,opa-testutil" --test opa -- $(tr '\n' ' ' < tests/opa.passing)
 fi

src/engine.rs

@@ -453,4 +453,9 @@ impl Engine {
     ) -> Result<()> {
         self.interpreter.add_extension(path, nargs, extension)
     }
+
+    #[cfg(feature = "coverage")]
+    pub fn get_coverage_report(&self) -> Result<crate::coverage::Report> {
+        self.interpreter.get_coverage_report()
+    }
 }

src/interpreter.rs

@@ -66,6 +66,9 @@ pub struct Interpreter {
     strict_builtin_errors: bool,
     imports: BTreeMap<String, Ref<Expr>>,
     extensions: HashMap<String, (u8, Rc<Box<dyn Extension>>)>,
+
+    #[cfg(feature = "coverage")]
+    coverage: HashMap<Source, Vec<bool>>,
 }
 
 impl Default for Interpreter {
@@ -177,6 +180,9 @@ impl Interpreter {
             strict_builtin_errors: true,
             imports: BTreeMap::default(),
             extensions: HashMap::new(),
+
+            #[cfg(feature = "coverage")]
+            coverage: HashMap::new(),
         }
     }
 
@@ -2583,6 +2589,31 @@ impl Interpreter {
             expr.span().text()
         );
 
+        #[cfg(feature = "coverage")]
+        {
+            let span = expr.span();
+            let source = &span.source;
+            let line = span.line as usize;
+            if line > 0 {
+                // Check if coverage table already exists for source.
+                match self.coverage.get_mut(source) {
+                    Some(c) => {
+                        // Ensure that current line is valid.
+                        if c.len() < line + 1 {
+                            c.resize(line + 1, false);
+                        }
+                        c[line] = true;
+                    }
+                    _ => {
+                        // Create new table.
+                        let mut c = vec![false; line + 1];
+                        c[line] = true;
+                        self.coverage.insert(source.clone(), c);
+                    }
+                }
+            }
+        }
+
         match expr.as_ref() {
             Expr::Null(_) => Ok(Value::Null),
             Expr::True(_) => Ok(Value::Bool(true)),
@@ -3475,4 +3506,123 @@ impl Interpreter {
             bail!("extension already added");
         }
     }
+
+    #[cfg(feature = "coverage")]
+    fn gather_uncovered_lines_in_query(
+        &self,
+        query: &Ref<Query>,
+        covered: &Vec<bool>,
+        uncovered: &mut BTreeSet<u32>,
+    ) -> Result<()> {
+        for stmt in &query.stmts {
+            // TODO: with mods
+            match &stmt.literal {
+                Literal::SomeVars { .. } => (),
+                Literal::SomeIn {
+                    value, collection, ..
+                } => {
+                    self.gather_uncovered_lines_in_expr(value, covered, uncovered)?;
+                    self.gather_uncovered_lines_in_expr(collection, covered, uncovered)?;
+                }
+                Literal::Expr { expr, .. } | Literal::NotExpr { expr, .. } => {
+                    self.gather_uncovered_lines_in_expr(expr, covered, uncovered)?;
+                }
+                Literal::Every { domain, query, .. } => {
+                    self.gather_uncovered_lines_in_expr(domain, covered, uncovered)?;
+                    self.gather_uncovered_lines_in_query(query, covered, uncovered)?;
+                }
+            }
+        }
+        Ok(())
+    }
+
+    #[cfg(feature = "coverage")]
+    fn gather_uncovered_lines_in_expr(
+        &self,
+        expr: &Ref<Expr>,
+        covered: &Vec<bool>,
+        uncovered: &mut BTreeSet<u32>,
+    ) -> Result<()> {
+        use Expr::*;
+        traverse(expr, &mut |e| {
+            Ok(match e.as_ref() {
+                ArrayCompr { query, .. } | SetCompr { query, .. } | ObjectCompr { query, .. } => {
+                    self.gather_uncovered_lines_in_query(query, covered, uncovered)?;
+                    false
+                }
+                _ => {
+                    let line = e.span().line as usize;
+                    if line >= covered.len() || !covered[line] {
+                        uncovered.insert(line as u32);
+                    }
+                    true
+                }
+            })
+        })?;
+        Ok(())
+    }
+
+    #[cfg(feature = "coverage")]
+    pub fn get_coverage_report(&self) -> Result<crate::coverage::Report> {
+        let mut report = crate::coverage::Report::default();
+
+        for module in self.modules.iter() {
+            let span = module.package.refr.span();
+
+            // Get coverage information for the module.
+            let Some(covered) = self.coverage.get(&span.source) else {
+                continue;
+            };
+
+            let mut uncovered = BTreeSet::new();
+
+            // Loop through each rule and figure out the lines that were not coverd.
+            for rule in &module.policy {
+                match rule.as_ref() {
+                    Rule::Spec { head, bodies, .. } => {
+                        match head {
+                            RuleHead::Compr { assign, .. } | RuleHead::Func { assign, .. } => {
+                                if let Some(a) = assign {
+                                    self.gather_uncovered_lines_in_expr(
+                                        &a.value,
+                                        covered,
+                                        &mut uncovered,
+                                    )?;
+                                }
+                            }
+                            RuleHead::Set { key, .. } => {
+                                if let Some(k) = key {
+                                    self.gather_uncovered_lines_in_expr(
+                                        k,
+                                        covered,
+                                        &mut uncovered,
+                                    )?;
+                                }
+                            }
+                        }
+                        for b in bodies {
+                            self.gather_uncovered_lines_in_query(
+                                &b.query,
+                                covered,
+                                &mut uncovered,
+                            )?;
+                        }
+                    }
+                    Rule::Default { value, .. } => {
+                        self.gather_uncovered_lines_in_expr(value, covered, &mut uncovered)?;
+                    }
+                }
+            }
+
+            let file = crate::coverage::PolicyFile {
+                path: span.source.file().clone(),
+                code: span.source.contents().clone(),
+                uncovered,
+            };
+
+            report.files.push(file);
+        }
+
+        Ok(report)
+    }
 }

src/lexer.rs

@@ -6,6 +6,7 @@ use core::iter::Peekable;
 use core::str::CharIndices;
 
 use std::convert::AsRef;
+use std::hash::{Hash, Hasher};
 use std::path::Path;
 
 use crate::Rc;
@@ -25,6 +26,38 @@ pub struct Source {
     src: Rc<SourceInternal>,
 }
 
+impl std::cmp::Ord for Source {
+    fn cmp(&self, other: &Source) -> std::cmp::Ordering {
+        Rc::as_ptr(&self.src).cmp(&Rc::as_ptr(&other.src))
+    }
+}
+
+impl std::cmp::PartialOrd for Source {
+    fn partial_cmp(&self, other: &Source) -> Option<std::cmp::Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl std::cmp::PartialEq for Source {
+    fn eq(&self, other: &Source) -> bool {
+        Rc::as_ptr(&self.src) == Rc::as_ptr(&other.src)
+    }
+}
+
+impl std::cmp::Eq for Source {}
+
+impl Hash for Source {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        Rc::as_ptr(&self.src).hash(state)
+    }
+}
+
+impl Debug for Source {
+    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), std::fmt::Error> {
+        self.src.file.fmt(f)
+    }
+}
+
 #[derive(Clone)]
 pub struct SourceStr {
     source: Source,

src/lib.rs

@@ -339,6 +339,21 @@ impl std::fmt::Debug for dyn Extension {
     }
 }
 
+#[cfg(feature = "coverage")]
+pub mod coverage {
+    #[derive(Default, serde::Serialize, serde::Deserialize, Eq, PartialEq)]
+    pub struct PolicyFile {
+        pub path: String,
+        pub code: String,
+        pub uncovered: std::collections::BTreeSet<u32>,
+    }
+
+    #[derive(Default, serde::Serialize, serde::Deserialize, Eq, PartialEq)]
+    pub struct Report {
+        pub files: Vec<PolicyFile>,
+    }
+}
+
 /// Items in `unstable` are likely to change.
 #[doc(hidden)]
 pub mod unstable {

tests/aci/main.rs

@@ -108,6 +108,70 @@ fn run_aci_tests(dir: &Path) -> Result<()> {
     Ok(())
 }
 
+#[cfg(feature = "coverage")]
+fn run_aci_tests_coverage(dir: &Path) -> Result<()> {
+    let mut engine = Engine::new();
+
+    let mut added = std::collections::BTreeSet::new();
+
+    for entry in WalkDir::new(dir)
+        .sort_by_file_name()
+        .into_iter()
+        .filter_map(|e| e.ok())
+    {
+        let path = entry.path();
+        if !path.to_string_lossy().ends_with(".yaml") {
+            continue;
+        }
+
+        let yaml = std::fs::read(&path)?;
+        let yaml = String::from_utf8_lossy(&yaml);
+        let test: YamlTest = serde_yaml::from_str(&yaml)?;
+
+        for case in &test.cases {
+            for (idx, rego) in case.modules.iter().enumerate() {
+                if rego.ends_with(".rego") {
+                    let path = dir.join(rego);
+                    let path = path.to_str().expect("not a valid path");
+                    let path = path.to_string();
+                    if !added.contains(&path) {
+                        engine.add_policy_from_file(path.to_string())?;
+                        added.insert(path);
+                    }
+                } else {
+                    engine.add_policy(format!("rego{idx}.rego"), rego.clone())?;
+                }
+            }
+
+            engine.clear_data();
+            engine.add_data(case.data.clone())?;
+            engine.set_input(case.input.clone());
+            let _query_results = engine.eval_query(case.query.clone(), true)?;
+        }
+    }
+
+    println!("\n\nCOVERAGE REPORT");
+    // Fetch coverage report.
+    let report = engine.get_coverage_report()?;
+    for file in report.files.into_iter() {
+        if file.uncovered.is_empty() {
+            println!("{} has full coverage", file.path);
+            continue;
+        }
+
+        println!("{}:", file.path);
+        for (line, code) in file.code.split('\n').enumerate() {
+            if file.uncovered.contains(&(line as u32 + 1)) {
+                println!("\x1b[31m {line:4}  {code}\x1b[0m");
+            } else {
+                println!(" {line:4}  {code}");
+            }
+        }
+    }
+
+    Ok(())
+}
+
 #[derive(clap::Parser)]
 #[command(author, version, about, long_about = None)]
 struct Cli {
@@ -119,5 +183,12 @@ struct Cli {
 
 fn main() -> Result<()> {
     let cli = Cli::parse();
-    run_aci_tests(&Path::new(&cli.test_dir))
+
+    cfg_if::cfg_if! {
+    if #[cfg(feature = "coverage")] {
+        run_aci_tests_coverage(&Path::new(&cli.test_dir))
+    } else {
+        run_aci_tests(&Path::new(&cli.test_dir))
+    }
+    }
 }

tests/coverage/mod.rs

@@ -0,0 +1,81 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+use std::collections::BTreeSet;
+
+use regorus::*;
+
+use anyhow::Result;
+use test_generator::test_resources;
+
+#[derive(serde::Deserialize)]
+struct TestCase {
+    data: Option<Value>,
+    input: Option<Value>,
+    modules: Vec<String>,
+    note: String,
+    query: String,
+    uncovered: Vec<BTreeSet<u32>>,
+    skip: Option<bool>,
+}
+
+#[derive(serde::Deserialize)]
+struct YamlTest {
+    cases: Vec<TestCase>,
+}
+
+fn yaml_test_impl(file: &str) -> Result<()> {
+    let yaml_str = std::fs::read_to_string(file)?;
+    let test: YamlTest = serde_yaml::from_str(&yaml_str)?;
+
+    println!("running {file}");
+
+    for case in test.cases.into_iter() {
+        print!("case {} ", case.note);
+        if case.skip == Some(true) {
+            println!("skipped");
+            continue;
+        }
+
+        let mut engine = Engine::new();
+
+        for (idx, rego) in case.modules.iter().enumerate() {
+            engine.add_policy(format!("rego_{idx}"), rego.clone())?;
+        }
+
+        if let Some(data) = case.data {
+            engine.add_data(data)?;
+        }
+
+        if let Some(input) = case.input {
+            engine.set_input(input);
+        }
+
+        let _ = engine.eval_query(case.query.clone(), false)?;
+        let report = engine.get_coverage_report()?;
+
+        for (idx, uncovered) in case.uncovered.into_iter().enumerate() {
+            assert_eq!(uncovered, report.files[idx].uncovered);
+        }
+
+        println!("passed");
+    }
+
+    Ok(())
+}
+
+fn yaml_test(file: &str) -> Result<()> {
+    match yaml_test_impl(file) {
+        Ok(_) => Ok(()),
+        Err(e) => {
+            // If Err is returned, it doesn't always get printed by cargo test.
+            // Therefore, panic with the error.
+            panic!("{}", e);
+        }
+    }
+}
+
+#[test_resources("tests/coverage/*.yaml")]
+fn run(path: &str) {
+    yaml_test(path).unwrap()
+}

tests/coverage/tests.yaml

@@ -0,0 +1,20 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+cases:
+  - note: basic
+    modules:
+      - |
+        package test
+
+        x = 1
+        
+        y = k {
+          input.x == 5
+          k = input.k
+        }
+    query: data.test
+    uncovered: [
+      [ 5, 7 ]
+    ]
+        

tests/mod.rs

@@ -1,6 +1,9 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
+#[cfg(feature = "coverage")]
+mod coverage;
+
 mod engine;
 mod lexer;
 mod parser;