Allow user specified cache directory (aquasecurity#12)

* Added cache-dir flag The flag cache-dir allows users to specify the cache directory the database is stored in. Signed-off-by: Ken Herner <[email protected]> * Moved dbDir definition into functions Need to move dbDir definition into the calling functions as the cache directory may have been specified. With dbDir defined at package level, it would always be instantiated to the default value and would never be updated to the user specifyed value. Signed-off-by: Ken Herner <[email protected]>
GuaoGuao · May 16, 2019 · 2d512c5 · 2d512c5
1 parent 936297a
commit 2d512c5
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 4 deletions.
diff --git a/cmd/trivy/main.go b/cmd/trivy/main.go
@@ -93,6 +93,10 @@ OPTIONS:
 			Name:  "debug, d",
 			Usage: "debug mode",
 		},
+		cli.StringFlag{
+			Name: "cache-dir",
+			Usage: "cache directory",
+		},
 	}
 
 	app.Action = func(c *cli.Context) error {

diff --git a/pkg/db/db.go b/pkg/db/db.go
@@ -16,10 +16,10 @@ import (
 
 var (
 	db    *bolt.DB
-	dbDir = filepath.Join(utils.CacheDir(), "db")
 )
 
 func Init() (err error) {
+	dbDir := filepath.Join(utils.CacheDir(), "db")
 	if err = os.MkdirAll(dbDir, 0700); err != nil {
 		return xerrors.Errorf("failed to mkdir: %w", err)
 	}
@@ -34,6 +34,7 @@ func Init() (err error) {
 }
 
 func Reset() error {
+	dbDir := filepath.Join(utils.CacheDir(), "db")
 	if err := os.RemoveAll(dbDir); err != nil {
 		return xerrors.Errorf("failed to reset DB: %w", err)
 	}

diff --git a/pkg/run.go b/pkg/run.go
@@ -29,6 +29,12 @@ func Run(c *cli.Context) (err error) {
 	if err = log.InitLogger(debug); err != nil {
 		l.Fatal(err)
 	}
+
+	cacheDir := c.String("cache-dir")
+	if cacheDir != "" {
+		utils.SetCacheDir(cacheDir)
+	}
+
 	log.Logger.Debugf("cache dir:  %s", utils.CacheDir())
 
 	if c.Bool("reset") {

diff --git a/pkg/utils/utils.go b/pkg/utils/utils.go
@@ -12,15 +12,24 @@ import (
 	"golang.org/x/xerrors"
 )
 
+var cacheDir string
+
 func CacheDir() string {
-	cacheDir, err := os.UserCacheDir()
-	if err != nil {
-		cacheDir = os.TempDir()
+	if cacheDir == "" {
+		var err error
+		cacheDir, err = os.UserCacheDir()
+		if err != nil {
+			cacheDir = os.TempDir()
+		}
 	}
 	dir := filepath.Join(cacheDir, "trivy")
 	return dir
 }
 
+func SetCacheDir(cd string) {
+	cacheDir = cd
+}
+
 func FileWalk(root string, targetFiles map[string]struct{}, walkFn func(r io.Reader, path string) error) error {
 	err := filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
 		if info.IsDir() {