From 331f19527ea037d7b270321050ba4e42a61c7a8d Mon Sep 17 00:00:00 2001
From: TristonianJones <tswadell@google.com>
Date: Wed, 20 Nov 2024 16:41:10 -0800
Subject: [PATCH 1/2] Add two-variable comprehension support to cel-policy

---
 policy/config.go                            |  3 +++
 policy/go.mod                               |  2 +-
 policy/helper_test.go                       | 12 ++++++------
 policy/parser.go                            |  8 ++++----
 policy/testdata/required_labels/config.yaml |  1 +
 policy/testdata/required_labels/policy.yaml |  4 ++--
 6 files changed, 17 insertions(+), 13 deletions(-)

diff --git a/policy/config.go b/policy/config.go
index d58c3b8c..f6cd3fdc 100644
--- a/policy/config.go
+++ b/policy/config.go
@@ -287,4 +287,7 @@ var extFactories = map[string]ExtensionFactory{
 	"strings": func(version uint32) cel.EnvOption {
 		return ext.Strings(ext.StringsVersion(version))
 	},
+	"two-var-comprehensions": func(version uint32) cel.EnvOption {
+		return ext.TwoVarComprehensions()
+	},
 }
diff --git a/policy/go.mod b/policy/go.mod
index 4d2a16f2..abc04514 100644
--- a/policy/go.mod
+++ b/policy/go.mod
@@ -3,7 +3,7 @@ module github.com/google/cel-go/policy
 go 1.22
 
 require (
-	github.com/google/cel-go v0.21.0
+	github.com/google/cel-go v0.22.0
 	google.golang.org/protobuf v1.34.2
 	gopkg.in/yaml.v3 v3.0.1
 )
diff --git a/policy/helper_test.go b/policy/helper_test.go
index d681fb9a..c621d7a1 100644
--- a/policy/helper_test.go
+++ b/policy/helper_test.go
@@ -118,12 +118,12 @@ var (
 	cel.@block([
 	  spec.labels,
 	  @index0.filter(l, !(l in resource.labels)),
-	  resource.labels.filter(l, l in @index0 && @index0[l] != resource.labels[l])],
-	(@index1.size() > 0)
-	  ? optional.of("missing one or more required labels: %s".format([@index1]))
-	  : ((@index2.size() > 0)
-	    ? optional.of("invalid values provided on one or more labels: %s".format([@index2]))
-		: optional.none()))`,
+	    resource.labels.transformList(l, value, l in @index0 && value != @index0[l], l)],
+      (@index1.size() > 0)
+	   ? optional.of("missing one or more required labels: %s".format([@index1]))
+	   : ((@index2.size() > 0)
+	     ? optional.of("invalid values provided on one or more labels: %s".format([@index2]))
+		 : optional.none()))`,
 		},
 		{
 			name: "restricted_destinations",
diff --git a/policy/parser.go b/policy/parser.go
index 306e1418..7b34b246 100644
--- a/policy/parser.go
+++ b/policy/parser.go
@@ -472,17 +472,17 @@ func (parser *Parser) Parse(src *Source) (*Policy, *cel.Issues) {
 	errs := common.NewErrors(src)
 	iss := cel.NewIssuesWithSourceInfo(errs, info)
 	p := newParserImpl(parser.TagVisitor, info, src, iss)
-	policy := p.parseYaml(src)
+	policy := p.parseYAML(src)
 	if iss.Err() != nil {
 		return nil, iss
 	}
 	return policy, nil
 }
 
-func (p *parserImpl) parseYaml(src *Source) *Policy {
+func (p *parserImpl) parseYAML(src *Source) *Policy {
 	// Parse yaml representation from the source to an object model.
 	var docNode yaml.Node
-	err := sourceToYaml(src, &docNode)
+	err := sourceToYAML(src, &docNode)
 	if err != nil {
 		p.iss.ReportErrorAtID(0, err.Error())
 		return nil
@@ -491,7 +491,7 @@ func (p *parserImpl) parseYaml(src *Source) *Policy {
 	return p.ParsePolicy(p, docNode.Content[0])
 }
 
-func sourceToYaml(src *Source, docNode *yaml.Node) error {
+func sourceToYAML(src *Source, docNode *yaml.Node) error {
 	err := yaml.Unmarshal([]byte(src.Content()), docNode)
 	if err != nil {
 		return err
diff --git a/policy/testdata/required_labels/config.yaml b/policy/testdata/required_labels/config.yaml
index 14311d76..1fae24d4 100644
--- a/policy/testdata/required_labels/config.yaml
+++ b/policy/testdata/required_labels/config.yaml
@@ -17,6 +17,7 @@ extensions:
   - name: "bindings"
   - name: "strings"
     version: 2
+  - name: "two-var-comprehensions"
 variables:
   - name: "spec"
     type:
diff --git a/policy/testdata/required_labels/policy.yaml b/policy/testdata/required_labels/policy.yaml
index 219c6cc9..8508564d 100644
--- a/policy/testdata/required_labels/policy.yaml
+++ b/policy/testdata/required_labels/policy.yaml
@@ -21,8 +21,8 @@ rule:
       expression: variables.want.filter(l, !(l in resource.labels))
     - name: invalid
       expression: >
-        resource.labels.filter(l,
-          l in variables.want && variables.want[l] != resource.labels[l])
+        resource.labels.transformList(l, value,
+          l in variables.want && value != variables.want[l], l)
   match:
     - condition: variables.missing.size() > 0
       output: |

From ceadd6d05eaaad475d839c860aa77f1cb39c91ed Mon Sep 17 00:00:00 2001
From: TristonianJones <tswadell@google.com>
Date: Wed, 20 Nov 2024 18:01:44 -0800
Subject: [PATCH 2/2] Fix bazel BUILD dep to include comprehensions and
 additional tests

---
 ext/BUILD.bazel | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/ext/BUILD.bazel b/ext/BUILD.bazel
index 1fece700..b764fa1f 100644
--- a/ext/BUILD.bazel
+++ b/ext/BUILD.bazel
@@ -8,6 +8,7 @@ go_library(
     name = "go_default_library",
     srcs = [
         "bindings.go",
+        "comprehensions.go",
         "encoders.go",
         "formatting.go",
         "guards.go",
@@ -45,7 +46,9 @@ go_test(
     name = "go_default_test",
     size = "small",
     srcs = [
-        "encoders_test.go",
+        "bindings_test.go",
+        "comprehensions_test.go",
+        "encoders_test.go",        
         "lists_test.go",
         "math_test.go",
         "native_test.go",