[WIP] Simple Policy Compiler

policy/compiler.go

@@ -0,0 +1,170 @@
+package policy
+
+import (
+	"github.com/google/cel-go/cel"
+	"github.com/google/cel-go/common"
+	"github.com/google/cel-go/common/ast"
+	"github.com/google/cel-go/common/operators"
+	"github.com/google/cel-go/common/types"
+)
+
+type compiler struct {
+	env  *cel.Env
+	info *ast.SourceInfo
+	src  *Source
+}
+
+type compiledRule struct {
+	variables []*compiledVariable
+	matches   []*compiledMatch
+}
+
+type compiledVariable struct {
+	name string
+	expr *cel.Ast
+}
+
+type compiledMatch struct {
+	cond       *cel.Ast
+	output     *cel.Ast
+	nestedRule *compiledRule
+}
+
+func compile(env *cel.Env, p *policy) (*cel.Ast, *cel.Issues) {
+	c := &compiler{
+		env:  env,
+		info: p.info,
+		src:  p.source,
+	}
+	errs := common.NewErrors(c.src)
+	iss := cel.NewIssuesWithSourceInfo(errs, c.info)
+	rule, ruleIss := c.compileRule(p.rule, c.env, iss)
+	iss = iss.Append(ruleIss)
+	if iss.Err() != nil {
+		return nil, iss
+	}
+	ruleRoot, _ := env.Compile("true")
+	opt := cel.NewStaticOptimizer(&ruleComposer{rule: rule})
+	ruleExprAST, iss := opt.Optimize(env, ruleRoot)
+	return ruleExprAST, iss.Append(iss)
+}
+
+func (c *compiler) compileRule(r *rule, ruleEnv *cel.Env, iss *cel.Issues) (*compiledRule, *cel.Issues) {
+	var err error
+	compiledVars := make([]*compiledVariable, len(r.variables))
+	for i, v := range r.variables {
+		exprSrc := c.relSource(v.expression)
+		varAST, exprIss := ruleEnv.CompileSource(exprSrc)
+		if exprIss.Err() == nil {
+			ruleEnv, err = ruleEnv.Extend(cel.Variable(v.name.value, varAST.OutputType()))
+			if err != nil {
+				iss.ReportErrorAtID(v.expression.id, "invalid variable declaration")
+			}
+			compiledVars[i] = &compiledVariable{
+				name: v.name.value,
+				expr: varAST,
+			}
+		}
+		iss = iss.Append(exprIss)
+	}
+	compiledMatches := []*compiledMatch{}
+	for _, m := range r.matches {
+		condSrc := c.relSource(m.condition)
+		condAST, condIss := ruleEnv.CompileSource(condSrc)
+		iss = iss.Append(condIss)
+		if m.output != nil && m.rule != nil {
+			iss.ReportErrorAtID(m.condition.id, "either output or rule may be set but not both")
+			continue
+		}
+		if m.output != nil {
+			outSrc := c.relSource(*m.output)
+			outAST, outIss := ruleEnv.CompileSource(outSrc)
+			iss = iss.Append(outIss)
+			compiledMatches = append(compiledMatches, &compiledMatch{
+				cond:   condAST,
+				output: outAST,
+			})
+			continue
+		}
+		if m.rule != nil {
+			nestedRule, ruleIss := c.compileRule(m.rule, ruleEnv, iss)
+			iss = iss.Append(ruleIss)
+			compiledMatches = append(compiledMatches, &compiledMatch{
+				cond:       condAST,
+				nestedRule: nestedRule,
+			})
+		}
+	}
+	return &compiledRule{
+		variables: compiledVars,
+		matches:   compiledMatches,
+	}, iss
+}
+
+func (c *compiler) relSource(pstr policyString) *RelativeSource {
+	line := 0
+	col := 1
+	if offset, found := c.info.GetOffsetRange(pstr.id); found {
+		if loc, found := c.src.OffsetLocation(offset.Start); found {
+			line = loc.Line()
+			col = loc.Column()
+		}
+	}
+	return c.src.Relative(pstr.value, line, col)
+}
+
+type ruleComposer struct {
+	rule *compiledRule
+}
+
+func (opt *ruleComposer) Optimize(ctx *cel.OptimizerContext, a *ast.AST) *ast.AST {
+	ruleExpr := optimizeRule(ctx, opt.rule)
+	ctx.UpdateExpr(a.Expr(), ruleExpr)
+	return ctx.NewAST(ruleExpr)
+}
+
+func optimizeRule(ctx *cel.OptimizerContext, r *compiledRule) ast.Expr {
+	matchExpr := ctx.NewCall("optional.none")
+	matches := r.matches
+	for i := len(matches) - 1; i >= 0; i-- {
+		m := matches[i]
+		cond := ctx.CopyASTAndMetadata(m.cond.NativeRep())
+		triviallyTrue := cond.Kind() == ast.LiteralKind && cond.AsLiteral() == types.True
+		if m.output != nil {
+			out := ctx.CopyASTAndMetadata(m.output.NativeRep())
+			if triviallyTrue {
+				matchExpr = out
+				continue
+			}
+			matchExpr = ctx.NewCall(
+				operators.Conditional,
+				cond,
+				ctx.NewCall("optional.of", out),
+				matchExpr)
+			continue
+		}
+		nestedRule := optimizeRule(ctx, m.nestedRule)
+		if triviallyTrue {
+			matchExpr = nestedRule
+			continue
+		}
+		matchExpr = ctx.NewCall(
+			operators.Conditional,
+			cond,
+			nestedRule,
+			matchExpr)
+	}
+
+	vars := r.variables
+	for i := len(vars) - 1; i >= 0; i-- {
+		v := vars[i]
+		varAST := ctx.CopyASTAndMetadata(v.expr.NativeRep())
+		// Build up the bindings in reverse order, starting from root, all the way up to the outermost
+		// binding:
+		//    currExpr = cel.bind(outerVar, outerExpr, currExpr)
+		inlined, bindMacro := ctx.NewBindMacro(matchExpr.ID(), v.name, varAST, matchExpr)
+		ctx.SetMacroCall(inlined.ID(), bindMacro)
+		matchExpr = inlined
+	}
+	return matchExpr
+}

policy/compiler_test.go

@@ -0,0 +1,32 @@
+package policy
+
+import (
+	"testing"
+
+	"github.com/google/cel-go/cel"
+)
+
+func TestCompile(t *testing.T) {
+	srcFile := readPolicy(t, "testdata/required_labels.yaml")
+	p, iss := parse(srcFile)
+	if iss.Err() != nil {
+		t.Fatalf("parse() failed: %v", iss.Err())
+	}
+	if p.name.value != "required_labels" {
+		t.Errorf("policy name is %v, wanted 'required_labels'", p.name)
+	}
+	env, err := cel.NewEnv(
+		cel.OptionalTypes(),
+		cel.EnableMacroCallTracking(),
+		cel.ExtendedValidations(),
+		cel.Variable("rule.labels", cel.MapType(cel.StringType, cel.StringType)),
+		cel.Variable("resource.labels", cel.MapType(cel.StringType, cel.StringType)),
+	)
+	if err != nil {
+		t.Fatalf("cel.NewEnv() failed: %v", err)
+	}
+	_, iss = compile(env, p)
+	if iss.Err() != nil {
+		t.Errorf("compile() failed: %v", iss.Err())
+	}
+}

policy/go.mod

@@ -0,0 +1,20 @@
+module github.com/google/cel-go/policy
+
+go 1.20
+
+require (
+	github.com/google/cel-go v0.20.1
+	gopkg.in/yaml.v3 v3.0.1
+)
+
+require (
+	github.com/antlr4-go/antlr/v4 v4.13.0 // indirect
+	github.com/stoewer/go-strcase v1.2.0 // indirect
+	golang.org/x/exp v0.0.0-20230515195305-f3d0a9c9a5cc // indirect
+	golang.org/x/text v0.9.0 // indirect
+	google.golang.org/genproto/googleapis/api v0.0.0-20230803162519-f966b187b2e5 // indirect
+	google.golang.org/genproto/googleapis/rpc v0.0.0-20230803162519-f966b187b2e5 // indirect
+	google.golang.org/protobuf v1.33.0 // indirect
+)
+
+replace github.com/google/cel-go => ../.

policy/go.sum

@@ -0,0 +1,28 @@
+github.com/antlr4-go/antlr/v4 v4.13.0 h1:lxCg3LAv+EUK6t1i0y1V6/SLeUi0eKEKdhQAlS8TVTI=
+github.com/antlr4-go/antlr/v4 v4.13.0/go.mod h1:pfChB/xh/Unjila75QW7+VU4TSnWnnk9UTnmpPaOR2g=
+github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8=
+github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/stoewer/go-strcase v1.2.0 h1:Z2iHWqGXH00XYgqDmNgQbIBxf3wrNq0F3feEy0ainaU=
+github.com/stoewer/go-strcase v1.2.0/go.mod h1:IBiWB2sKIp3wVVQ3Y035++gc+knqhUQag1KpM8ahLw8=
+github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/testify v1.5.1 h1:nOGnQDM7FYENwehXlg/kFVnos3rEvtKTjRvOWSzb6H4=
+github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5cxcmMvtA=
+golang.org/x/exp v0.0.0-20230515195305-f3d0a9c9a5cc h1:mCRnTeVUjcrhlRmO0VK8a6k6Rrf6TF9htwo2pJVSjIU=
+golang.org/x/exp v0.0.0-20230515195305-f3d0a9c9a5cc/go.mod h1:V1LtkGg67GoY2N1AnLN78QLrzxkLyJw7RJb1gzOOz9w=
+golang.org/x/text v0.9.0 h1:2sjJmO8cDvYveuX97RDLsxlyUxLl+GHoLxBiRdHllBE=
+golang.org/x/text v0.9.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8=
+google.golang.org/genproto/googleapis/api v0.0.0-20230803162519-f966b187b2e5 h1:nIgk/EEq3/YlnmVVXVnm14rC2oxgs1o0ong4sD/rd44=
+google.golang.org/genproto/googleapis/api v0.0.0-20230803162519-f966b187b2e5/go.mod h1:5DZzOUPCLYL3mNkQ0ms0F3EuUNZ7py1Bqeq6sxzI7/Q=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20230803162519-f966b187b2e5 h1:eSaPbMR4T7WfH9FvABk36NBMacoTUKdWCvV0dx+KfOg=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20230803162519-f966b187b2e5/go.mod h1:zBEcrKX2ZOcEkHWxBPAIvYUWOKKMIhYcmNiUIu2ji3I=
+google.golang.org/protobuf v1.33.0 h1:uNO2rsAINq/JlFpSdYEKIZ0uKD/R9cpdv0T+yoGwGmI=
+google.golang.org/protobuf v1.33.0/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHhKbcUYpos=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v2 v2.2.2 h1:ZCJp+EgiOT7lHqUV2J862kp8Qj64Jo6az82+3Td9dZw=
+gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
+gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
+gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=