Make library prefix constant

constraint/pkg/client/client_addtemplate_bench_test.go

@@ -8,35 +8,43 @@ import (
 
 	"github.com/open-policy-agent/frameworks/constraint/pkg/client/clienttest"
 	"github.com/open-policy-agent/frameworks/constraint/pkg/core/templates"
+	"github.com/open-policy-agent/frameworks/constraint/pkg/handler/handlertest"
 )
 
 var modules = []struct {
-	name       string
-	makeModule func(i int) string
-}{{
-	name:       "Simple",
-	makeModule: makeModuleSimple,
-}, {
-	name:       "Complex",
-	makeModule: makeModuleComplex,
-}}
+	name   string
+	module string
+	libs   []string
+}{
+	{
+		name:   "Simple",
+		module: moduleSimple,
+		libs:   nil,
+	},
+	{
+		name:   "Complex",
+		module: moduleComplex,
+		libs:   nil,
+	},
+	{
+		name:   "Very Complex",
+		module: moduleVeryComplex,
+		libs:   []string{libVeryComplex},
+	},
+}
 
 func makeKind(i int) string {
 	return fmt.Sprintf("foo%d", i)
 }
 
-func makeModuleSimple(i int) string {
-	kind := makeKind(i)
-	return fmt.Sprintf(`package %s
+const moduleSimple = `package foo
+
 violation[{"msg": msg}] {
   input.review.object.foo == input.parameters.foo
   msg := sprintf("input.foo is %%v", [input.parameters.foo])
-}`, kind)
-}
+}`
 
-func makeModuleComplex(i int) string {
-	kind := makeKind(i)
-	return fmt.Sprintf(`package %s
+const moduleComplex = `package foo
 
 identical(obj, review) {
   obj.metadata.namespace == review.object.metadata.namespace
@@ -52,17 +60,189 @@ violation[{"msg": msg}] {
   other.spec.rules[_].host == host
   not identical(other, input.review)
   msg := sprintf("Ingress host conflicts with an existing Ingress <%%v>", [host])
-}`, kind)
+}`
+
+const libVeryComplex = `package lib.helpers
+
+missing(obj, field) = true {
+    not obj[field]
+}
+
+missing(obj, field) = true {
+  obj[field] == ""
+}
+
+canonify_cpu(orig) = new {
+  is_number(orig)
+  new := orig * 1000
+}
+
+canonify_cpu(orig) = new {
+  not is_number(orig)
+  endswith(orig, "m")
+  new := to_number(replace(orig, "m", ""))
+}
+
+canonify_cpu(orig) = new {
+  not is_number(orig)
+  not endswith(orig, "m")
+  re_match("^[0-9]+$", orig)
+  new := to_number(orig) * 1000
+}
+
+# 10 ** 21
+mem_multiple("E") = 1000000000000000000000 { true }
+
+# 10 ** 18
+mem_multiple("P") = 1000000000000000000 { true }
+
+# 10 ** 15
+mem_multiple("T") = 1000000000000000 { true }
+
+# 10 ** 12
+mem_multiple("G") = 1000000000000 { true }
+
+# 10 ** 9
+mem_multiple("M") = 1000000000 { true }
+
+# 10 ** 6
+mem_multiple("k") = 1000000 { true }
+
+# 10 ** 3
+mem_multiple("") = 1000 { true }
+
+# Kubernetes accepts millibyte precision when it probably shouldn't.
+# https://github.com/kubernetes/kubernetes/issues/28741
+# 10 ** 0
+mem_multiple("m") = 1 { true }
+
+# 1000 * 2 ** 10
+mem_multiple("Ki") = 1024000 { true }
+
+# 1000 * 2 ** 20
+mem_multiple("Mi") = 1048576000 { true }
+
+# 1000 * 2 ** 30
+mem_multiple("Gi") = 1073741824000 { true }
+
+# 1000 * 2 ** 40
+mem_multiple("Ti") = 1099511627776000 { true }
+
+# 1000 * 2 ** 50
+mem_multiple("Pi") = 1125899906842624000 { true }
+
+# 1000 * 2 ** 60
+mem_multiple("Ei") = 1152921504606846976000 { true }
+
+get_suffix(mem) = suffix {
+  not is_string(mem)
+  suffix := ""
+}
+
+get_suffix(mem) = suffix {
+  is_string(mem)
+  suffix := substring(mem, count(mem) - 1, -1)
+  mem_multiple(suffix)
 }
 
-func makeConstraintTemplate(i int, makeModule func(i int) string) *templates.ConstraintTemplate {
+get_suffix(mem) = suffix {
+  is_string(mem)
+  suffix := substring(mem, count(mem) - 2, -1)
+  mem_multiple(suffix)
+}
+
+get_suffix(mem) = suffix {
+  is_string(mem)
+  not substring(mem, count(mem) - 1, -1)
+  not substring(mem, count(mem) - 2, -1)
+  suffix := ""
+}
+
+canonify_mem(orig) = new {
+  is_number(orig)
+  new := orig * 1000
+}
+
+canonify_mem(orig) = new {
+  not is_number(orig)
+  suffix := get_suffix(orig)
+  raw := replace(orig, suffix, "")
+  new := to_number(raw) * mem_multiple(suffix)
+}`
+
+const moduleVeryComplex = `package k8scontainerlimits
+import data.lib.helpers
+
+
+violation[{"msg": msg}] {
+  container := input.review.object.spec.containers[_]
+  cpu_orig := container.resources.limits.cpu
+  not helpers.canonify_cpu(cpu_orig)
+  msg := sprintf("container <%v> cpu limit <%v> could not be parsed", [container.name, cpu_orig])
+}
+
+violation[{"msg": msg}] {
+  container := input.review.object.spec.containers[_]
+  mem_orig := container.resources.limits.memory
+  not helpers.canonify_mem(mem_orig)
+  msg := sprintf("container <%v> memory limit <%v> could not be parsed", [container.name, mem_orig])
+}
+
+violation[{"msg": msg}] {
+  container := input.review.object.spec.containers[_]
+  not container.resources
+  msg := sprintf("container <%v> has no resource limits", [container.name])
+}
+
+violation[{"msg": msg}] {
+  container := input.review.object.spec.containers[_]
+  not container.resources.limits
+  msg := sprintf("container <%v> has no resource limits", [container.name])
+}
+
+violation[{"msg": msg}] {
+  container := input.review.object.spec.containers[_]
+  helpers.missing(container.resources.limits, "cpu")
+  msg := sprintf("container <%v> has no cpu limit", [container.name])
+}
+
+violation[{"msg": msg}] {
+  container := input.review.object.spec.containers[_]
+  helpers.missing(container.resources.limits, "memory")
+  msg := sprintf("container <%v> has no memory limit", [container.name])
+}
+
+violation[{"msg": msg}] {
+  container := input.review.object.spec.containers[_]
+  cpu_orig := container.resources.limits.cpu
+  cpu := helpers.canonify_cpu(cpu_orig)
+  max_cpu_orig := input.parameters.cpu
+  max_cpu := helpers.canonify_cpu(max_cpu_orig)
+  cpu > max_cpu
+  msg := sprintf("container <%v> cpu limit <%v> is higher than the maximum allowed of <%v>", [container.name, cpu_orig, max_cpu_orig])
+}
+
+violation[{"msg": msg}] {
+  container := input.review.object.spec.containers[_]
+  mem_orig := container.resources.limits.memory
+  mem := helpers.canonify_mem(mem_orig)
+  max_mem_orig := input.parameters.memory
+  max_mem := helpers.canonify_mem(max_mem_orig)
+  mem > max_mem
+  msg := sprintf("container <%v> memory limit <%v> is higher than the maximum allowed of <%v>", [container.name, mem_orig, max_mem_orig])
+}
+
+`
+
+func makeConstraintTemplate(i int, module string, libs ...string) *templates.ConstraintTemplate {
 	kind := makeKind(i)
 	ct := &templates.ConstraintTemplate{}
 	ct.SetName(kind)
 	ct.Spec.CRD.Spec.Names.Kind = kind
 	ct.Spec.Targets = []templates.Target{{
-		Target: "test.target",
-		Rego:   makeModule(i),
+		Target: handlertest.TargetName,
+		Rego:   module,
+		Libs:   libs,
 	}}
 
 	return ct
@@ -76,7 +256,7 @@ func BenchmarkClient_AddTemplate(b *testing.B) {
 					ctx := context.Background()
 					cts := make([]*templates.ConstraintTemplate, n)
 					for i := range cts {
-						cts[i] = makeConstraintTemplate(i, tc.makeModule)
+						cts[i] = makeConstraintTemplate(i, tc.module, tc.libs...)
 					}
 
 					for i := 0; i < b.N; i++ {
@@ -108,7 +288,7 @@ func BenchmarkClient_AddTemplate_Parallel(b *testing.B) {
 				b.Run(fmt.Sprintf("%d Templates", n), func(b *testing.B) {
 					cts := make([]*templates.ConstraintTemplate, n)
 					for i := range cts {
-						cts[i] = makeConstraintTemplate(i, tc.makeModule)
+						cts[i] = makeConstraintTemplate(i, tc.module, tc.libs...)
 					}
 
 					for i := 0; i < b.N; i++ {

constraint/pkg/client/client_test.go

@@ -370,6 +370,13 @@ r = 5
 			wantHandled: nil,
 			wantError:   clienterrors.ErrInvalidConstraintTemplate,
 		},
+		{
+			name:        "Very Complex Template",
+			targets:     []handler.TargetHandler{&handlertest.Handler{}},
+			template:    cts.New(cts.OptTargets(cts.Target(handlertest.TargetName, moduleVeryComplex, libVeryComplex))),
+			wantHandled: map[string]bool{handlertest.TargetName: true},
+			wantError:   nil,
+		},
 	}
 
 	for _, tc := range tcs {

constraint/pkg/client/drivers/local/compilers.go

@@ -124,17 +124,14 @@ type TargetModule struct {
 // parseConstraintTemplate validates the rego in template target by parsing
 // rego modules.
 func parseConstraintTemplate(templ *templates.ConstraintTemplate, externs []string) (map[string]TargetModule, error) {
-	kind := templ.Spec.CRD.Spec.Names.Kind
-	pkgPrefix := templateLibPrefix(kind)
-
-	rr, err := regorewriter.New(regorewriter.NewPackagePrefixer(pkgPrefix), []string{libRoot}, externs)
+	rr, err := regorewriter.New(regorewriter.NewPackagePrefixer(templateLibPrefix), []string{libRoot}, externs)
 	if err != nil {
 		return nil, fmt.Errorf("creating rego rewriter: %w", err)
 	}
 
 	mods := make(map[string]TargetModule)
 	for _, target := range templ.Spec.Targets {
-		targetMods, err := parseConstraintTemplateTarget(rr, pkgPrefix, target)
+		targetMods, err := parseConstraintTemplateTarget(rr, target)
 		if err != nil {
 			return nil, err
 		}
@@ -148,7 +145,7 @@ func parseConstraintTemplate(templ *templates.ConstraintTemplate, externs []stri
 	return mods, nil
 }
 
-func parseConstraintTemplateTarget(rr *regorewriter.RegoRewriter, pkgPrefix string, targetSpec templates.Target) ([]string, error) {
+func parseConstraintTemplateTarget(rr *regorewriter.RegoRewriter, targetSpec templates.Target) ([]string, error) {
 	entryPoint, err := parseModule(targetSpec.Rego)
 	if err != nil {
 		return nil, fmt.Errorf("%w: %v", clienterrors.ErrInvalidConstraintTemplate, err)
@@ -172,7 +169,7 @@ func parseConstraintTemplateTarget(rr *regorewriter.RegoRewriter, pkgPrefix stri
 
 	rr.AddEntryPointModule(templatePath, entryPoint)
 	for idx, libSrc := range targetSpec.Libs {
-		libPath := fmt.Sprintf(`%s["lib_%d"]`, pkgPrefix, idx)
+		libPath := fmt.Sprintf(`%s["lib_%d"]`, templateLibPrefix, idx)
 		if err = rr.AddLib(libPath, libSrc); err != nil {
 			return nil, fmt.Errorf("%w: %v",
 				clienterrors.ErrInvalidConstraintTemplate, err)
@@ -216,12 +213,6 @@ func compileTemplateTarget(module TargetModule, capabilities *ast.Capabilities,
 	}
 	modules[hookModulePath] = builtinModule
 
-	regoModule, err := ast.ParseModule(templatePath, module.Rego)
-	if err != nil {
-		return nil, fmt.Errorf("%w: %v", clienterrors.ErrParse, err)
-	}
-	modules[templatePath] = regoModule
-
 	for i, lib := range module.Libs {
 		libPath := fmt.Sprintf("%s%d", templatePath, i)
 		libModule, err := ast.ParseModule(libPath, lib)

constraint/pkg/client/drivers/local/driver.go

@@ -307,11 +307,6 @@ func (d *Driver) Dump(ctx context.Context) (string, error) {
 	return string(b), nil
 }
 
-// templateLibPrefix returns the new lib prefix for the libs that are specified in the CT.
-func templateLibPrefix(name string) string {
-	return fmt.Sprintf("libs.%s", name)
-}
-
 // parseModule parses the module and also fails empty modules.
 func parseModule(rego string) (*ast.Module, error) {
 	module, err := ast.ParseModule(templatePath, rego)

constraint/pkg/client/drivers/local/rego.go

@@ -5,6 +5,10 @@ const (
 	// Must match the "data.xxx.violation[r]" path in hookModule.
 	templatePath = "template"
 
+	// templateLibPrefix is the path under which library Rego code is stored.
+	// Must match "data.xxx.[library package]" path.
+	templateLibPrefix = "libs"
+
 	// hookModulePath.
 	hookModulePath = "hooks.hooks_builtin"