From f0df725c5a217bc7324a14d0cac6e2a91a97c022 Mon Sep 17 00:00:00 2001
From: Teppei Fukuda <knqyf263@gmail.com>
Date: Mon, 10 Apr 2023 15:37:13 +0300
Subject: [PATCH] fix: lock downloading policies and database (#4017)

---
 pkg/commands/artifact/run.go        |  6 +++---
 pkg/commands/operation/operation.go | 12 ++++++++++--
 pkg/commands/server/run.go          |  2 +-
 3 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/pkg/commands/artifact/run.go b/pkg/commands/artifact/run.go
index 24947788e0..b408f97603 100644
--- a/pkg/commands/artifact/run.go
+++ b/pkg/commands/artifact/run.go
@@ -126,7 +126,7 @@ func NewRunner(ctx context.Context, cliOptions flag.Options, opts ...runnerOptio
 	}
 
 	// Update the vulnerability database if needed.
-	if err := r.initDB(cliOptions); err != nil {
+	if err := r.initDB(ctx, cliOptions); err != nil {
 		return nil, xerrors.Errorf("DB error: %w", err)
 	}
 
@@ -302,7 +302,7 @@ func (r *runner) Report(opts flag.Options, report types.Report) error {
 	return nil
 }
 
-func (r *runner) initDB(opts flag.Options) error {
+func (r *runner) initDB(ctx context.Context, opts flag.Options) error {
 	if err := r.initJavaDB(opts); err != nil {
 		return err
 	}
@@ -314,7 +314,7 @@ func (r *runner) initDB(opts flag.Options) error {
 
 	// download the database file
 	noProgress := opts.Quiet || opts.NoProgress
-	if err := operation.DownloadDB(opts.AppVersion, opts.CacheDir, opts.DBRepository, noProgress, opts.SkipDBUpdate, opts.Remote()); err != nil {
+	if err := operation.DownloadDB(ctx, opts.AppVersion, opts.CacheDir, opts.DBRepository, noProgress, opts.SkipDBUpdate, opts.Remote()); err != nil {
 		return err
 	}
 
diff --git a/pkg/commands/operation/operation.go b/pkg/commands/operation/operation.go
index 1982e9af55..eb512ef2b0 100644
--- a/pkg/commands/operation/operation.go
+++ b/pkg/commands/operation/operation.go
@@ -6,6 +6,7 @@ import (
 	"crypto/x509"
 	"os"
 	"strings"
+	"sync"
 
 	"github.com/go-redis/redis/v8"
 	"github.com/google/wire"
@@ -22,6 +23,8 @@ import (
 	"github.com/aquasecurity/trivy/pkg/utils/fsutils"
 )
 
+var mu sync.Mutex
+
 // SuperSet binds cache dependencies
 var SuperSet = wire.NewSet(
 	cache.NewFSCache,
@@ -106,9 +109,11 @@ func (c Cache) ClearArtifacts() error {
 }
 
 // DownloadDB downloads the DB
-func DownloadDB(appVersion, cacheDir, dbRepository string, quiet, skipUpdate bool, opt types.RemoteOptions) error {
+func DownloadDB(ctx context.Context, appVersion, cacheDir, dbRepository string, quiet, skipUpdate bool, opt types.RemoteOptions) error {
+	mu.Lock()
+	defer mu.Unlock()
+
 	client := db.NewClient(cacheDir, quiet, db.WithDBRepository(dbRepository))
-	ctx := context.Background()
 	needsUpdate, err := client.NeedsUpdate(appVersion, skipUpdate)
 	if err != nil {
 		return xerrors.Errorf("database error: %w", err)
@@ -143,6 +148,9 @@ func showDBInfo(cacheDir string) error {
 
 // InitBuiltinPolicies downloads the built-in policies and loads them
 func InitBuiltinPolicies(ctx context.Context, cacheDir string, quiet, skipUpdate bool) ([]string, error) {
+	mu.Lock()
+	defer mu.Unlock()
+
 	client, err := policy.NewClient(cacheDir, quiet)
 	if err != nil {
 		return nil, xerrors.Errorf("policy client error: %w", err)
diff --git a/pkg/commands/server/run.go b/pkg/commands/server/run.go
index 16e92b4177..5d13584ac5 100644
--- a/pkg/commands/server/run.go
+++ b/pkg/commands/server/run.go
@@ -34,7 +34,7 @@ func Run(ctx context.Context, opts flag.Options) (err error) {
 	}
 
 	// download the database file
-	if err = operation.DownloadDB(opts.AppVersion, opts.CacheDir, opts.DBRepository,
+	if err = operation.DownloadDB(ctx, opts.AppVersion, opts.CacheDir, opts.DBRepository,
 		true, opts.SkipDBUpdate, opts.Remote()); err != nil {
 		return err
 	}