refactor: add hook interface for extended functionality (#8585)

Author: knqyf263

integration/module_test.go

@@ -3,12 +3,12 @@
 package integration
 
 import (
-	"github.com/aquasecurity/trivy/pkg/types"
 	"path/filepath"
 	"testing"
 
+	"github.com/aquasecurity/trivy/pkg/extension"
 	"github.com/aquasecurity/trivy/pkg/fanal/analyzer"
-	"github.com/aquasecurity/trivy/pkg/scan/post"
+	"github.com/aquasecurity/trivy/pkg/types"
 )
 
 func TestModule(t *testing.T) {
@@ -52,7 +52,7 @@ func TestModule(t *testing.T) {
 
 			t.Cleanup(func() {
 				analyzer.DeregisterAnalyzer("spring4shell")
-				post.DeregisterPostScanner("spring4shell")
+				extension.DeregisterHook("spring4shell")
 			})
 
 			// Run Trivy

internal/hooktest/hook.go

@@ -0,0 +1,96 @@
+package hooktest
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/aquasecurity/trivy/pkg/extension"
+	"github.com/aquasecurity/trivy/pkg/flag"
+	"github.com/aquasecurity/trivy/pkg/types"
+)
+
+type testHook struct{}
+
+func (*testHook) Name() string {
+	return "test"
+}
+
+func (*testHook) Version() int {
+	return 1
+}
+
+// RunHook implementation
+func (*testHook) PreRun(ctx context.Context, opts flag.Options) error {
+	if opts.GlobalOptions.ConfigFile == "bad-config" {
+		return errors.New("bad pre-run")
+	}
+	return nil
+}
+
+func (*testHook) PostRun(ctx context.Context, opts flag.Options) error {
+	if opts.GlobalOptions.ConfigFile == "bad-config" {
+		return errors.New("bad post-run")
+	}
+	return nil
+}
+
+// ScanHook implementation
+func (*testHook) PreScan(ctx context.Context, target *types.ScanTarget, options types.ScanOptions) error {
+	if target.Name == "bad-pre" {
+		return errors.New("bad pre-scan")
+	}
+	target.Name += " (pre-scan)"
+	return nil
+}
+
+func (*testHook) PostScan(ctx context.Context, results types.Results) (types.Results, error) {
+	for i, r := range results {
+		if r.Target == "bad" {
+			return nil, errors.New("bad")
+		}
+		for j := range r.Vulnerabilities {
+			results[i].Vulnerabilities[j].References = []string{
+				"https://example.com/post-scan",
+			}
+		}
+	}
+	return results, nil
+}
+
+// ReportHook implementation
+func (*testHook) PreReport(ctx context.Context, report *types.Report, opts flag.Options) error {
+	if report.ArtifactName == "bad-report" {
+		return errors.New("bad pre-report")
+	}
+
+	// Modify the report
+	for i := range report.Results {
+		for j := range report.Results[i].Vulnerabilities {
+			report.Results[i].Vulnerabilities[j].Title = "Modified by pre-report hook"
+		}
+	}
+	return nil
+}
+
+func (*testHook) PostReport(ctx context.Context, report *types.Report, opts flag.Options) error {
+	if report.ArtifactName == "bad-report" {
+		return errors.New("bad post-report")
+	}
+
+	// Modify the report
+	for i := range report.Results {
+		for j := range report.Results[i].Vulnerabilities {
+			report.Results[i].Vulnerabilities[j].Description = "Modified by post-report hook"
+		}
+	}
+	return nil
+}
+
+func Init(t *testing.T) {
+	h := &testHook{}
+	extension.RegisterHook(h)
+	t.Cleanup(func() {
+		extension.DeregisterHook(h.Name())
+	})
+}

pkg/commands/artifact/run.go

@@ -15,6 +15,7 @@ import (
 	"github.com/aquasecurity/trivy/pkg/cache"
 	"github.com/aquasecurity/trivy/pkg/commands/operation"
 	"github.com/aquasecurity/trivy/pkg/db"
+	"github.com/aquasecurity/trivy/pkg/extension"
 	"github.com/aquasecurity/trivy/pkg/fanal/analyzer"
 	"github.com/aquasecurity/trivy/pkg/fanal/artifact"
 	ftypes "github.com/aquasecurity/trivy/pkg/fanal/types"
@@ -277,7 +278,6 @@ func (r *runner) Report(ctx context.Context, opts flag.Options, report types.Rep
 	if err := pkgReport.Write(ctx, report, opts); err != nil {
 		return xerrors.Errorf("unable to write results: %w", err)
 	}
-
 	return nil
 }
 
@@ -375,12 +375,32 @@ func Run(ctx context.Context, opts flag.Options, targetKind TargetKind) (err err
 		return v.SafeWriteConfigAs("trivy-default.yaml")
 	}
 
+	// Call pre-run hooks
+	if err := extension.PreRun(ctx, opts); err != nil {
+		return xerrors.Errorf("pre run error: %w", err)
+	}
+
+	// Run the application
+	report, err := run(ctx, opts, targetKind)
+	if err != nil {
+		return xerrors.Errorf("run error: %w", err)
+	}
+
+	// Call post-run hooks
+	if err := extension.PostRun(ctx, opts); err != nil {
+		return xerrors.Errorf("post run error: %w", err)
+	}
+
+	return operation.Exit(opts, report.Results.Failed(), report.Metadata)
+}
+
+func run(ctx context.Context, opts flag.Options, targetKind TargetKind) (types.Report, error) {
 	r, err := NewRunner(ctx, opts)
 	if err != nil {
 		if errors.Is(err, SkipScan) {
-			return nil
+			return types.Report{}, nil
 		}
-		return xerrors.Errorf("init error: %w", err)
+		return types.Report{}, xerrors.Errorf("init error: %w", err)
 	}
 	defer r.Close(ctx)
 
@@ -395,24 +415,27 @@ func Run(ctx context.Context, opts flag.Options, targetKind TargetKind) (err err
 
 	scanFunction, exists := scans[targetKind]
 	if !exists {
-		return xerrors.Errorf("unknown target kind: %s", targetKind)
+		return types.Report{}, xerrors.Errorf("unknown target kind: %s", targetKind)
 	}
 
+	// 1. Scan the artifact
 	report, err := scanFunction(ctx, opts)
 	if err != nil {
-		return xerrors.Errorf("%s scan error: %w", targetKind, err)
+		return types.Report{}, xerrors.Errorf("%s scan error: %w", targetKind, err)
 	}
 
+	// 2. Filter the results
 	report, err = r.Filter(ctx, opts, report)
 	if err != nil {
-		return xerrors.Errorf("filter error: %w", err)
+		return types.Report{}, xerrors.Errorf("filter error: %w", err)
 	}
 
+	// 3. Report the results
 	if err = r.Report(ctx, opts, report); err != nil {
-		return xerrors.Errorf("report error: %w", err)
+		return types.Report{}, xerrors.Errorf("report error: %w", err)
 	}
 
-	return operation.Exit(opts, report.Results.Failed(), report.Metadata)
+	return report, nil
 }
 
 func disabledAnalyzers(opts flag.Options) []analyzer.Type {

pkg/extension/hook.go

@@ -0,0 +1,162 @@
+package extension
+
+import (
+	"context"
+	"sort"
+
+	"github.com/samber/lo"
+	"golang.org/x/xerrors"
+
+	"github.com/aquasecurity/trivy/pkg/flag"
+	"github.com/aquasecurity/trivy/pkg/types"
+)
+
+var hooks = make(map[string]Hook)
+
+func RegisterHook(s Hook) {
+	// Avoid duplication
+	hooks[s.Name()] = s
+}
+
+func DeregisterHook(name string) {
+	delete(hooks, name)
+}
+
+// Hook is an interface that defines the methods for a hook.
+type Hook interface {
+	// Name returns the name of the extension.
+	Name() string
+}
+
+// RunHook is a extension that is called before and after all the processes.
+type RunHook interface {
+	Hook
+
+	// PreRun is called before all the processes.
+	PreRun(ctx context.Context, opts flag.Options) error
+
+	// PostRun is called after all the processes.
+	PostRun(ctx context.Context, opts flag.Options) error
+}
+
+// ScanHook is a extension that is called before and after the scan.
+type ScanHook interface {
+	Hook
+
+	// PreScan is called before the scan. It can modify the scan target.
+	// It may be called on the server side in client/server mode.
+	PreScan(ctx context.Context, target *types.ScanTarget, opts types.ScanOptions) error
+
+	// PostScan is called after the scan. It can modify the results.
+	// It may be called on the server side in client/server mode.
+	// NOTE: Wasm modules cannot directly modify the passed results,
+	//       so it returns a copy of the results.
+	PostScan(ctx context.Context, results types.Results) (types.Results, error)
+}
+
+// ReportHook is a extension that is called before and after the report is written.
+type ReportHook interface {
+	Hook
+
+	// PreReport is called before the report is written.
+	// It can modify the report. It is called on the client side.
+	PreReport(ctx context.Context, report *types.Report, opts flag.Options) error
+
+	// PostReport is called after the report is written.
+	// It can modify the report. It is called on the client side.
+	PostReport(ctx context.Context, report *types.Report, opts flag.Options) error
+}
+
+func PreRun(ctx context.Context, opts flag.Options) error {
+	for _, e := range Hooks() {
+		h, ok := e.(RunHook)
+		if !ok {
+			continue
+		}
+		if err := h.PreRun(ctx, opts); err != nil {
+			return xerrors.Errorf("%s pre run error: %w", e.Name(), err)
+		}
+	}
+	return nil
+}
+
+// PostRun is a hook that is called after all the processes.
+func PostRun(ctx context.Context, opts flag.Options) error {
+	for _, e := range Hooks() {
+		h, ok := e.(RunHook)
+		if !ok {
+			continue
+		}
+		if err := h.PostRun(ctx, opts); err != nil {
+			return xerrors.Errorf("%s post run error: %w", e.Name(), err)
+		}
+	}
+	return nil
+}
+
+// PreScan is a hook that is called before the scan.
+func PreScan(ctx context.Context, target *types.ScanTarget, options types.ScanOptions) error {
+	for _, e := range Hooks() {
+		h, ok := e.(ScanHook)
+		if !ok {
+			continue
+		}
+		if err := h.PreScan(ctx, target, options); err != nil {
+			return xerrors.Errorf("%s pre scan error: %w", e.Name(), err)
+		}
+	}
+	return nil
+}
+
+// PostScan is a hook that is called after the scan.
+func PostScan(ctx context.Context, results types.Results) (types.Results, error) {
+	var err error
+	for _, e := range Hooks() {
+		h, ok := e.(ScanHook)
+		if !ok {
+			continue
+		}
+		results, err = h.PostScan(ctx, results)
+		if err != nil {
+			return nil, xerrors.Errorf("%s post scan error: %w", e.Name(), err)
+		}
+	}
+	return results, nil
+}
+
+// PreReport is a hook that is called before the report is written.
+func PreReport(ctx context.Context, report *types.Report, opts flag.Options) error {
+	for _, e := range Hooks() {
+		h, ok := e.(ReportHook)
+		if !ok {
+			continue
+		}
+		if err := h.PreReport(ctx, report, opts); err != nil {
+			return xerrors.Errorf("%s pre report error: %w", e.Name(), err)
+		}
+	}
+	return nil
+}
+
+// PostReport is a hook that is called after the report is written.
+func PostReport(ctx context.Context, report *types.Report, opts flag.Options) error {
+	for _, e := range Hooks() {
+		h, ok := e.(ReportHook)
+		if !ok {
+			continue
+		}
+		if err := h.PostReport(ctx, report, opts); err != nil {
+			return xerrors.Errorf("%s post report error: %w", e.Name(), err)
+		}
+	}
+	return nil
+}
+
+// Hooks returns the list of hooks.
+func Hooks() []Hook {
+	hooks := lo.Values(hooks)
+	sort.Slice(hooks, func(i, j int) bool {
+		return hooks[i].Name() < hooks[j].Name()
+	})
+	return hooks
+}