Repository: plexsystems/konstraint (go)

Issue:
Input parameters verification ignores assignment rules
Parameter annotation checker only works for main rego file and doesn't work for imported files.

For example, running `create` on following files produces no errors.
```rego
# src.rego -- annotation @parameter test is missing
package example
import data.lib.foo
violation[{"msg": "msg"}] {
  test = foo.test
  true
}

# lib/foo.rego
package foo
test = input.parameters.test
```


Pull request description:
Parse parameter names from rule header values in addition to the body
Signed-off-by: James Alseth <james@jalseth.me>

---

Fixes #276

Patch:
diff --git a/docs/constraint_creation.md b/docs/constraint_creation.md
index d8388ae..0bc6da6 100644
--- a/docs/constraint_creation.md
+++ b/docs/constraint_creation.md
@@ -121,12 +121,15 @@ Gatekeeper has the ability for a single `ConstraintTemplate` resource to be used
 
 To use parameters, add one or more `@parameter <name> <type>` statements where `<name>` is the name of the parameter and `<type>` is the OpenAPI v3 type of the parameter (string, integer, etc.). Arrays are supported via `@parameter <name> array <type>`. Each parameter tag must be on its own line. Each parameter in the rule body must have a `@parameter` tag in the comment block header.
 
+Parameter description can be added after `--`, and and can fold on multiple lines following the parameter.
+
 ```rego
 # @title Required Labels
 #
 # This policy allows you to require certain labels are set on a resource.
 #
-# @parameter labels array string
+# @parameter labels array string --
+# -- array of required label keys
 package required_labels
 
 import data.lib.core
diff --git a/examples/policies-no-rego.md b/examples/policies-no-rego.md
index 2b6926d..557a9d7 100644
--- a/examples/policies-no-rego.md
+++ b/examples/policies-no-rego.md
@@ -37,8 +37,8 @@
 
 **Parameters:**
 
-* labels: array of string
-
+* labels: array of string<br>
+  array of required label keys
 
 This policy allows you to require certain labels are set on a resource.
 Adapted from https://github.com/open-policy-agent/gatekeeper/blob/master/example/templates/k8srequiredlabels_template.yaml
diff --git a/examples/policies.md b/examples/policies.md
index 046f012..869e913 100755
--- a/examples/policies.md
+++ b/examples/policies.md
@@ -37,8 +37,8 @@
 
 **Parameters:**
 
-* labels: array of string
-
+* labels: array of string<br>
+  array of required label keys
 
 This policy allows you to require certain labels are set on a resource.
 Adapted from https://github.com/open-policy-agent/gatekeeper/blob/master/example/templates/k8srequiredlabels_template.yaml
diff --git a/examples/required-labels/src.rego b/examples/required-labels/src.rego
index 5cb21e7..2f61035 100644
--- a/examples/required-labels/src.rego
+++ b/examples/required-labels/src.rego
@@ -3,7 +3,8 @@
 # This policy allows you to require certain labels are set on a resource.
 # Adapted from https://github.com/open-policy-agent/gatekeeper/blob/master/example/templates/k8srequiredlabels_template.yaml
 #
-# @parameter labels array string
+# @parameter labels array string --
+# -- array of required label keys
 
 package required_labels
 
diff --git a/examples/required-labels/template.yaml b/examples/required-labels/template.yaml
index a3a6c72..6378cf3 100755
--- a/examples/required-labels/template.yaml
+++ b/examples/required-labels/template.yaml
@@ -12,6 +12,7 @@ spec:
         openAPIV3Schema:
           properties:
             labels:
+              description: array of required label keys
               items:
                 type: string
               type: array
diff --git a/internal/commands/create.go b/internal/commands/create.go
index f396d55..9226375 100644
--- a/internal/commands/create.go
+++ b/internal/commands/create.go
@@ -232,13 +232,17 @@ func getOpenAPISchemaProperties(r rego.Rego) map[string]apiextensionsv1.JSONSche
 	for _, p := range r.Parameters() {
 		if p.IsArray {
 			properties[p.Name] = apiextensionsv1.JSONSchemaProps{
-				Type: "array",
+				Type:        "array",
+				Description: p.Description,
 				Items: &apiextensionsv1.JSONSchemaPropsOrArray{
 					Schema: &apiextensionsv1.JSONSchemaProps{Type: p.Type},
 				},
 			}
 		} else {
-			properties[p.Name] = apiextensionsv1.JSONSchemaProps{Type: p.Type}
+			properties[p.Name] = apiextensionsv1.JSONSchemaProps{
+				Type:        p.Type,
+				Description: p.Description,
+			}
 		}
 	}
 
diff --git a/internal/commands/document_template.go b/internal/commands/document_template.go
index 8e22c09..338b611 100644
--- a/internal/commands/document_template.go
+++ b/internal/commands/document_template.go
@@ -26,7 +26,10 @@ const docTemplate = `# Policies
 
 **Parameters:**
 
-{{ range .Header.Parameters }}* {{ .Name }}: {{ if .IsArray }}array of {{ end }}{{ .Type }}
+{{ range .Header.Parameters -}}
+* {{ .Name }}: {{ if .IsArray }}array of {{ end }}{{ .Type }}
+{{- if .Description }}<br>
+  {{ .Description }}{{- end -}}
 {{ end }}
 {{- end }}
 
diff --git a/internal/rego/rego.go b/internal/rego/rego.go
index 142fdfc..e81f63b 100644
--- a/internal/rego/rego.go
+++ b/internal/rego/rego.go
@@ -42,9 +42,10 @@ type Rego struct {
 
 // Parameter represents a parameter that the policy uses
 type Parameter struct {
-	Name    string
-	Type    string
-	IsArray bool
+	Name        string
+	Type        string
+	IsArray     bool
+	Description string
 }
 
 // GetAllSeverities gets all of the rego files found in the given
@@ -194,8 +195,16 @@ func (r Rego) PolicyID() string {
 func (r Rego) Description() string {
 	var description string
 	var handlingCodeBlock bool
+	var handlingParamDescription bool
+
 	for _, comment := range r.headerComments {
-		if commentStartsWith(comment, "@") {
+		if !handlingCodeBlock && !handlingParamDescription && commentStartsWith(comment, "@parameter") && strings.Contains(comment, "--") {
+			handlingParamDescription = true
+		} else if !handlingCodeBlock && handlingParamDescription && !commentStartsWith(comment, "--") {
+			handlingParamDescription = false
+		}
+
+		if handlingParamDescription || commentStartsWith(comment, "@") {
 			continue
 		}
 
@@ -327,7 +336,7 @@ func parseDirectory(directory string, parseImports bool) ([]Rego, error) {
 			}
 		}
 
-		bodyParams := getBodyParamNames(file.Parsed.Rules)
+		bodyParams := getRuleParamNames(file.Parsed.Rules)
 		headerParams, err := getHeaderParams(headerComments)
 		if err != nil {
 			return nil, fmt.Errorf("parse header parameters: %w", err)
@@ -375,29 +384,33 @@ func parseDirectory(directory string, parseImports bool) ([]Rego, error) {
 	return regos, nil
 }
 
-func getBodyParamNames(rules []*ast.Rule) []string {
-	r := regexp.MustCompile(`(core|input)\.parameters\.([a-zA-Z0-9_-]+)`)
-	var bodyParams []string
-	for _, rule := range rules {
-		matches := r.FindAllStringSubmatch(rule.Body.String(), -1)
+func getRuleParamNames(rules []*ast.Rule) []string {
+	re := regexp.MustCompile(`input\.parameters\.([a-zA-Z0-9_-]+)`)
+	var ruleParams []string
+	for _, r := range rules {
+		matches := re.FindAllStringSubmatch(r.String(), -1)
 		for _, match := range matches {
-			if !contains(bodyParams, match[2]) {
-				bodyParams = append(bodyParams, match[2])
+			if !contains(ruleParams, match[1]) {
+				ruleParams = append(ruleParams, match[1])
 			}
 		}
 	}
 
-	return bodyParams
+	return ruleParams
 }
 
 func getHeaderParams(comments []string) ([]Parameter, error) {
 	var parameters []Parameter
-	for _, comment := range comments {
+	for i := 0; i < len(comments); i++ {
+		comment := comments[i]
+
 		if !commentStartsWith(comment, "@parameter ") {
 			continue
 		}
 
 		params := strings.SplitAfter(comment, "@parameter ")[1]
+		paramsDesc := strings.SplitN(params, " --", 2)
+		params = paramsDesc[0]
 		paramsSplit := strings.Split(params, " ")
 		if len(paramsSplit) == 0 {
 			return nil, fmt.Errorf("parameter name and type must be specified")
@@ -417,6 +430,22 @@ func getHeaderParams(comments []string) ([]Parameter, error) {
 			p.Type = paramsSplit[1]
 		}
 
+		if len(paramsDesc) > 1 {
+			p.Description = strings.TrimSpace(paramsDesc[1])
+
+			for i++; i != len(comments); i++ {
+				extraComment := strings.TrimSpace(comments[i])
+				if !strings.HasPrefix(extraComment, "--") {
+					i--
+					break
+				}
+				extraComment = strings.TrimSpace(extraComment[2:])
+				p.Description += " " + extraComment
+			}
+
+			p.Description = strings.TrimSpace(p.Description)
+		}
+
 		parameters = append(parameters, p)
 	}
 


Test patch:
diff --git a/internal/rego/rego_test.go b/internal/rego/rego_test.go
index 09c16ff..4ee098a 100644
--- a/internal/rego/rego_test.go
+++ b/internal/rego/rego_test.go
@@ -162,43 +162,52 @@ func TestGetPolicyID_Null(t *testing.T) {
 	}
 }
 
-func TestGetBodyParamNamesFromInput(t *testing.T) {
-	ruleString := `violation[msg] {
-	foo := "bar"
-	bar := input.parameters.baz
-	baz := input.parameters.foobars[_]
-	box := input.parameters.baz
-}`
-
-	rule, err := ast.ParseRule(ruleString)
-	if err != nil {
-		t.Fatalf("parse rule: %s", err)
-	}
-
-	expected := []string{"baz", "foobars"}
-	actual := getBodyParamNames([]*ast.Rule{rule})
-	if !(reflect.DeepEqual(expected, actual)) {
-		t.Errorf("unexpected bodyParams. expected %+v, actual %+v", expected, actual)
+func TestGetRuleParamNamesFromInput(t *testing.T) {
+	testCases := []struct {
+		desc string
+		rule string
+		want []string
+	}{
+		{
+			desc: "No Parameters",
+			rule: `foo = "bar" { true }`,
+		},
+		{
+			desc: "Parameters in rule body",
+			rule: `violation[msg] {
+				foo := "bar"
+				bar := input.parameters.baz
+				baz := input.parameters.foobars[_]
+				box := input.parameters.baz
+			}`,
+			want: []string{"baz", "foobars"},
+		},
+		{
+			desc: "Parameters in rule value",
+			rule: `foo = input.parameters.bar { true }`,
+			want: []string{"bar"},
+		},
+		{
+			desc: "Parameters in body and value",
+			rule: `foo = input.parameters.bar {
+				x := input.parameters.baz
+			}`,
+			want: []string{"bar", "baz"},
+		},
 	}
-}
 
-func TestGetBodyParamNamesFromCore(t *testing.T) {
-	ruleString := `violation[msg] {
-	foo := "bar"
-	bar := core.parameters.baz
-	baz := core.parameters.foobars[_]
-	box := core.parameters.baz
-}`
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			rule, err := ast.ParseRule(tc.rule)
+			if err != nil {
+				t.Fatalf("parse rule: %s", err)
+			}
 
-	rule, err := ast.ParseRule(ruleString)
-	if err != nil {
-		t.Fatalf("parse rule: %s", err)
-	}
-
-	expected := []string{"baz", "foobars"}
-	actual := getBodyParamNames([]*ast.Rule{rule})
-	if !(reflect.DeepEqual(expected, actual)) {
-		t.Errorf("unexpected bodyParams. expected %+v, actual %+v", expected, actual)
+			actual := getRuleParamNames([]*ast.Rule{rule})
+			if !(reflect.DeepEqual(tc.want, actual)) {
+				t.Errorf("unexpected bodyParams. expected %+v, actual %+v", tc.want, actual)
+			}
+		})
 	}
 }
 
@@ -206,21 +215,30 @@ func TestGetHeaderParams(t *testing.T) {
 	comments := []string{
 		"@title Title",
 		"Description",
-		"@parameter foo string",
+		"@parameter foo string -- with description",
 		"@parameter bar array string",
+		"@parameter baz array string -- with multiline",
+		"-- description",
 		"@kinds another/thing",
 	}
 
 	expected := []Parameter{
 		{
-			Name: "foo",
-			Type: "string",
+			Name:        "foo",
+			Type:        "string",
+			Description: "with description",
 		},
 		{
 			Name:    "bar",
 			Type:    "string",
 			IsArray: true,
 		},
+		{
+			Name:        "baz",
+			Type:        "string",
+			IsArray:     true,
+			Description: "with multiline description",
+		},
 	}
 
 	actual, err := getHeaderParams(comments)
diff --git a/test/src.rego b/test/src.rego
index 4b291ee..10f9d60 100644
--- a/test/src.rego
+++ b/test/src.rego
@@ -7,6 +7,8 @@
 # @matchExpression doggos Exists
 # @namespaces dev stage prod
 # @excludedNamespaces kube-system gatekeeper-system
+# @parameter super string -- super duper cool parameter
+# -- with a description on two strings
 package test
 
 import data.lib.libraryA
diff --git a/test/template_Test.yaml b/test/template_Test.yaml
index 6678bed..9c550de 100644
--- a/test/template_Test.yaml
+++ b/test/template_Test.yaml
@@ -8,6 +8,12 @@ spec:
     spec:
       names:
         kind: Test
+      validation:
+        openAPIV3Schema:
+          properties:
+            super:
+              description: super duper cool parameter with a description on two strings
+              type: string
   targets:
   - libs:
     - |-