From cc38ccde3608df3370d4eb8e71a12d08f46dde1b Mon Sep 17 00:00:00 2001
From: Juan Antonio Osorio <ozz@stacklok.com>
Date: Tue, 14 Jan 2025 10:10:17 +0200
Subject: [PATCH] Introduce rego function to parse TOML files (#5294)

This introduces `parse_toml` to our rego library. It allows us to parse
TOML files in order to more easily write policies for things like
`pyproject.toml`.

Signed-off-by: Juan Antonio Osorio <ozz@stacklok.com>
---
 internal/engine/eval/rego/lib.go      | 39 ++++++++++++++
 internal/engine/eval/rego/lib_test.go | 75 +++++++++++++++++++++++++++
 2 files changed, 114 insertions(+)

diff --git a/internal/engine/eval/rego/lib.go b/internal/engine/eval/rego/lib.go
index e7f3f2697e..18454e5d7a 100644
--- a/internal/engine/eval/rego/lib.go
+++ b/internal/engine/eval/rego/lib.go
@@ -23,6 +23,7 @@ import (
 	"github.com/open-policy-agent/opa/v1/ast"
 	"github.com/open-policy-agent/opa/v1/rego"
 	"github.com/open-policy-agent/opa/v1/types"
+	"github.com/pelletier/go-toml/v2"
 	"github.com/stacklok/frizbee/pkg/replacer"
 	"github.com/stacklok/frizbee/pkg/utils/config"
 	"gopkg.in/yaml.v3"
@@ -42,6 +43,7 @@ var MinderRegoLib = []func(res *interfaces.Result) func(*rego.Rego){
 	FileWalk,
 	ListGithubActions,
 	ParseYaml,
+	ParseToml,
 	JQIsTrue,
 }
 
@@ -724,3 +726,40 @@ func ParseYaml(_ *interfaces.Result) func(*rego.Rego) {
 		},
 	)
 }
+
+// ParseToml is a rego function that parses a TOML configuration string into a structured data format.
+// It takes one argument: the TOML content as a string.
+// It returns the parsed TOML data as an AST term.
+// It's exposed as `parse_toml`.
+func ParseToml(_ *interfaces.Result) func(*rego.Rego) {
+	return rego.Function1(
+		&rego.Function{
+			Name: "parse_toml",
+			// Takes one string argument (the YAML content) and returns any type
+			Decl: types.NewFunction(types.Args(types.S), types.A),
+		},
+		func(_ rego.BuiltinContext, content *ast.Term) (*ast.Term, error) {
+			var tomlStr string
+
+			// Convert the YAML input from the term into a string
+			if err := ast.As(content.Value, &tomlStr); err != nil {
+				return nil, err
+			}
+
+			// Convert the YAML string into a Go map
+			var jsonObj any
+			err := toml.Unmarshal([]byte(tomlStr), &jsonObj)
+			if err != nil {
+				return nil, fmt.Errorf("error converting YAML to JSON: %w", err)
+			}
+
+			// Convert the Go value to an ast.Value
+			value, err := ast.InterfaceToValue(jsonObj)
+			if err != nil {
+				return nil, fmt.Errorf("error converting to AST value: %w", err)
+			}
+
+			return ast.NewTerm(value), nil
+		},
+	)
+}
diff --git a/internal/engine/eval/rego/lib_test.go b/internal/engine/eval/rego/lib_test.go
index e30579bc3d..9052650c2d 100644
--- a/internal/engine/eval/rego/lib_test.go
+++ b/internal/engine/eval/rego/lib_test.go
@@ -1122,3 +1122,78 @@ allow {
 		})
 	}
 }
+
+func TestParseToml(t *testing.T) {
+	t.Parallel()
+
+	scenario := []struct {
+		name    string
+		toml    string
+		want    string
+		wantErr bool
+	}{
+		{
+			name: "simple key-value",
+			toml: "foo = \"bar\"",
+			want: `{"foo": "bar"}`,
+		},
+		{
+			name: "nested structure",
+			toml: `
+[foo.bar]
+baz = "qux"`,
+			want: `{"foo": {"bar": {"baz": "qux"}}}`,
+		},
+		{
+			name: "array values",
+			toml: `
+items = ["foo", "bar", "baz"]`,
+			want: `{"items": ["foo", "bar", "baz"]}`,
+		},
+		{
+			name: "parse array of tables",
+			toml: `
+[[items]]
+name = "foo"
+[[items]]
+name = "bar"`,
+			want: `{"items": [{"name": "foo"}, {"name": "bar"}]}`,
+		},
+	}
+
+	for _, s := range scenario {
+		t.Run(s.name, func(t *testing.T) {
+			t.Parallel()
+
+			regoCode := fmt.Sprintf(`
+package minder
+
+default allow = false
+
+allow {
+	parsed := parse_toml(%q)
+	print(parsed)
+	expected := json.unmarshal(%q)
+	parsed == expected
+}`, s.toml, s.want)
+
+			e, err := rego.NewRegoEvaluator(
+				&minderv1.RuleType_Definition_Eval_Rego{
+					Type: rego.DenyByDefaultEvaluationType.String(),
+					Def:  regoCode,
+				},
+				nil,
+			)
+
+			require.NoError(t, err, "could not create evaluator")
+
+			_, err = e.Eval(context.Background(), map[string]any{}, nil, &interfaces.Result{})
+
+			if s.wantErr {
+				require.Error(t, err)
+			} else {
+				require.NoError(t, err)
+			}
+		})
+	}
+}