From 5bcdb8b52c864ca39243a1e31bb94a229f87394a Mon Sep 17 00:00:00 2001
From: Tristan Swadell <tswadell@google.com>
Date: Fri, 2 Aug 2024 11:12:32 -0700
Subject: [PATCH] Check for output type agreement during the compile phase
 (#992)

* Check for output type agreement during the compile phase
* Updated comment and nil safety of OutputType
---
 policy/compiler.go                 | 50 ++++++++++++++++++++++++++++--
 policy/helper_test.go              |  5 ++-
 policy/testdata/errors/policy.yaml |  4 +++
 3 files changed, 56 insertions(+), 3 deletions(-)

diff --git a/policy/compiler.go b/policy/compiler.go
index 1ee12e9a..de2100ec 100644
--- a/policy/compiler.go
+++ b/policy/compiler.go
@@ -53,6 +53,16 @@ func (r *CompiledRule) Matches() []*CompiledMatch {
 	return r.matches[:]
 }
 
+// OutputType returns the output type of the first match clause as all match clauses
+// are validated for agreement prior to construction fo the CompiledRule.
+func (r *CompiledRule) OutputType() *cel.Type {
+	// It's a compilation error if the output types of the matches don't agree
+	for _, m := range r.Matches() {
+		return m.OutputType()
+	}
+	return cel.DynType
+}
+
 // CompiledVariable represents the variable name, expression, and associated type-check declaration.
 type CompiledVariable struct {
 	id      int64
@@ -105,6 +115,17 @@ func (m *CompiledMatch) NestedRule() *CompiledRule {
 	return m.nestedRule
 }
 
+// OutputType returns the cel.Type associated with output expression.
+func (m *CompiledMatch) OutputType() *cel.Type {
+	if m.output != nil {
+		return m.output.Expr().OutputType()
+	}
+	if m.nestedRule != nil {
+		return m.nestedRule.OutputType()
+	}
+	return cel.DynType
+}
+
 // OutputValue represents the output expression associated with a match block.
 type OutputValue struct {
 	id   int64
@@ -263,11 +284,36 @@ func (c *compiler) compileRule(r *Rule, ruleEnv *cel.Env, iss *cel.Issues) (*Com
 			}
 		}
 	}
-	return &CompiledRule{
+
+	rule := &CompiledRule{
 		id:        r.id,
 		variables: compiledVars,
 		matches:   compiledMatches,
-	}, iss
+	}
+	// Validate type agreement between the different match outputs
+	c.checkMatchOutputTypesAgree(rule, iss)
+	return rule, iss
+}
+
+func (c *compiler) checkMatchOutputTypesAgree(rule *CompiledRule, iss *cel.Issues) {
+	var outputType *cel.Type
+	for _, m := range rule.Matches() {
+		if outputType == nil {
+			outputType = m.OutputType()
+			if outputType.TypeName() == "error" {
+				outputType = nil
+				continue
+			}
+		}
+		matchOutputType := m.OutputType()
+		if matchOutputType.TypeName() == "error" {
+			continue
+		}
+		if !outputType.IsAssignableType(matchOutputType) {
+			iss.ReportErrorAtID(m.Output().ID(), "incompatible output types: %s not assignable to %s", outputType, matchOutputType)
+			return
+		}
+	}
 }
 
 func (c *compiler) relSource(pstr ValueString) *RelativeSource {
diff --git a/policy/helper_test.go b/policy/helper_test.go
index afd2ae14..f3c84b40 100644
--- a/policy/helper_test.go
+++ b/policy/helper_test.go
@@ -159,7 +159,10 @@ ERROR: testdata/errors/policy.yaml:31:75: Syntax error: extraneous input ']' exp
  | ..........................................................................^
 ERROR: testdata/errors/policy.yaml:34:67: undeclared reference to 'format' (in container '')
  |         "invalid values provided on one or more labels: %s".format([variables.invalid])
- | ..................................................................^`,
+ | ..................................................................^
+ERROR: testdata/errors/policy.yaml:38:16: incompatible output types: bool not assignable to string
+ |       output: "'false'"
+ | ...............^`,
 		},
 		{
 			name: "limits",
diff --git a/policy/testdata/errors/policy.yaml b/policy/testdata/errors/policy.yaml
index 338ba399..e43c9453 100644
--- a/policy/testdata/errors/policy.yaml
+++ b/policy/testdata/errors/policy.yaml
@@ -32,3 +32,7 @@ rule:
     - condition: variables.invalid.size() > 0
       output: |
         "invalid values provided on one or more labels: %s".format([variables.invalid])
+    - condition: "1 > 0"
+      output: "true"
+    - condition: "1 < 0"
+      output: "'false'"