Add CEL variables support (#37)

* feat: add variables support for checks

* feat: use variables in M-113

* feat: use lazy map for variables

* chore: rename compiledVariable struct

* chore: fix typo

internal/builtins/embed.go

@@ -19,4 +19,4 @@ import (
 )
 
 //go:embed *
-var EmbbedChecksFS embed.FS
+var EmbedChecksFS embed.FS

internal/builtins/embed_test.go

@@ -20,8 +20,8 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-func TestEmbbedChecks(t *testing.T) {
-	entries, err := EmbbedChecksFS.ReadDir(".")
+func TestEmbedChecks(t *testing.T) {
+	entries, err := EmbedChecksFS.ReadDir(".")
 	assert.NoError(t, err)
 	assert.Greater(t, len(entries), 0)
 }

internal/builtins/pss/restricted/M-113_run_as_non_root.yml

@@ -42,20 +42,23 @@ match:
     - group: batch
       version: v1
       resource: jobs
-validations:
-  # Pod or Containers must set `securityContext.runAsNonRoot`
-  - expression: >
-      (has(podSpec.securityContext) && has(podSpec.securityContext.runAsNonRoot)) ||
-      allContainers.all(container,
-        has(container.securityContext) && has(container.securityContext.runAsNonRoot)
+variables:
+  - name: podRunAsNonRoot
+    expression: podSpec.?securityContext.?runAsNonRoot.orValue(false)
+
+  # containers that explicitly set runAsNonRoot=false
+  - name: explicitlyBadContainers
+    expression: >
+      allContainers.filter(c,
+          has(c.securityContext) && has(c.securityContext.runAsNonRoot) && c.securityContext.runAsNonRoot == false
       )
-    message: "Container could be running as root user"
 
-  # Neither Pod nor Containers should set `securityContext.runAsNonRoot` to false
-  - expression: >
-      (!has(podSpec.securityContext) || !has(podSpec.securityContext.runAsNonRoot) || podSpec.securityContext.runAsNonRoot != false)
-      &&
-      allContainers.all(container,
-        !has(container.securityContext) || !has(container.securityContext.runAsNonRoot) || container.securityContext.runAsNonRoot != false
+  # containers that didn't set runAsNonRoot and aren't caught by a pod-level runAsNonRoot=true
+  - name: implicitlyBadContainers
+    expression: >
+      allContainers.filter(c,
+          !variables.podRunAsNonRoot && (!has(c.securityContext) || !has(c.securityContext.runAsNonRoot)) 
       )
-    message: "Container could be running as root user"
+
+validations:
+  - expression: variables.explicitlyBadContainers.size() == 0 && variables.implicitlyBadContainers.size() == 0

internal/builtins/pss/restricted/M-113_run_as_non_root_test.yml

@@ -14,7 +14,6 @@
 
 - name: "securityContext not specified"
   pass: false
-  message: "Container could be running as root user"
   input: |
     apiVersion: apps/v1
     kind: Deployment
@@ -34,9 +33,31 @@
         matchLabels:
           app: nginx
 
-- name: "Pod explicitly set runAsNonRoot to false and Container to true"
+- name: "Pod set runAsNonRoot to false"
   pass: false
-  message: "Container could be running as root user"
+  input: |
+    apiVersion: apps/v1
+    kind: Deployment
+    metadata:
+      name: nginx
+    spec:
+      template:
+        metadata:
+          name: nginx
+          labels:
+            app: nginx
+        spec:
+          securityContext:
+            runAsNonRoot: false
+          containers:
+            - name: nginx
+              image: nginx
+      selector:
+        matchLabels:
+          app: nginx
+
+- name: "Pod set runAsNonRoot to false and container to true"
+  pass: true
   input: |
     apiVersion: apps/v1
     kind: Deployment
@@ -60,9 +81,8 @@
         matchLabels:
           app: nginx
 
-- name: "Container explicitly set runAsNonRoot to false and Pod to true"
+- name: "container set runAsNonRoot to false and Pod to true"
   pass: false
-  message: "Container could be running as root user"
   input: |
     apiVersion: apps/v1
     kind: Deployment
@@ -86,7 +106,7 @@
         matchLabels:
           app: nginx
 
-- name: "Container set runAsNonRoot to true and Pod not specified"
+- name: "container set runAsNonRoot to true and Pod not specified"
   pass: true
   input: |
     apiVersion: apps/v1
@@ -109,7 +129,7 @@
         matchLabels:
           app: nginx
 
-- name: "Pod set runAsNonRoot to true and Container not specified"
+- name: "Pod set runAsNonRoot to true and container not specified"
   pass: true
   input: |
     apiVersion: apps/v1

pkg/loader/builtin.go

@@ -26,8 +26,8 @@ import (
 var Builtins []types.Check
 
 func init() {
-	c, _, walkFn := walkDir(builtins.EmbbedChecksFS.ReadFile, true)
-	err := fs.WalkDir(builtins.EmbbedChecksFS, ".", walkFn)
+	c, _, walkFn := walkDir(builtins.EmbedChecksFS.ReadFile, true)
+	err := fs.WalkDir(builtins.EmbedChecksFS, ".", walkFn)
 	if err != nil {
 		log.Fatal(err)
 	}

pkg/types/check.go

@@ -23,6 +23,7 @@ type Check struct {
 	ID          string            `json:"id"`
 	Match       Match             `json:"match"`
 	Validations []Validation      `json:"validations"`
+	Variables   []Variable        `json:"variables"`
 	Params      map[string]any    `json:"params"`
 	Severity    Severity          `json:"severity"`
 	Message     string            `json:"message"`
@@ -51,6 +52,11 @@ type Validation struct {
 	Message    string `json:"message,omitempty"`
 }
 
+type Variable struct {
+	Name       string `json:"name"`
+	Expression string `json:"expression"`
+}
+
 type Test struct {
 	Name        string        `json:"name"`
 	Input       string        `json:"input"`

pkg/validator/activation.go

@@ -26,6 +26,7 @@ const (
 	AllContainersVarName = "allContainers"
 	APIVersionsVarName   = "apiVersions"
 	KubeVersionVarName   = "kubeVersion"
+	VariableVarName      = "variables"
 )
 
 // activation implements the interpreter.Activation
@@ -37,6 +38,7 @@ type activation struct {
 	params        any
 	apiVersions   []string
 	kubeVersion   any
+	variables     any
 }
 
 func (a *activation) ResolveName(name string) (any, bool) {
@@ -55,6 +57,8 @@ func (a *activation) ResolveName(name string) (any, bool) {
 		return a.apiVersions, true
 	case KubeVersionVarName:
 		return a.kubeVersion, true
+	case VariableVarName:
+		return a.variables, true
 	default:
 		return nil, false
 	}

pkg/validator/compiler.go

@@ -43,53 +43,102 @@ var baseEnvOptions = []cel.EnvOption{
 	k8scellib.Quantity(),
 
 	cel.Variable(ObjectVarName, cel.DynType),
-	cel.Variable(ParamsVarName, cel.DynType),
 	cel.Variable(APIVersionsVarName, cel.ListType(cel.StringType)),
 	cel.Variable(KubeVersionVarName, cel.DynType),
 }
 
+var programOptions = []cel.ProgramOption{
+	cel.EvalOptions(cel.OptOptimize),
+	cel.CostLimit(1000000),
+	cel.InterruptCheckFrequency(100),
+}
+
 var podSpecEnvOptions = []cel.EnvOption{
 	cel.Variable(PodMetaVarName, cel.DynType),
 	cel.Variable(PodSpecVarName, cel.DynType),
 	cel.Variable(AllContainersVarName, cel.ListType(cel.DynType)),
 }
 
-// Compile compiles the expressions of the given check and returns a Validator
+// Compile compiles variables and expressions of the given check and returns a Validator
 func Compile(check types.Check, apiResources []*metav1.APIResourceList, kubeVersion *version.Info) (Validator, error) {
 	if len(check.Validations) == 0 {
 		return nil, errors.New("invalid check: a check must have at least 1 validation")
 	}
-	hasPodSpec := MatchesPodSpec(check.Match.Resources)
-	env, err := newEnv(hasPodSpec)
+	env, err := newEnv(check)
 	if err != nil {
 		return nil, fmt.Errorf("environment construction error %s", err.Error())
 	}
-	prgs := make([]cel.Program, 0, len(check.Validations))
-	for i, v := range check.Validations {
-		ast, issues := env.Compile(v.Expression)
-		if issues != nil && issues.Err() != nil {
-			return nil, fmt.Errorf("validation[%d].expression: type-check error: %s", i, issues.Err())
-		}
-		if ast.OutputType() != cel.BoolType {
-			return nil, fmt.Errorf("validation[%d].expression: cel expression must evaluate to a bool", i)
-		}
-		prg, err := env.Program(ast, cel.EvalOptions(cel.OptOptimize))
-		if err != nil {
-			return nil, fmt.Errorf("validation[%d].expression: program construction error: %s", i, err)
-		}
-		prgs = append(prgs, prg)
-	}
+
+	variables, err := compileVariables(env, check.Variables)
+
+	prgs, err := compileValidations(env, check.Validations)
+
 	apiVersions := make([]string, 0, len(apiResources))
 	for _, resource := range apiResources {
 		apiVersions = append(apiVersions, resource.GroupVersion)
 	}
-	return &CELValidator{check: check, programs: prgs, hasPodSpec: hasPodSpec, apiVersions: apiVersions, kubeVersion: kubeVersion}, nil
+	return &CELValidator{check: check, programs: prgs, apiVersions: apiVersions, kubeVersion: kubeVersion, variables: variables}, nil
 }
 
-func newEnv(podSpec bool) (*cel.Env, error) {
+func newEnv(check types.Check) (*cel.Env, error) {
 	opts := baseEnvOptions
-	if podSpec {
+	if MatchesPodSpec(check.Match.Resources) {
 		opts = append(opts, podSpecEnvOptions...)
 	}
+	if len(check.Variables) > 0 {
+		opts = append(opts, cel.Variable(VariableVarName, cel.MapType(cel.StringType, cel.DynType)))
+	}
+	if len(check.Params) > 0 {
+		opts = append(opts, cel.Variable(ParamsVarName, cel.DynType))
+	}
 	return cel.NewEnv(opts...)
 }
+
+func compileVariables(env *cel.Env, vars []types.Variable) ([]compiledVariable, error) {
+	variables := make([]compiledVariable, 0, len(vars))
+	for _, v := range vars {
+		prg, err := compileExpression(env, v.Expression, cel.AnyType)
+		if err != nil {
+			return nil, fmt.Errorf("variables[%q].expression: %s", v.Name, err)
+		}
+		variables = append(variables, compiledVariable{name: v.Name, program: prg})
+	}
+	return variables, nil
+}
+
+func compileValidations(env *cel.Env, vals []types.Validation) ([]cel.Program, error) {
+	prgs := make([]cel.Program, 0, len(vals))
+	for i, v := range vals {
+		prg, err := compileExpression(env, v.Expression, cel.BoolType)
+		if err != nil {
+			return nil, fmt.Errorf("validations[%d].expression: %s", i, err)
+		}
+		prgs = append(prgs, prg)
+	}
+	return prgs, nil
+}
+
+func compileExpression(env *cel.Env, exp string, allowedTypes ...*cel.Type) (cel.Program, error) {
+	ast, issues := env.Compile(exp)
+	if issues != nil && issues.Err() != nil {
+		return nil, fmt.Errorf("type-check error: %s", issues.Err())
+	}
+	found := false
+	for _, t := range allowedTypes {
+		if ast.OutputType() == t || cel.AnyType == t {
+			found = true
+			break
+		}
+	}
+	if !found {
+		if len(allowedTypes) == 1 {
+			return nil, fmt.Errorf("must evaluate to %v", allowedTypes[0].String())
+		}
+		return nil, fmt.Errorf("must evaluate to one of %v", allowedTypes)
+	}
+	prg, err := env.Program(ast, programOptions...)
+	if err != nil {
+		return nil, fmt.Errorf("program construction error: %s", err)
+	}
+	return prg, nil
+}

pkg/validator/validator.go

@@ -19,20 +19,26 @@ import (
 
 	"github.com/google/cel-go/cel"
 	"github.com/google/cel-go/common/types"
+	"github.com/google/cel-go/common/types/ref"
+	marvin "github.com/undistro/marvin/pkg/types"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/version"
-
-	marvin "github.com/undistro/marvin/pkg/types"
+	"k8s.io/apiserver/pkg/cel/lazy"
 )
 
 // CELValidator is a Validator that performs CEL expressions
 type CELValidator struct {
 	check       marvin.Check
 	programs    []cel.Program
-	hasPodSpec  bool
 	apiVersions []string
 	kubeVersion *version.Info
+	variables   []compiledVariable
+}
+
+type compiledVariable struct {
+	name    string
+	program cel.Program
 }
 
 func (r *CELValidator) SetAPIVersions(apiVersions []string) {
@@ -47,10 +53,15 @@ func (r *CELValidator) Validate(obj unstructured.Unstructured, params any) (bool
 	if params == nil {
 		params = r.check.Params
 	}
-	input := &activation{object: obj.UnstructuredContent(), apiVersions: r.apiVersions, params: params}
+	input := &activation{object: obj.UnstructuredContent(), apiVersions: r.apiVersions, params: params, variables: make(map[string]any)}
 	if err := r.setPodSpecParams(obj, input); err != nil {
 		return false, "", err
 	}
+	lazyMap := lazy.NewMapValue(types.MapType)
+	for _, v := range r.variables {
+		lazyMap.Append(v.name, callback(v, input))
+	}
+	input.variables = lazyMap
 	for i, prg := range r.programs {
 		out, _, err := prg.Eval(input)
 		if err != nil {
@@ -63,8 +74,18 @@ func (r *CELValidator) Validate(obj unstructured.Unstructured, params any) (bool
 	return true, "", nil
 }
 
+func callback(v compiledVariable, activation any) lazy.GetFieldFunc {
+	return func(_ *lazy.MapValue) ref.Val {
+		val, _, err := v.program.Eval(activation)
+		if err != nil {
+			return types.NewErr("variable %q fails to evaluate: %v", v.name, err)
+		}
+		return val
+	}
+}
+
 func (r *CELValidator) setPodSpecParams(obj unstructured.Unstructured, input *activation) error {
-	if !r.hasPodSpec || !HasPodSpec(obj) {
+	if !HasPodSpec(obj) {
 		return nil
 	}
 	meta, spec, err := ExtractPodSpec(obj)