feat(misconf): loading embedded checks as a fallback

pkg/cloud/report/convert.go

@@ -8,6 +8,7 @@ import (
 	"github.com/aws/aws-sdk-go-v2/aws/arn"
 
 	ftypes "github.com/aquasecurity/trivy/pkg/fanal/types"
+	"github.com/aquasecurity/trivy/pkg/iac/rego"
 	"github.com/aquasecurity/trivy/pkg/iac/scan"
 	"github.com/aquasecurity/trivy/pkg/types"
 )
@@ -57,7 +58,7 @@ func ConvertResults(results scan.Results, provider string, scoped []string) map[
 
 				// empty namespace implies a go rule from defsec, "builtin" refers to a built-in rego rule
 				// this ensures we don't generate bad links for custom policies
-				if result.RegoNamespace() == "" || strings.HasPrefix(result.RegoNamespace(), "builtin.") {
+				if result.RegoNamespace() == "" || rego.IsBuiltinNamespace(result.RegoNamespace()) {
 					primaryURL = fmt.Sprintf("https://avd.aquasec.com/misconfig/%s", strings.ToLower(result.Rule().AVDID))
 				}
 

pkg/commands/artifact/run.go

@@ -20,6 +20,7 @@ import (
 	"github.com/aquasecurity/trivy/pkg/fanal/cache"
 	ftypes "github.com/aquasecurity/trivy/pkg/fanal/types"
 	"github.com/aquasecurity/trivy/pkg/flag"
+	"github.com/aquasecurity/trivy/pkg/iac/rego"
 	"github.com/aquasecurity/trivy/pkg/javadb"
 	"github.com/aquasecurity/trivy/pkg/log"
 	"github.com/aquasecurity/trivy/pkg/misconf"
@@ -49,11 +50,6 @@ const (
 )
 
 var (
-	defaultPolicyNamespaces = []string{
-		"appshield",
-		"defsec",
-		"builtin",
-	}
 	SkipScan = errors.New("skip subsequent processes")
 )
 
@@ -597,7 +593,7 @@ func initScannerConfig(opts flag.Options, cacheClient cache.Cache) (ScannerConfi
 		configScannerOptions = misconf.ScannerOption{
 			Debug:                    opts.Debug,
 			Trace:                    opts.Trace,
-			Namespaces:               append(opts.PolicyNamespaces, defaultPolicyNamespaces...),
+			Namespaces:               append(opts.PolicyNamespaces, rego.BuiltinNamespaces()...),
 			PolicyPaths:              append(opts.PolicyPaths, downloadedPolicyPaths...),
 			DataPaths:                opts.DataPaths,
 			HelmValues:               opts.HelmValues,

pkg/iac/rego/embed.go

@@ -8,7 +8,7 @@ import (
 
 	"github.com/open-policy-agent/opa/ast"
 
-	rules2 "github.com/aquasecurity/trivy-policies"
+	checks "github.com/aquasecurity/trivy-policies"
 	"github.com/aquasecurity/trivy/pkg/iac/rules"
 )
 
@@ -62,11 +62,11 @@ func RegisterRegoRules(modules map[string]*ast.Module) {
 }
 
 func LoadEmbeddedPolicies() (map[string]*ast.Module, error) {
-	return LoadPoliciesFromDirs(rules2.EmbeddedPolicyFileSystem, ".")
+	return LoadPoliciesFromDirs(checks.EmbeddedPolicyFileSystem, ".")
 }
 
 func LoadEmbeddedLibraries() (map[string]*ast.Module, error) {
-	return LoadPoliciesFromDirs(rules2.EmbeddedLibraryFileSystem, ".")
+	return LoadPoliciesFromDirs(checks.EmbeddedLibraryFileSystem, ".")
 }
 
 func LoadPoliciesFromDirs(target fs.FS, paths ...string) (map[string]*ast.Module, error) {

pkg/iac/rego/load.go

@@ -5,12 +5,30 @@ import (
 	"fmt"
 	"io"
 	"io/fs"
+	"path/filepath"
 	"strings"
 
 	"github.com/open-policy-agent/opa/ast"
 	"github.com/open-policy-agent/opa/bundle"
+	"github.com/samber/lo"
 )
 
+var builtinNamespaces = map[string]struct{}{
+	"builtin":   {},
+	"defsec":    {},
+	"appshield": {},
+}
+
+func BuiltinNamespaces() []string {
+	return lo.Keys(builtinNamespaces)
+}
+
+func IsBuiltinNamespace(namespace string) bool {
+	return lo.ContainsBy(BuiltinNamespaces(), func(ns string) bool {
+		return strings.HasPrefix(namespace, ns+".")
+	})
+}
+
 func IsRegoFile(name string) bool {
 	return strings.HasSuffix(name, bundle.RegoExt) && !strings.HasSuffix(name, "_test"+bundle.RegoExt)
 }
@@ -38,28 +56,20 @@ func (s *Scanner) loadPoliciesFromReaders(readers []io.Reader) (map[string]*ast.
 	return modules, nil
 }
 
-func (s *Scanner) loadEmbedded(enableEmbeddedLibraries, enableEmbeddedPolicies bool) error {
-	if enableEmbeddedLibraries {
-		loadedLibs, errLoad := LoadEmbeddedLibraries()
-		if errLoad != nil {
-			return fmt.Errorf("failed to load embedded rego libraries: %w", errLoad)
-		}
-		for name, policy := range loadedLibs {
-			s.policies[name] = policy
-		}
-		s.debug.Log("Loaded %d embedded libraries.", len(loadedLibs))
+func (s *Scanner) loadEmbedded() error {
+	loaded, err := LoadEmbeddedLibraries()
+	if err != nil {
+		return fmt.Errorf("failed to load embedded rego libraries: %w", err)
 	}
+	s.embeddedLibs = loaded
+	s.debug.Log("Loaded %d embedded libraries.", len(loaded))
 
-	if enableEmbeddedPolicies {
-		loaded, err := LoadEmbeddedPolicies()
-		if err != nil {
-			return fmt.Errorf("failed to load embedded rego policies: %w", err)
-		}
-		for name, policy := range loaded {
-			s.policies[name] = policy
-		}
-		s.debug.Log("Loaded %d embedded policies.", len(loaded))
+	loaded, err = LoadEmbeddedPolicies()
+	if err != nil {
+		return fmt.Errorf("failed to load embedded rego policies: %w", err)
 	}
+	s.embeddedChecks = loaded
+	s.debug.Log("Loaded %d embedded policies.", len(loaded))
 
 	return nil
 }
@@ -75,10 +85,18 @@ func (s *Scanner) LoadPolicies(enableEmbeddedLibraries, enableEmbeddedPolicies b
 		srcFS = s.policyFS
 	}
 
-	if err := s.loadEmbedded(enableEmbeddedLibraries, enableEmbeddedPolicies); err != nil {
+	if err := s.loadEmbedded(); err != nil {
 		return err
 	}
 
+	if enableEmbeddedPolicies {
+		s.policies = lo.Assign(s.policies, s.embeddedChecks)
+	}
+
+	if enableEmbeddedLibraries {
+		s.policies = lo.Assign(s.policies, s.embeddedLibs)
+	}
+
 	var err error
 	if len(paths) > 0 {
 		loaded, err := LoadPoliciesFromDirs(srcFS, paths...)
@@ -127,6 +145,60 @@ func (s *Scanner) LoadPolicies(enableEmbeddedLibraries, enableEmbeddedPolicies b
 	return s.compilePolicies(srcFS, paths)
 }
 
+func (s *Scanner) fallbackChecks(compiler *ast.Compiler) {
+
+	var excludedFiles []string
+
+	for _, e := range compiler.Errors {
+		if _, ok := e.Details.(*ast.RefErrInvalidDetail); !ok {
+			continue
+		}
+
+		loc := e.Location.File
+
+		if lo.Contains(excludedFiles, loc) {
+			continue
+		}
+
+		badPolicy, exists := s.policies[loc]
+		if !exists || badPolicy == nil {
+			continue
+		}
+
+		if !IsBuiltinNamespace(getModuleNamespace(badPolicy)) {
+			continue
+		}
+
+		s.debug.Log("Error occurred while parsing: %s, %s. \nTry loading embedded policy.", loc, e.Error())
+
+		embedded := s.findMatchedEmbeddedCheck(badPolicy)
+		if embedded == nil {
+			s.debug.Log("Failed to find embedded policy: %s", loc)
+			continue
+		}
+
+		s.debug.Log("Found embedded policy: %s", embedded.Package.Location.File)
+		delete(s.policies, loc) // remove bad policy
+		s.policies[embedded.Package.Location.File] = embedded
+		delete(s.embeddedChecks, embedded.Package.Location.File) // avoid infinite loop if embedded check contains ref error
+		excludedFiles = append(excludedFiles, e.Location.File)
+	}
+
+	compiler.Errors = lo.Filter(compiler.Errors, func(e *ast.Error, _ int) bool {
+		return !lo.Contains(excludedFiles, e.Location.File)
+	})
+}
+
+func (s *Scanner) findMatchedEmbeddedCheck(module *ast.Module) *ast.Module {
+	for _, policy := range s.embeddedChecks {
+		if policy.Package.Path.String() == module.Package.Path.String() ||
+			filepath.Base(policy.Package.Location.File) == filepath.Base(module.Package.Location.File) {
+			return policy
+		}
+	}
+	return nil
+}
+
 func (s *Scanner) prunePoliciesWithError(compiler *ast.Compiler) error {
 	if len(compiler.Errors) > s.regoErrorLimit {
 		s.debug.Log("Error(s) occurred while loading policies")
@@ -157,6 +229,7 @@ func (s *Scanner) compilePolicies(srcFS fs.FS, paths []string) error {
 
 	compiler.Compile(s.policies)
 	if compiler.Failed() {
+		s.fallbackChecks(compiler)
 		if err := s.prunePoliciesWithError(compiler); err != nil {
 			return err
 		}

pkg/iac/rego/load_test.go

@@ -8,6 +8,7 @@ import (
 	"testing"
 	"testing/fstest"
 
+	trivy_policies "github.com/aquasecurity/trivy-policies"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
@@ -19,6 +20,9 @@ import (
 //go:embed all:testdata/policies
 var testEmbedFS embed.FS
 
+//go:embed testdata/embedded
+var embeddedPoliciesFS embed.FS
+
 func Test_RegoScanning_WithSomeInvalidPolicies(t *testing.T) {
 	t.Run("allow no errors", func(t *testing.T) {
 		var debugBuf bytes.Buffer
@@ -95,3 +99,37 @@ deny {
 	})
 
 }
+
+func Test_FallbackToEmbedded(t *testing.T) {
+	scanner := rego.NewScanner(
+		types.SourceDockerfile,
+		options.ScannerWithRegoErrorLimits(0),
+	)
+	fsys := fstest.MapFS{
+		"policies/my-policy2.rego": &fstest.MapFile{
+			Data: []byte(`# METADATA
+# schemas:
+# - input: schema["fooschema"]
+
+package builtin.test
+
+deny {
+input.evil == "foo bar"
+}`),
+		},
+		"schemas/fooschema.json": &fstest.MapFile{
+			Data: []byte(`{
+				"$schema": "http://json-schema.org/draft-07/schema#",
+				"type": "object",
+				"properties": {
+					"foo": {
+						"type": "string"
+					}
+				}
+			}`),
+		},
+	}
+	trivy_policies.EmbeddedPolicyFileSystem = embeddedPoliciesFS
+	err := scanner.LoadPolicies(false, false, fsys, []string{"."}, nil)
+	assert.NoError(t, err)
+}

pkg/iac/rego/scanner.go

@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"io"
 	"io/fs"
+	"maps"
 	"strings"
 
 	"github.com/open-policy-agent/opa/ast"
@@ -41,6 +42,9 @@ type Scanner struct {
 	spec           string
 	inputSchema    interface{} // unmarshalled into this from a json schema document
 	sourceType     types.Source
+
+	embeddedLibs   map[string]*ast.Module
+	embeddedChecks map[string]*ast.Module
 }
 
 func (s *Scanner) SetUseEmbeddedLibraries(b bool) {
@@ -135,13 +139,12 @@ func NewScanner(source types.Source, opts ...options.ScannerOption) *Scanner {
 	s := &Scanner{
 		regoErrorLimit: ast.CompileErrorLimitDefault,
 		sourceType:     source,
-		ruleNamespaces: map[string]struct{}{
-			"builtin":   {},
-			"appshield": {},
-			"defsec":    {},
-		},
-		runtimeValues: addRuntimeValues(),
+		ruleNamespaces: make(map[string]struct{}),
+		runtimeValues:  addRuntimeValues(),
 	}
+
+	maps.Copy(s.ruleNamespaces, builtinNamespaces)
+
 	for _, opt := range opts {
 		opt(s)
 	}

pkg/iac/rego/testdata/embedded/my-policy1.rego

@@ -0,0 +1,9 @@
+# METADATA
+# schemas:
+# - input: schema["fooschema"]
+
+package builtin.test
+
+deny {
+    input.foo == "foo bar"
+}

pkg/scanner/local/scan.go

@@ -17,6 +17,7 @@ import (
 	"github.com/aquasecurity/trivy/pkg/fanal/analyzer"
 	"github.com/aquasecurity/trivy/pkg/fanal/applier"
 	ftypes "github.com/aquasecurity/trivy/pkg/fanal/types"
+	"github.com/aquasecurity/trivy/pkg/iac/rego"
 	"github.com/aquasecurity/trivy/pkg/licensing"
 	"github.com/aquasecurity/trivy/pkg/log"
 	"github.com/aquasecurity/trivy/pkg/scanner/langpkg"
@@ -383,7 +384,7 @@ func toDetectedMisconfiguration(res ftypes.MisconfResult, defaultSeverity dbType
 
 	// empty namespace implies a go rule from defsec, "builtin" refers to a built-in rego rule
 	// this ensures we don't generate bad links for custom policies
-	if res.Namespace == "" || strings.HasPrefix(res.Namespace, "builtin.") {
+	if res.Namespace == "" || rego.IsBuiltinNamespace(res.Namespace) {
 		primaryURL = fmt.Sprintf("https://avd.aquasec.com/misconfig/%s", strings.ToLower(res.ID))
 		res.References = append(res.References, primaryURL)
 	}