feat(rego): extend engine with custom matchers (#2396)

Signed-off-by: Sylwester Piskozub <sylwesterpiskozub@gmail.com>

Author: Piskoo

pkg/policies/engine/engine.go

@@ -24,23 +24,29 @@ import (
 type PolicyEngine interface {
 	// Verify verifies an input against a policy
 	Verify(ctx context.Context, policy *Policy, input []byte, args map[string]any) (*EvaluationResult, error)
+	// MatchesParameters evaluates the matches_parameters rule to determine if evaluation parameters match expected parameters
+	MatchesParameters(ctx context.Context, policy *Policy, evaluationParams, expectedParams map[string]string) (bool, error)
+	// MatchesEvaluation evaluates the matches_evaluation rule using a PolicyEvaluation result and evaluation parameters
+	MatchesEvaluation(ctx context.Context, policy *Policy, evaluation *EvaluationResult, evaluationParams map[string]string) (bool, error)
 }
 
 type EvaluationResult struct {
-	Violations []*PolicyViolation
-	Skipped    bool
-	SkipReason string
-	Ignore     bool
-	RawData    *RawData
+	Violations []*PolicyViolation `json:"violations"`
+	Skipped    bool               `json:"skipped"`
+	SkipReason string             `json:"skipReason"`
+	Ignore     bool               `json:"ignore"`
+	RawData    *RawData           `json:"rawData"`
 }
+
 type RawData struct {
-	Input  json.RawMessage
-	Output json.RawMessage
+	Input  json.RawMessage `json:"input"`
+	Output json.RawMessage `json:"output"`
 }
 
 // PolicyViolation represents a policy failure
 type PolicyViolation struct {
-	Subject, Violation string
+	Subject   string `json:"subject"`
+	Violation string `json:"violation"`
 }
 
 // Policy represents a loaded policy in any of the supported engines.

pkg/policies/engine/rego/rego.go

@@ -111,9 +111,13 @@ const (
 	// EnvironmentModePermissive allows all operations on the compiler
 	EnvironmentModePermissive EnvironmentMode = 1
 	inputArgs                                 = "args"
+	expectedArgs                              = "expected_args"
+	evalResult                                = "evaluation_result"
 	inputElements                             = "elements"
 	deprecatedRule                            = "violations"
 	mainRule                                  = "result"
+	matchesParametersRule                     = "matches_parameters"
+	matchesEvaluationRule                     = "matches_evaluation"
 )
 
 // builtinFuncNotAllowed is a list of builtin functions that are not allowed in the compiler
@@ -383,3 +387,85 @@ func getRuleName(packagePath ast.Ref, rule string) string {
 	}
 	return fmt.Sprintf("%v.%s\n", packagePath, rule)
 }
+
+// MatchesParameters evaluates the matches_parameters rule in a rego policy.
+// The function creates an input object with policy parameters and expected parameters.
+// Returns true if the policy's matches_parameters rule evaluates to true, false otherwise.
+func (r *Engine) MatchesParameters(ctx context.Context, policy *engine.Policy, evaluationParams, expectedParams map[string]string) (bool, error) {
+	policyString := string(policy.Source)
+	parsedModule, err := ast.ParseModule(policy.Name, policyString)
+	if err != nil {
+		return false, fmt.Errorf("failed to parse rego policy: %w", err)
+	}
+
+	// Create input with policy and expected parameters
+	inputMap := make(map[string]interface{})
+	inputMap[inputArgs] = evaluationParams
+	inputMap[expectedArgs] = expectedParams
+
+	// Evaluate matches_parameters rule
+	matchesParameters, err := r.evaluateMatchingRule(ctx, getRuleName(parsedModule.Package.Path, matchesParametersRule), parsedModule, inputMap)
+	if err != nil {
+		// Defaults to false
+		return false, err
+	}
+
+	return matchesParameters, nil
+}
+
+// MatchesEvaluation evaluates the matches_evaluation rule in a rego policy.
+// The function creates an input object with policy parameters and evaluation result.
+// Returns true if the policy's matches_evaluation rule evaluates to true, false otherwise.
+// If the rule is not found or evaluation fails, it defaults to false.
+func (r *Engine) MatchesEvaluation(ctx context.Context, policy *engine.Policy, ev *engine.EvaluationResult, evaluationParams map[string]string) (bool, error) {
+	policyString := string(policy.Source)
+	parsedModule, err := ast.ParseModule(policy.Name, policyString)
+	if err != nil {
+		return false, fmt.Errorf("failed to parse rego policy: %w", err)
+	}
+
+	// Create input with the policy evaluation data
+	inputMap := make(map[string]interface{})
+	inputMap[inputArgs] = evaluationParams
+	inputMap[evalResult] = ev
+
+	// Evaluate matches_parameters rule
+	matchesEvaluation, err := r.evaluateMatchingRule(ctx, getRuleName(parsedModule.Package.Path, matchesEvaluationRule), parsedModule, inputMap)
+	if err != nil {
+		// Defaults to false
+		return false, err
+	}
+
+	return matchesEvaluation, nil
+}
+
+// Evaluates a single rule and returns its boolean result
+func (r *Engine) evaluateMatchingRule(ctx context.Context, ruleName string, parsedModule *ast.Module, decodedInput interface{}) (bool, error) {
+	// Add input
+	regoInput := rego.Input(decodedInput)
+
+	// Add module
+	regoFunc := rego.ParsedModule(parsedModule)
+	options := []func(r *rego.Rego){regoInput, regoFunc, rego.Capabilities(r.Capabilities())}
+
+	if r.operatingMode == EnvironmentModeRestrictive {
+		options = append(options, rego.StrictBuiltinErrors(true))
+	}
+
+	res, err := queryRego(ctx, ruleName, options...)
+	if err != nil {
+		return false, err
+	}
+
+	// Parse the boolean result
+	for _, exp := range res {
+		for _, val := range exp.Expressions {
+			if boolResult, ok := val.Value.(bool); ok {
+				return boolResult, nil
+			}
+		}
+	}
+
+	// No valid boolean result found
+	return false, nil
+}

pkg/policies/engine/rego/rego_test.go

@@ -330,3 +330,130 @@ func TestRego_WithPermissiveMode(t *testing.T) {
 		assert.NotContains(t, err.Error(), "rego_type_error: undefined function rego.parse_module")
 	})
 }
+
+func TestRego_MatchesParameters(t *testing.T) {
+	regoContent, err := os.ReadFile("testfiles/matches_parameters.rego")
+	require.NoError(t, err)
+
+	r := NewEngine()
+	policy := &engine.Policy{
+		Name:   "matches-parameters-test",
+		Source: regoContent,
+	}
+
+	t.Run("high severity matches medium expectation", func(t *testing.T) {
+		matches, err := r.MatchesParameters(context.TODO(), policy,
+			map[string]string{"severity": "high"},
+			map[string]string{"severity": "medium"})
+		require.NoError(t, err)
+		assert.True(t, matches)
+	})
+
+	t.Run("low severity does not match high expectation", func(t *testing.T) {
+		matches, err := r.MatchesParameters(context.TODO(), policy,
+			map[string]string{"severity": "low"},
+			map[string]string{"severity": "high"})
+		require.NoError(t, err)
+		assert.False(t, matches)
+	})
+
+	t.Run("critical severity matches critical expectation", func(t *testing.T) {
+		matches, err := r.MatchesParameters(context.TODO(), policy,
+			map[string]string{"severity": "critical"},
+			map[string]string{"severity": "critical"})
+		require.NoError(t, err)
+		assert.True(t, matches)
+	})
+
+	t.Run("unknown severity parameter", func(t *testing.T) {
+		matches, err := r.MatchesParameters(context.TODO(), policy,
+			map[string]string{"severity": "unknown"},
+			map[string]string{"severity": "medium"})
+		require.NoError(t, err)
+		assert.False(t, matches)
+	})
+
+	t.Run("empty parameters", func(t *testing.T) {
+		matches, err := r.MatchesParameters(context.TODO(), policy,
+			map[string]string{},
+			map[string]string{})
+		require.NoError(t, err)
+		assert.False(t, matches)
+	})
+}
+
+func TestRego_MatchesEvaluation(t *testing.T) {
+	regoContent, err := os.ReadFile("testfiles/matches_evaluation.rego")
+	require.NoError(t, err)
+
+	r := NewEngine()
+	policy := &engine.Policy{
+		Name:   "matches-evaluation-test",
+		Source: regoContent,
+	}
+
+	t.Run("evaluation with violations and high severity matches", func(t *testing.T) {
+		evaluation := &engine.EvaluationResult{
+			Violations: []*engine.PolicyViolation{
+				{Subject: "test", Violation: "test violation"},
+			},
+			Skipped:    false,
+			SkipReason: "",
+			Ignore:     false,
+		}
+		evaluationParams := map[string]string{"severity": "high"}
+		matches, err := r.MatchesEvaluation(context.TODO(), policy, evaluation, evaluationParams)
+		require.NoError(t, err)
+		assert.True(t, matches)
+	})
+
+	t.Run("evaluation without violations does not match", func(t *testing.T) {
+		evaluation := &engine.EvaluationResult{
+			Violations: []*engine.PolicyViolation{},
+			Skipped:    false,
+			SkipReason: "",
+			Ignore:     false,
+		}
+		evaluationParams := map[string]string{"severity": "high"}
+		matches, err := r.MatchesEvaluation(context.TODO(), policy, evaluation, evaluationParams)
+		require.NoError(t, err)
+		assert.False(t, matches)
+	})
+
+	t.Run("evaluation with violations but wrong severity does not match", func(t *testing.T) {
+		evaluation := &engine.EvaluationResult{
+			Violations: []*engine.PolicyViolation{
+				{Subject: "test", Violation: "test violation"},
+			},
+			Skipped:    false,
+			SkipReason: "",
+			Ignore:     false,
+		}
+		evaluationParams := map[string]string{"severity": "low"}
+		matches, err := r.MatchesEvaluation(context.TODO(), policy, evaluation, evaluationParams)
+		require.NoError(t, err)
+		assert.False(t, matches)
+	})
+
+	t.Run("nil evaluation does not match", func(t *testing.T) {
+		evaluationParams := map[string]string{"severity": "high"}
+		matches, err := r.MatchesEvaluation(context.TODO(), policy, nil, evaluationParams)
+		require.NoError(t, err)
+		assert.False(t, matches)
+	})
+
+	t.Run("empty evaluation params", func(t *testing.T) {
+		evaluation := &engine.EvaluationResult{
+			Violations: []*engine.PolicyViolation{
+				{Subject: "test", Violation: "test violation"},
+			},
+			Skipped:    false,
+			SkipReason: "",
+			Ignore:     false,
+		}
+		evaluationParams := map[string]string{}
+		matches, err := r.MatchesEvaluation(context.TODO(), policy, evaluation, evaluationParams)
+		require.NoError(t, err)
+		assert.False(t, matches)
+	})
+}

pkg/policies/engine/rego/testfiles/matches_evaluation.rego

@@ -0,0 +1,12 @@
+package test_matches_evaluation
+
+matches_evaluation := result {
+    # Check if the evaluation contains violations 
+    count(input.evaluation_result.violations) > 0
+    
+    # Check if we have the expected parameter
+    input.args.severity == "high"
+    
+    # If both conditions are met, return true
+    result := true
+}

pkg/policies/engine/rego/testfiles/matches_parameters.rego

@@ -0,0 +1,20 @@
+package test_matches_parameters
+
+severity_levels := ["low", "medium", "high", "critical"]
+
+severity_index(level) := index {
+    some i
+    severity_levels[i] == level
+    index := i
+}
+
+matches_parameters := result {
+    eval_severity := input.args.severity
+    expected_severity := input.expected_args.severity
+    
+    eval_idx := severity_index(eval_severity)
+    expected_idx := severity_index(expected_severity)
+    
+    # Evaluation severity must be >= expected severity
+    result := eval_idx >= expected_idx
+}