fix(misconf): avoid panic if the scheme is not valid (aquasecurity#6496)

Author: nikpivkin

pkg/iac/rego/build.go

@@ -1,6 +1,7 @@
 package rego
 
 import (
+	"fmt"
 	"io/fs"
 	"path/filepath"
 	"strings"
@@ -16,28 +17,37 @@ func BuildSchemaSetFromPolicies(policies map[string]*ast.Module, paths []string,
 	schemaSet := ast.NewSchemaSet()
 	schemaSet.Put(ast.MustParseRef("schema.input"), make(map[string]interface{})) // for backwards compat only
 	var customFound bool
+
 	for _, policy := range policies {
 		for _, annotation := range policy.Annotations {
 			for _, ss := range annotation.Schemas {
 				schemaName, err := ss.Schema.Ptr()
-				if err != nil {
+				if err != nil || schemaName == "input" {
 					continue
 				}
-				if schemaName != "input" {
-					if schema, ok := schemas.SchemaMap[types.Source(schemaName)]; ok {
-						customFound = true
-						schemaSet.Put(ast.MustParseRef(ss.Schema.String()), util.MustUnmarshalJSON([]byte(schema)))
-					} else {
-						b, err := findSchemaInFS(paths, fsys, schemaName)
-						if err != nil {
-							return schemaSet, true, err
-						}
-						if b != nil {
-							customFound = true
-							schemaSet.Put(ast.MustParseRef(ss.Schema.String()), util.MustUnmarshalJSON(b))
-						}
+
+				var schema []byte
+				if s, ok := schemas.SchemaMap[types.Source(schemaName)]; ok {
+					schema = []byte(s)
+				} else {
+					b, err := findSchemaInFS(paths, fsys, schemaName)
+					if err != nil {
+						return schemaSet, true, err
+					}
+
+					if b == nil {
+						return nil, false, fmt.Errorf("could not find schema %q", schemaName)
 					}
+
+					schema = b
+				}
+
+				var rawSchema any
+				if err := util.UnmarshalJSON(schema, &rawSchema); err != nil {
+					return schemaSet, false, fmt.Errorf("could not parse schema %q: %w", schemaName, err)
 				}
+				customFound = true
+				schemaSet.Put(ast.MustParseRef(ss.Schema.String()), rawSchema)
 			}
 		}
 	}

pkg/iac/rego/embed_test.go

@@ -80,24 +80,6 @@ deny[res]{
 	res := true
 }`,
 		},
-		{
-			name: "sad path schema does not exist",
-			inputPolicy: `# METADATA
-# title: "dummy title"
-# description: "some description"
-# scope: package
-# schemas:
-# - input: schema["invalid schema"]
-# custom:
-#   input:
-#     selector:
-#     - type: dockerfile
-package builtin.dockerfile.DS1234
-deny[res]{
-	res := true
-}`,
-			expectedError: true,
-		},
 	}
 
 	for _, tc := range testCases {

pkg/iac/rego/load_test.go

@@ -1,13 +1,19 @@
-package rego
+package rego_test
 
 import (
 	"bytes"
 	"embed"
+	"io"
+	"strings"
 	"testing"
+	"testing/fstest"
 
-	"github.com/aquasecurity/trivy/pkg/iac/types"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+
+	"github.com/aquasecurity/trivy/pkg/iac/rego"
+	"github.com/aquasecurity/trivy/pkg/iac/scanners/options"
+	"github.com/aquasecurity/trivy/pkg/iac/types"
 )
 
 //go:embed all:testdata/policies
@@ -16,31 +22,76 @@ var testEmbedFS embed.FS
 func Test_RegoScanning_WithSomeInvalidPolicies(t *testing.T) {
 	t.Run("allow no errors", func(t *testing.T) {
 		var debugBuf bytes.Buffer
-		scanner := NewScanner(types.SourceDockerfile)
-		scanner.SetRegoErrorLimit(0)
-		scanner.SetDebugWriter(&debugBuf)
-		p, _ := LoadPoliciesFromDirs(testEmbedFS, ".")
-		require.NotNil(t, p)
-
-		scanner.policies = p
-		err := scanner.compilePolicies(testEmbedFS, []string{"policies"})
+		scanner := rego.NewScanner(
+			types.SourceDockerfile,
+			options.ScannerWithRegoErrorLimits(0),
+			options.ScannerWithDebug(&debugBuf),
+		)
+
+		err := scanner.LoadPolicies(false, false, testEmbedFS, []string{"."}, nil)
 		require.ErrorContains(t, err, `want (one of): ["Cmd" "EndLine" "Flags" "JSON" "Original" "Path" "Stage" "StartLine" "SubCmd" "Value"]`)
 		assert.Contains(t, debugBuf.String(), "Error(s) occurred while loading policies")
 	})
 
 	t.Run("allow up to max 1 error", func(t *testing.T) {
 		var debugBuf bytes.Buffer
-		scanner := NewScanner(types.SourceDockerfile)
-		scanner.SetRegoErrorLimit(1)
-		scanner.SetDebugWriter(&debugBuf)
-
-		p, _ := LoadPoliciesFromDirs(testEmbedFS, ".")
-		scanner.policies = p
+		scanner := rego.NewScanner(
+			types.SourceDockerfile,
+			options.ScannerWithRegoErrorLimits(1),
+			options.ScannerWithDebug(&debugBuf),
+		)
 
-		err := scanner.compilePolicies(testEmbedFS, []string{"policies"})
+		err := scanner.LoadPolicies(false, false, testEmbedFS, []string{"."}, nil)
 		require.NoError(t, err)
 
 		assert.Contains(t, debugBuf.String(), "Error occurred while parsing: testdata/policies/invalid.rego, testdata/policies/invalid.rego:7")
 	})
 
+	t.Run("schema does not exist", func(t *testing.T) {
+		check := `# METADATA
+# schemas:
+# - input: schema["fooschema"]
+package mypackage
+
+deny {
+    input.evil == "foo bar"
+}`
+		scanner := rego.NewScanner(types.SourceJSON)
+
+		err := scanner.LoadPolicies(false, false, fstest.MapFS{}, []string{"."}, []io.Reader{strings.NewReader(check)})
+		assert.ErrorContains(t, err, "could not find schema \"fooschema\"")
+	})
+
+	t.Run("schema is invalid", func(t *testing.T) {
+		check := `# METADATA
+# schemas:
+# - input: schema["fooschema"]
+package mypackage
+
+deny {
+    input.evil == "foo bar"
+}`
+		scanner := rego.NewScanner(types.SourceJSON)
+
+		fsys := fstest.MapFS{
+			"schemas/fooschema.json": &fstest.MapFile{
+				Data: []byte("bad json"),
+			},
+		}
+
+		err := scanner.LoadPolicies(false, false, fsys, []string{"."}, []io.Reader{strings.NewReader(check)})
+		assert.ErrorContains(t, err, "could not parse schema \"fooschema\"")
+	})
+
+	t.Run("schema is not specified", func(t *testing.T) {
+		check := `package mypackage
+
+deny {
+    input.evil == "foo bar"
+}`
+		scanner := rego.NewScanner(types.SourceJSON)
+		err := scanner.LoadPolicies(false, false, fstest.MapFS{}, []string{"."}, []io.Reader{strings.NewReader(check)})
+		assert.NoError(t, err)
+	})
+
 }

pkg/iac/scanners/dockerfile/scanner_test.go

@@ -546,7 +546,7 @@ package builtin.dockerfile.DS006
 deny[res]{
 res := true
 }`,
-			expectedError: `1 error occurred: rules/rule.rego:12: rego_type_error: undefined schema: schema["spooky-schema"]`,
+			expectedError: "could not find schema \"spooky-schema\"",
 		},
 	}
 

pkg/iac/scanners/terraform/scanner_integration_test.go

@@ -161,6 +161,7 @@ deny[cause] {
 
 	t.Run("without skip", func(t *testing.T) {
 		scanner := New(
+			ScannerWithSkipCachedModules(true),
 			options.ScannerWithPolicyDirs("rules"),
 			options.ScannerWithRegoOnly(true),
 			options.ScannerWithEmbeddedPolicies(false),