Enhanced type checking using Regal AST schema (#201)

Currently enabled for the `regal test` command, which is the most
likely entrypoint for rule authors, whether for built-in Regal rules
or custom ones.

While it would be useful to provide this also for `regal lint` — we'll
probably want to do this together with the strictness rules, as we
currently don't require compilation, and it's not clear that we'll want
to unless a rule is configured that needs this, like most of the strictness
rules would.

Additionally, the schema is provided when running the Rego tests via
`go test ./...`, which should help catch errors we make ourselves.

It would be nice if custom rule authors didn't have to include the
schema in the metadata annotation, but I've yet to figure out the best
way of doing that, as we can't bundle a custom policy and also have them
read from disk without the bundle "owning" that path. I suppose we could
just parse a policy from a string, including a subpackage scoped
custom.regal.rules package annotation... but I'll need to think more
about this.

Fixes #52

Signed-off-by: Anders Eknert <anders@styra.com>

bundle/regal/regal.rego

@@ -4,4 +4,6 @@
 # - Styra
 # related_resources:
 # - https://www.styra.com
+# schemas:
+# - input: schema.regal.ast
 package regal

cmd/lint.go

@@ -2,7 +2,6 @@ package cmd
 
 import (
 	"context"
-	"embed"
 	"errors"
 	"fmt"
 	"io"
@@ -17,6 +16,7 @@ import (
 	"github.com/spf13/cobra"
 	"gopkg.in/yaml.v3"
 
+	"github.com/styrainc/regal/internal/embeds"
 	rio "github.com/styrainc/regal/internal/io"
 	"github.com/styrainc/regal/pkg/config"
 	"github.com/styrainc/regal/pkg/linter"
@@ -63,8 +63,6 @@ func (f *repeatedStringFlag) Set(s string) error {
 	return nil
 }
 
-var EmbedBundleFS embed.FS //nolint:gochecknoglobals
-
 func init() {
 	params := lintCommandParams{}
 
@@ -173,7 +171,7 @@ func lint(args []string, params lintCommandParams) (report.Report, error) {
 
 	// Create new fs from root of bundle, to avoid having to deal with
 	// "bundle" in paths (i.e. `data.bundle.regal`)
-	bfs, err := fs.Sub(EmbedBundleFS, "bundle")
+	bfs, err := fs.Sub(embeds.EmbedBundleFS, "bundle")
 	if err != nil {
 		return report.Report{}, fmt.Errorf("failed reading embedded bundle %w", err)
 	}

cmd/test.go

@@ -26,6 +26,7 @@ import (
 	"github.com/open-policy-agent/opa/version"
 
 	"github.com/styrainc/regal/internal/compile"
+	"github.com/styrainc/regal/internal/embeds"
 	rio "github.com/styrainc/regal/internal/io"
 	"github.com/styrainc/regal/pkg/builtins"
 )
@@ -132,7 +133,7 @@ func opaTest(args []string) int {
 
 	// Create new fs from root of bundle, to avoid having to deal with
 	// "bundle" in paths (i.e. `data.bundle.regal`)
-	bfs, err := fs.Sub(EmbedBundleFS, "bundle")
+	bfs, err := fs.Sub(embeds.EmbedBundleFS, "bundle")
 	if err != nil {
 		fmt.Fprintln(os.Stderr, fmt.Errorf("failed reading embedded bundle %w", err))
 
@@ -158,7 +159,9 @@ func opaTest(args []string) int {
 
 	compiler := compile.NewCompilerWithRegalBuiltins().
 		WithPathConflictsCheck(storage.NonEmpty(ctx, store, txn)).
-		WithEnablePrintStatements(!testParams.benchmark)
+		WithEnablePrintStatements(!testParams.benchmark).
+		WithSchemas(compile.SchemaSet(embeds.ASTSchema)).
+		WithUseTypeCheckAnnotations(true)
 
 	if testParams.threshold > 0 && !testParams.coverage {
 		testParams.coverage = true

docs/custom-rules.md

@@ -73,6 +73,8 @@ An example policy to implement this requirement might look something like this:
 # related_resources:
 # - description: documentation
 #   ref: https://www.acmecorp.example.org/docs/regal/package
+# schemas:
+# - input: schema.regal.ast
 package custom.regal.rules.naming["acme-corp-package"]
 
 import future.keywords.contains
@@ -107,6 +109,11 @@ Starting from top to bottom, these are the components comprising our custom rule
    order to document the purpose of the rule, along with any other information that could potentially be useful.
    All rule packages **must** have a `description`. Providing links to additional documentation under
    `related_resources` is recommended, but not required.
+1. Note the `schema` attribute present in the metadata annotation. Adding this is optional, but highly recommended, as
+   it will make the compiler aware of the structure of the input, i.e. the AST. This allows the compiler to fail when
+   unknown attributes are referenced, due to typos or other mistakes. The compiler will also fail when an attribute is
+   referenced using a type it does not have, like referring to a string as if it was a number. Set to `schema.regal.ast`
+   to use the AST schema provided by Regal.
 1. Regal will evaluate any rule named `report` in each linter policy, so at least one `report` rule **must** be present.
 1. In our example `report` rule, we evaluate another rule (`acme_corp_package`) in order to know if the package name
    starts with `acme.corp`, and another rule (`system_log_package`) to know if it starts with `system.log`. If neither

e2e/cli_test.go

@@ -10,12 +10,14 @@ import (
 	"os"
 	"os/exec"
 	"path/filepath"
+	"strings"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
 	"gopkg.in/yaml.v3"
 
 	"github.com/open-policy-agent/opa/tester"
+
 	"github.com/styrainc/regal/pkg/config"
 	"github.com/styrainc/regal/pkg/report"
 )
@@ -209,6 +211,33 @@ func TestTestRegalBundledRules(t *testing.T) {
 	}
 }
 
+func TestTestRegalTestWithExtendedASTTypeChecking(t *testing.T) {
+	t.Parallel()
+
+	stdout := bytes.Buffer{}
+	stderr := bytes.Buffer{}
+
+	cwd := must(os.Getwd)
+
+	err := regal(&stdout, &stderr)("test", cwd+"/testdata/ast_type_failure")
+
+	if exp, act := 1, ExitStatus(err); exp != act {
+		t.Errorf("expected exit status %d, got %d", exp, act)
+	}
+
+	expStart := "1 error occurred: "
+	expEnd := "rego_type_error: undefined ref: input.foo\n\tinput.foo\n\t      ^\n\t      " +
+		"have: \"foo\"\n\t      want (one of): [\"annotations\" \"comments\" \"imports\" \"package\" \"regal\" \"rules\"]\n"
+
+	if !strings.HasPrefix(stdout.String(), expStart) {
+		t.Errorf("expected stdout error message starting with %q, got %q", expStart, stdout.String())
+	}
+
+	if !strings.HasSuffix(stdout.String(), expEnd) {
+		t.Errorf("expected stdout error message ending with %q, got %q", expEnd, stdout.String())
+	}
+}
+
 func binary() string {
 	if b := os.Getenv("REGAL_BIN"); b != "" {
 		return b

e2e/testdata/ast_type_failure/custom_type_fail.rego

@@ -0,0 +1,15 @@
+# METADATA
+# description: Custom rule with type checker failure
+# schemas:
+# - input: schema.regal.ast
+package custom.regal.rules.naming.type_fail
+
+import future.keywords.contains
+import future.keywords.if
+
+report contains foo if {
+	# There is no "foo" attrinbute in the AST,
+	# so this should fail at compile time
+	foo := input.foo
+	foo == "bar"
+}

internal/compile/compile.go

@@ -2,11 +2,12 @@ package compile
 
 import (
 	"github.com/open-policy-agent/opa/ast"
+	"github.com/open-policy-agent/opa/util"
 
 	"github.com/styrainc/regal/pkg/builtins"
 )
 
-func NewCompilerWithRegalBuiltins() *ast.Compiler {
+func Capabilities() *ast.Capabilities {
 	caps := ast.CapabilitiesForThisVersion()
 	caps.Builtins = append(caps.Builtins, &ast.Builtin{
 		Name: builtins.RegalParseModuleMeta.Name,
@@ -21,5 +22,17 @@ func NewCompilerWithRegalBuiltins() *ast.Compiler {
 		Decl: builtins.RegalLastMeta.Decl,
 	})
 
-	return ast.NewCompiler().WithCapabilities(caps)
+	return caps
+}
+
+func SchemaSet(s []byte) *ast.SchemaSet {
+	schema := util.MustUnmarshalJSON(s)
+	schemaSet := ast.NewSchemaSet()
+	schemaSet.Put(ast.MustParseRef("schema.regal.ast"), schema)
+
+	return schemaSet
+}
+
+func NewCompilerWithRegalBuiltins() *ast.Compiler {
+	return ast.NewCompiler().WithCapabilities(Capabilities())
 }

internal/embeds/embeds.go

@@ -0,0 +1,8 @@
+//nolint:gochecknoglobals
+package embeds
+
+import "embed"
+
+var EmbedBundleFS embed.FS
+
+var ASTSchema []byte

internal/io/io.go

@@ -62,16 +62,6 @@ func MustLoadRegalBundlePath(path string) bundle.Bundle {
 	return regalBundle
 }
 
-// MustJSON parses JSON from reader, exit on failure.
-func MustJSON(from any) []byte {
-	bs, err := json.MarshalIndent(from, "", "   ")
-	if err != nil {
-		log.Fatal(err)
-	}
-
-	return bs
-}
-
 // JSONRoundTrip convert any value to JSON and back again.
 func JSONRoundTrip(from any, to any) error {
 	bs, err := json.Marshal(from)

internal/test/rego_test.go

@@ -3,6 +3,7 @@ package test
 import (
 	"context"
 	"io/fs"
+	"os"
 	"path/filepath"
 	"testing"
 
@@ -11,6 +12,7 @@ import (
 	"github.com/open-policy-agent/opa/storage/inmem"
 	"github.com/open-policy-agent/opa/tester"
 
+	"github.com/styrainc/regal/internal/compile"
 	"github.com/styrainc/regal/pkg/builtins"
 )
 
@@ -38,11 +40,22 @@ func TestRunRegoUnitTests(t *testing.T) {
 		store.Abort(ctx, txn)
 	})
 
+	schema, err := os.ReadFile("../../schemas/regal-ast.json")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	compiler := compile.NewCompilerWithRegalBuiltins().
+		WithSchemas(compile.SchemaSet(schema)).
+		WithUseTypeCheckAnnotations(true)
+
 	runner := tester.NewRunner().
+		SetCompiler(compiler).
 		SetStore(store).
 		CapturePrintOutput(true).
 		SetRuntime(ast.NewTerm(ast.NewObject())).
 		SetBundles(bundle).
+		// TODO: Not needed?
 		AddCustomBuiltins(builtins.TestContextBuiltins())
 
 	ch, err := runner.RunTests(ctx, txn)

main.go

@@ -6,6 +6,7 @@ import (
 	"os"
 
 	"github.com/styrainc/regal/cmd"
+	"github.com/styrainc/regal/internal/embeds"
 )
 
 // Note: this will bundle the tests as well, but since that has negligible impact on the size of the binary,
@@ -14,14 +15,19 @@ import (
 //go:embed all:bundle
 var bundle embed.FS
 
+//go:embed schemas/regal-ast.json
+var schema []byte
+
 func main() {
 	// Remove date and time from any `log.*` calls, as that doesn't add much of value here
 	// Evaluate options for logging later..
 	log.SetFlags(0)
 
-	// The embedded FS can't point to parent directories, so while not pretty, we'll
-	// need to pass it from here to the next command
-	cmd.EmbedBundleFS = bundle
+	// Embedded files and directories can only traverse down in the tree, i.e. no parent paths,
+	// so while this isn't pretty, we'll use a separate package to carry state from here. If you
+	// know of a better way of doing this, don't hesitate to fix this.
+	embeds.EmbedBundleFS = bundle
+	embeds.ASTSchema = schema
 
 	if err := cmd.RootCommand.Execute(); err != nil {
 		os.Exit(1)