refactor(misconf): type-safe parser results in generic scanner (aquasecurity#9685)

Signed-off-by: nikpivkin <[email protected]>

Author: nikpivkin

pkg/iac/scanners/dockerfile/parser/parser.go

@@ -12,7 +12,7 @@ import (
 	"github.com/aquasecurity/trivy/pkg/iac/providers/dockerfile"
 )
 
-func Parse(_ context.Context, r io.Reader, path string) (any, error) {
+func Parse(_ context.Context, r io.Reader, path string) ([]*dockerfile.Dockerfile, error) {
 	parsed, err := parser.Parse(r)
 	if err != nil {
 		return nil, fmt.Errorf("dockerfile parse error: %w", err)
@@ -86,7 +86,7 @@ func Parse(_ context.Context, r io.Reader, path string) (any, error) {
 		parsedFile.Stages = append(parsedFile.Stages, stage)
 	}
 
-	return &parsedFile, nil
+	return []*dockerfile.Dockerfile{&parsedFile}, nil
 }
 
 func originalFromHeredoc(node *parser.Node) string {

pkg/iac/scanners/dockerfile/parser/parser_test.go

@@ -7,7 +7,6 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
-	"github.com/aquasecurity/trivy/pkg/iac/providers/dockerfile"
 	"github.com/aquasecurity/trivy/pkg/iac/scanners/dockerfile/parser"
 )
 
@@ -18,12 +17,11 @@ RUN make /app
 CMD python /app/app.py
 `
 
-	res, err := parser.Parse(t.Context(), strings.NewReader(input), "Dockerfile")
+	dfs, err := parser.Parse(t.Context(), strings.NewReader(input), "Dockerfile")
 	require.NoError(t, err)
+	require.Len(t, dfs, 1)
 
-	df, ok := res.(*dockerfile.Dockerfile)
-	require.True(t, ok)
-
+	df := dfs[0]
 	assert.Len(t, df.Stages, 1)
 
 	assert.Equal(t, "ubuntu:18.04", df.Stages[0].Name)
@@ -102,11 +100,12 @@ EOF`,
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			res, err := parser.Parse(t.Context(), strings.NewReader(tt.src), "Dockerfile")
+			dfs, err := parser.Parse(t.Context(), strings.NewReader(tt.src), "Dockerfile")
 			require.NoError(t, err)
+			require.Len(t, dfs, 1)
 
-			df, ok := res.(*dockerfile.Dockerfile)
-			require.True(t, ok)
+			df := dfs[0]
+			require.Len(t, df.Stages, 1)
 
 			cmd := df.Stages[0].Commands[0]
 

pkg/iac/scanners/dockerfile/scanner.go

@@ -1,12 +1,18 @@
 package dockerfile
 
 import (
+	"github.com/aquasecurity/trivy/pkg/iac/providers/dockerfile"
 	"github.com/aquasecurity/trivy/pkg/iac/scanners/dockerfile/parser"
 	"github.com/aquasecurity/trivy/pkg/iac/scanners/generic"
 	"github.com/aquasecurity/trivy/pkg/iac/scanners/options"
 	"github.com/aquasecurity/trivy/pkg/iac/types"
 )
 
-func NewScanner(opts ...options.ScannerOption) *generic.GenericScanner {
-	return generic.NewScanner("Dockerfile", types.SourceDockerfile, generic.ParseFunc(parser.Parse), opts...)
+func NewScanner(opts ...options.ScannerOption) *generic.GenericScanner[*dockerfile.Dockerfile] {
+	defaultOpts := []options.ScannerOption{
+		generic.WithSupportsInlineIgnore[*dockerfile.Dockerfile](true),
+	}
+	p := generic.ParseFunc[*dockerfile.Dockerfile](parser.Parse)
+	return generic.NewScanner("Dockerfile", types.SourceDockerfile, p, append(defaultOpts, opts...)...,
+	)
 }

pkg/iac/scanners/generic/scanner.go

@@ -22,41 +22,67 @@ import (
 	"github.com/aquasecurity/trivy/pkg/log"
 )
 
-func NewJsonScanner(opts ...options.ScannerOption) *GenericScanner {
-	return NewScanner("JSON", types.SourceJSON, ParseFunc(parseJson), opts...)
+func NewJsonScanner(opts ...options.ScannerOption) *GenericScanner[*identityMarshaler] {
+	return NewScanner("JSON", types.SourceJSON, ParseFunc[*identityMarshaler](parseJson), opts...)
 }
 
-func NewYamlScanner(opts ...options.ScannerOption) *GenericScanner {
-	return NewScanner("YAML", types.SourceYAML, ParseFunc(parseYaml), opts...)
+func NewYamlScanner(opts ...options.ScannerOption) *GenericScanner[*identityMarshaler] {
+	return NewScanner("YAML", types.SourceYAML, ParseFunc[*identityMarshaler](parseYaml), opts...)
 }
 
-func NewTomlScanner(opts ...options.ScannerOption) *GenericScanner {
-	return NewScanner("TOML", types.SourceTOML, ParseFunc(parseTOML), opts...)
+func NewTomlScanner(opts ...options.ScannerOption) *GenericScanner[*identityMarshaler] {
+	return NewScanner("TOML", types.SourceTOML, ParseFunc[*identityMarshaler](parseTOML), opts...)
 }
 
-type configParser interface {
-	Parse(ctx context.Context, r io.Reader, path string) (any, error)
+type RegoMarshaler interface {
+	ToRego() any
+}
+
+type identityMarshaler struct {
+	value any
+}
+
+func (r identityMarshaler) ToRego() any {
+	return r.value
+}
+
+type configParser[T RegoMarshaler] interface {
+	Parse(ctx context.Context, r io.Reader, path string) ([]T, error)
+}
+
+type ParseFunc[T RegoMarshaler] func(ctx context.Context, r io.Reader, path string) ([]T, error)
+
+func (f ParseFunc[T]) Parse(ctx context.Context, r io.Reader, path string) ([]T, error) {
+	return f(ctx, r, path)
+}
+
+func WithSupportsInlineIgnore[T RegoMarshaler](supports bool) options.ScannerOption {
+	return func(s options.ConfigurableScanner) {
+		if ss, ok := s.(*GenericScanner[T]); ok {
+			ss.supportsInlineIgnore = supports
+		}
+	}
 }
 
 // GenericScanner is a scanner that scans a file as is without processing it
-type GenericScanner struct {
+type GenericScanner[T RegoMarshaler] struct {
 	*rego.RegoScannerProvider
 	name    string
 	source  types.Source
 	logger  *log.Logger
 	options []options.ScannerOption
 
-	parser configParser
+	parser               configParser[T]
+	supportsInlineIgnore bool
 }
 
-type ParseFunc func(ctx context.Context, r io.Reader, path string) (any, error)
-
-func (f ParseFunc) Parse(ctx context.Context, r io.Reader, path string) (any, error) {
-	return f(ctx, r, path)
-}
-
-func NewScanner(name string, source types.Source, parser configParser, opts ...options.ScannerOption) *GenericScanner {
-	s := &GenericScanner{
+func NewScanner[T RegoMarshaler](
+	name string,
+	source types.Source,
+	parser configParser[T],
+	opts ...options.ScannerOption,
+) *GenericScanner[T] {
+	s := &GenericScanner[T]{
 		RegoScannerProvider: rego.NewRegoScannerProvider(opts...),
 		name:                name,
 		options:             opts,
@@ -72,41 +98,26 @@ func NewScanner(name string, source types.Source, parser configParser, opts ...o
 	return s
 }
 
-func (s *GenericScanner) Name() string {
+func (s *GenericScanner[T]) Name() string {
 	return s.name
 }
 
-func (s *GenericScanner) ScanFS(ctx context.Context, fsys fs.FS, dir string) (scan.Results, error) {
-	fileset, err := s.parseFS(ctx, fsys, dir)
+func (s *GenericScanner[T]) ScanFS(ctx context.Context, fsys fs.FS, rootDir string) (scan.Results, error) {
+	fileMap, err := s.parseFS(ctx, fsys, rootDir)
 	if err != nil {
 		return nil, err
 	}
 
-	if len(fileset) == 0 {
+	if len(fileMap) == 0 {
 		return nil, nil
 	}
 
 	var inputs []rego.Input
-	for path, val := range fileset {
-		switch v := val.(type) {
-		case interface{ ToRego() any }:
+	for filePath, parsedObjects := range fileMap {
+		for _, parsedObj := range parsedObjects {
 			inputs = append(inputs, rego.Input{
-				Path:     path,
-				Contents: v.ToRego(),
-				FS:       fsys,
-			})
-		case []any:
-			for _, file := range v {
-				inputs = append(inputs, rego.Input{
-					Path:     path,
-					Contents: file,
-					FS:       fsys,
-				})
-			}
-		default:
-			inputs = append(inputs, rego.Input{
-				Path:     path,
-				Contents: v,
+				Path:     filePath,
+				Contents: parsedObj.ToRego(),
 				FS:       fsys,
 			})
 		}
@@ -131,13 +142,9 @@ func (s *GenericScanner) ScanFS(ctx context.Context, fsys fs.FS, dir string) (sc
 	return results, nil
 }
 
-func (s *GenericScanner) supportsIgnoreRules() bool {
-	return s.source == types.SourceDockerfile
-}
-
-func (s *GenericScanner) parseFS(ctx context.Context, fsys fs.FS, path string) (map[string]any, error) {
-	files := make(map[string]any)
-	if err := fs.WalkDir(fsys, filepath.ToSlash(path), func(path string, entry fs.DirEntry, err error) error {
+func (s *GenericScanner[T]) parseFS(ctx context.Context, fsys fs.FS, rootDir string) (map[string][]T, error) {
+	parsedFiles := make(map[string][]T)
+	walkFn := func(filePath string, entry fs.DirEntry, err error) error {
 		select {
 		case <-ctx.Done():
 			return ctx.Err()
@@ -150,27 +157,29 @@ func (s *GenericScanner) parseFS(ctx context.Context, fsys fs.FS, path string) (
 			return nil
 		}
 
-		f, err := fsys.Open(filepath.ToSlash(path))
+		f, err := fsys.Open(filepath.ToSlash(filePath))
 		if err != nil {
 			return err
 		}
 		defer f.Close()
 
-		df, err := s.parser.Parse(ctx, f, path)
+		parsedObjects, err := s.parser.Parse(ctx, f, filePath)
 		if err != nil {
-			s.logger.Error("Failed to parse file", log.FilePath(path), log.Err(err))
+			s.logger.Error("Failed to parse file", log.FilePath(filePath), log.Err(err))
 			return nil
 		}
-		files[path] = df
+		parsedFiles[filePath] = parsedObjects
 		return nil
-	}); err != nil {
+	}
+
+	if err := fs.WalkDir(fsys, filepath.ToSlash(rootDir), walkFn); err != nil {
 		return nil, err
 	}
-	return files, nil
+	return parsedFiles, nil
 }
 
-func (s *GenericScanner) applyIgnoreRules(fsys fs.FS, results scan.Results) error {
-	if !s.supportsIgnoreRules() {
+func (s *GenericScanner[T]) applyIgnoreRules(fsys fs.FS, results scan.Results) error {
+	if !s.supportsInlineIgnore {
 		return nil
 	}
 
@@ -190,21 +199,21 @@ func (s *GenericScanner) applyIgnoreRules(fsys fs.FS, results scan.Results) erro
 	return nil
 }
 
-func parseJson(_ context.Context, r io.Reader, _ string) (any, error) {
+func parseJson(_ context.Context, r io.Reader, _ string) ([]*identityMarshaler, error) {
 	var target any
 	if err := json.NewDecoder(r).Decode(&target); err != nil {
 		return nil, err
 	}
-	return target, nil
+	return []*identityMarshaler{{value: target}}, nil
 }
 
-func parseYaml(_ context.Context, r io.Reader, _ string) (any, error) {
+func parseYaml(_ context.Context, r io.Reader, _ string) ([]*identityMarshaler, error) {
 	contents, err := io.ReadAll(r)
 	if err != nil {
 		return nil, err
 	}
 
-	var results []any
+	var documents []*identityMarshaler
 
 	marker := "\n---\n"
 	altMarker := "\r\n---\r\n"
@@ -217,16 +226,16 @@ func parseYaml(_ context.Context, r io.Reader, _ string) (any, error) {
 		if err := yaml.Unmarshal([]byte(partial), &target); err != nil {
 			return nil, err
 		}
-		results = append(results, target)
+		documents = append(documents, &identityMarshaler{target})
 	}
 
-	return results, nil
+	return documents, nil
 }
 
-func parseTOML(_ context.Context, r io.Reader, _ string) (any, error) {
+func parseTOML(_ context.Context, r io.Reader, _ string) ([]*identityMarshaler, error) {
 	var target any
 	if _, err := toml.NewDecoder(r).Decode(&target); err != nil {
 		return nil, err
 	}
-	return target, nil
+	return []*identityMarshaler{{value: target}}, nil
 }

pkg/iac/scanners/helm/scanner.go

@@ -140,7 +140,7 @@ func (s *Scanner) getScanResults(ctx context.Context, path string, target fs.FS)
 		for _, manifest := range manifests {
 			fileResults, err := rs.ScanInput(ctx, types.SourceKubernetes, rego.Input{
 				Path:     file.TemplateFilePath,
-				Contents: manifest,
+				Contents: manifest.ToRego(),
 				FS:       manifestFS,
 			})
 			if err != nil {

pkg/iac/scanners/kubernetes/parser/parser.go

@@ -8,7 +8,7 @@ import (
 	"strings"
 )
 
-func Parse(_ context.Context, r io.Reader, path string) ([]any, error) {
+func Parse(_ context.Context, r io.Reader, path string) ([]*Manifest, error) {
 	contents, err := io.ReadAll(r)
 	if err != nil {
 		return nil, err
@@ -23,10 +23,10 @@ func Parse(_ context.Context, r io.Reader, path string) ([]any, error) {
 		if err != nil {
 			return nil, err
 		}
-		return []any{manifest.ToRego()}, nil
+		return []*Manifest{manifest}, nil
 	}
 
-	var manifests []any
+	var manifests []*Manifest
 
 	re := regexp.MustCompile(`(?m:^---\r?\n)`)
 	offset := 0
@@ -35,10 +35,9 @@ func Parse(_ context.Context, r io.Reader, path string) ([]any, error) {
 		if err != nil {
 			return nil, err
 		}
-
 		if manifest.Content != nil {
 			manifest.Content.Offset = offset
-			manifests = append(manifests, manifest.ToRego())
+			manifests = append(manifests, manifest)
 		}
 
 		offset += countLines(partial)

pkg/iac/scanners/kubernetes/parser/parser_test.go

@@ -65,13 +65,9 @@ kind: Pod
 			require.NoError(t, err)
 			require.Len(t, manifests, len(tt.expectedOffsets))
 
-			for i, m := range manifests {
-				manifest := m.(map[string]any)
-				metadata, ok := manifest["__defsec_metadata"].(map[string]any)
-				require.True(t, ok)
-				offset, ok := metadata["offset"].(int)
-				require.True(t, ok)
-				require.Equal(t, tt.expectedOffsets[i], offset)
+			for i, manifest := range manifests {
+				require.NotNil(t, manifest.Content)
+				require.Equal(t, tt.expectedOffsets[i], manifest.Content.Offset)
 			}
 		})
 	}

pkg/iac/scanners/kubernetes/scanner.go

@@ -10,10 +10,11 @@ import (
 	"github.com/aquasecurity/trivy/pkg/iac/types"
 )
 
-func NewScanner(opts ...options.ScannerOption) *generic.GenericScanner {
-	return generic.NewScanner("Kubernetes", types.SourceKubernetes, generic.ParseFunc(parse), opts...)
+func NewScanner(opts ...options.ScannerOption) *generic.GenericScanner[*parser.Manifest] {
+	p := generic.ParseFunc[*parser.Manifest](parse)
+	return generic.NewScanner("Kubernetes", types.SourceKubernetes, p, opts...)
 }
 
-func parse(ctx context.Context, r io.Reader, path string) (any, error) {
+func parse(ctx context.Context, r io.Reader, path string) ([]*parser.Manifest, error) {
 	return parser.Parse(ctx, r, path)
 }