chore(detector/vuls2): check downloaded time before updating db (#2077)

* chore(detector/vuls2): check downloaded time before updating db

Co-authored-by: MaineK00n <[email protected]>

detector/vuls2/db.go

@@ -1,35 +1,23 @@
 package vuls2
 
 import (
-	"context"
-	"encoding/json"
-	"io"
+	"fmt"
 	"os"
 	"path/filepath"
-	"strconv"
 	"time"
 
+	"github.com/future-architect/vuls/config"
 	"github.com/future-architect/vuls/logging"
-	"github.com/klauspost/compress/zstd"
-	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/pkg/errors"
-	progressbar "github.com/schollz/progressbar/v3"
 	"golang.org/x/xerrors"
-	oras "oras.land/oras-go/v2"
-	"oras.land/oras-go/v2/content"
-	"oras.land/oras-go/v2/content/memory"
-	"oras.land/oras-go/v2/registry/remote"
 
 	db "github.com/MaineK00n/vuls2/pkg/db/common"
-)
-
-const (
-	dbMediaType = "application/vnd.vulsio.vuls.db.layer.v1+zstd"
+	"github.com/MaineK00n/vuls2/pkg/db/fetch"
 )
 
 var (
 	// DefaultGHCRRepository is GitHub Container Registry for vuls2 db
-	DefaultGHCRRepository = "ghcr.io/vulsio/vuls-nightly-db"
+	DefaultGHCRRepository = fmt.Sprintf("%s:%d", "ghcr.io/vulsio/vuls-nightly-db", db.SchemaVersion)
 
 	// DefaultPath is the path for vuls2 db file
 	DefaultPath = func() string {
@@ -38,160 +26,68 @@ var (
 	}()
 )
 
-type Config struct {
-	Repository string
-	Path       string
-	SkipUpdate bool
-	Quiet      bool
-}
-
-func (c Config) Refresh() error {
-	lastModified, fileExists, err := c.loadLastModified()
+func newDBConnection(vuls2Cnf config.Vuls2DictConf, noProgress bool) (db.DB, error) {
+	willDownload, err := shouldDownload(vuls2Cnf, time.Now())
 	if err != nil {
-		return xerrors.Errorf("Failed to load vuls2 db metadata. err: %w", err)
-	}
-
-	if fileExists && time.Now().Before(lastModified.Add(6*time.Hour)) {
-		return nil
+		return nil, xerrors.Errorf("Failed to check whether to download vuls2 db. err: %w", err)
 	}
 
-	if c.SkipUpdate {
-		if !fileExists {
-			return xerrors.New("Vuls2 db not found, cannot skip update")
+	if willDownload {
+		logging.Log.Infof("Downloading vuls2 db. repository: %s", vuls2Cnf.Repository)
+		if err := fetch.Fetch(fetch.WithRepository(vuls2Cnf.Repository), fetch.WithDBPath(vuls2Cnf.Path), fetch.WithNoProgress(noProgress)); err != nil {
+			return nil, xerrors.Errorf("Failed to fetch vuls2 db. err: %w", err)
 		}
-		return nil
 	}
 
-	logging.Log.Infof("Downloading vuls2 db. repository: %s", c.Repository)
-	if err := c.fetch(c.Repository); err != nil {
-		return xerrors.Errorf("Failed to fetch vuls2 db. repository: %s, err: %w", c.Repository, err)
-	}
-
-	return nil
-}
-
-func (c Config) New() (db.DB, error) {
-	vuls2Config := db.Config{
+	dbc, err := (&db.Config{
 		Type: "boltdb",
-		Path: c.Path,
-	}
-
-	dbc, err := vuls2Config.New()
+		Path: vuls2Cnf.Path,
+	}).New()
 	if err != nil {
-		return nil, xerrors.Errorf("Failed to new vuls2 db. err: %w", err)
+		return nil, xerrors.Errorf("Failed to new vuls2 db connection. err: %w", err)
 	}
 
 	return dbc, nil
 }
 
-func (c Config) fetch(repoPath string) error {
-	logging.Log.Infof("Fetch vuls.db from %s", repoPath)
-
-	ctx := context.TODO()
-
-	ms := memory.New()
-
-	repo, err := remote.NewRepository(repoPath)
-	if err != nil {
-		return xerrors.Errorf("Failed to create repository client. repository: %s, err: %w", repoPath, err)
-	}
-
-	manifestDescriptor, err := oras.Copy(ctx, repo, strconv.Itoa(db.SchemaVersion), ms, "", oras.DefaultCopyOptions)
-	if err != nil {
-		return xerrors.Errorf("Failed to copy. repository: %s, err: %w", repoPath, err)
-	}
-
-	r, err := ms.Fetch(ctx, manifestDescriptor)
-	if err != nil {
-		return xerrors.Errorf("Failed to fetch manifest. err: %w", err)
-	}
-	defer r.Close()
-
-	var manifest ocispec.Manifest
-	if err := json.NewDecoder(content.NewVerifyReader(r, manifestDescriptor)).Decode(&manifest); err != nil {
-		return xerrors.Errorf("Failed to decode manifest. err: %w", err)
-	}
-
-	l := func() *ocispec.Descriptor {
-		for _, l := range manifest.Layers {
-			if l.MediaType == dbMediaType {
-				return &l
+func shouldDownload(vuls2Cnf config.Vuls2DictConf, now time.Time) (bool, error) {
+	if _, err := os.Stat(vuls2Cnf.Path); err != nil {
+		if errors.Is(err, os.ErrNotExist) {
+			if vuls2Cnf.SkipUpdate {
+				return false, xerrors.Errorf("%s not found, cannot skip update", vuls2Cnf.Path)
 			}
+			return true, nil
 		}
-		return nil
-	}()
-	if l == nil {
-		return xerrors.Errorf("Failed to find digest and filename from layers. actual layers: %#v", manifest.Layers)
-	}
-
-	r, err = repo.Fetch(ctx, *l)
-	if err != nil {
-		return xerrors.Errorf("Failed to fetch content. err: %w", err)
+		return false, xerrors.Errorf("Failed to stat vuls2 db file. err: %w", err)
 	}
-	defer r.Close()
 
-	d, err := zstd.NewReader(content.NewVerifyReader(r, *l))
-	if err != nil {
-		return errors.Wrap(err, "new zstd reader")
-	}
-	defer d.Close()
-
-	if err := os.MkdirAll(filepath.Dir(c.Path), 0755); err != nil {
-		return errors.Wrapf(err, "mkdir %s", filepath.Dir(c.Path))
-	}
-
-	f, err := os.Create(c.Path)
-	if err != nil {
-		return xerrors.Errorf("Failed to create. file: %s, err:%w", c.Path, err)
+	if vuls2Cnf.SkipUpdate {
+		return false, nil
 	}
-	defer f.Close()
 
-	var pb *progressbar.ProgressBar
-	pb = progressbar.DefaultBytesSilent(-1)
-	if !c.Quiet {
-		pb = progressbar.DefaultBytes(-1, "downloading")
-	}
-	if _, err := d.WriteTo(io.MultiWriter(f, pb)); err != nil {
-		return xerrors.Errorf("Failed to write. filename: %s. err: %w", f.Name(), err)
-	}
-	_ = pb.Finish()
-
-	return nil
-}
-
-func (c Config) loadLastModified() (time.Time, bool, error) {
-	if _, err := os.Stat(c.Path); errors.Is(err, os.ErrNotExist) {
-		return time.Time{}, false, nil
-	}
-
-	conf := db.Config{
+	dbc, err := (&db.Config{
 		Type: "boltdb",
-		Path: c.Path,
-	}
-
-	dbc, err := conf.New()
+		Path: vuls2Cnf.Path,
+	}).New()
 	if err != nil {
-		return time.Time{}, false, xerrors.Errorf("Failed to new vuls2 db. path: %s, err: %w", c.Path, err)
+		return false, xerrors.Errorf("Failed to new vuls2 db connection. path: %s, err: %w", vuls2Cnf.Path, err)
 	}
 
 	if err := dbc.Open(); err != nil {
-		return time.Time{}, false, xerrors.Errorf("Failed to open vuls2 db. path: %s, err: %w", c.Path, err)
+		return false, xerrors.Errorf("Failed to open vuls2 db. path: %s, err: %w", vuls2Cnf.Path, err)
 	}
-	defer func() {
-		_ = dbc.Close()
-	}()
+	defer dbc.Close()
 
 	metadata, err := dbc.GetMetadata()
 	if err != nil {
-		return time.Time{}, false, xerrors.Errorf("Failed to get vuls2 db metadata. path: %s, err: %w", c.Path, err)
+		return false, xerrors.Errorf("Failed to get vuls2 db metadata. path: %s, err: %w", vuls2Cnf.Path, err)
 	}
 	if metadata == nil {
-		return time.Time{}, false, xerrors.Errorf("Unexpected Vuls2 db metadata. metadata: nil,. path: %s", c.Path)
+		return false, xerrors.Errorf("Unexpected Vuls2 db metadata. metadata: nil,. path: %s", vuls2Cnf.Path)
 	}
 
-	if metadata.SchemaVersion != db.SchemaVersion {
-		return time.Time{}, false, xerrors.Errorf("Unexpected schema version. expected: %d, actual: %d", db.SchemaVersion, metadata.SchemaVersion)
+	if metadata.Downloaded != nil && now.Before((*metadata.Downloaded).Add(1*time.Hour)) {
+		return false, nil
 	}
-
-	return metadata.LastModified, true, nil
+	return metadata.LastModified.Add(6 * time.Hour).Before(now), nil
 }

detector/vuls2/db_test.go

@@ -0,0 +1,149 @@
+package vuls2
+
+import (
+	"path/filepath"
+	"reflect"
+	"testing"
+	"time"
+
+	"golang.org/x/xerrors"
+
+	"github.com/MaineK00n/vuls2/pkg/db/common"
+	"github.com/MaineK00n/vuls2/pkg/db/common/types"
+	"github.com/future-architect/vuls/config"
+)
+
+func Test_shouldDownload(t *testing.T) {
+	type args struct {
+		vuls2Cnf config.Vuls2DictConf
+		now      time.Time
+	}
+	tests := []struct {
+		name     string
+		args     args
+		metadata *types.Metadata
+		want     bool
+		wantErr  bool
+	}{
+		{
+			name: "no db file",
+			args: args{
+				vuls2Cnf: config.Vuls2DictConf{},
+				now:      *parse("2024-01-02T00:00:00Z"),
+			},
+			want: true,
+		},
+		{
+			name: "no db file, but skip update",
+			args: args{
+				vuls2Cnf: config.Vuls2DictConf{
+					SkipUpdate: true,
+				},
+				now: *parse("2024-01-02T00:00:00Z"),
+			},
+			wantErr: true,
+		},
+		{
+			name: "just created",
+			args: args{
+				vuls2Cnf: config.Vuls2DictConf{},
+				now:      *parse("2024-01-02T00:00:00Z"),
+			},
+			metadata: &types.Metadata{
+				LastModified:  *parse("2024-01-02T00:00:00Z"),
+				Downloaded:    parse("2024-01-02T00:00:00Z"),
+				SchemaVersion: common.SchemaVersion,
+			},
+			want: false,
+		},
+		{
+			name: "8 hours old",
+			args: args{
+				vuls2Cnf: config.Vuls2DictConf{},
+				now:      *parse("2024-01-02T08:00:00Z"),
+			},
+			metadata: &types.Metadata{
+				LastModified:  *parse("2024-01-02T00:00:00Z"),
+				Downloaded:    parse("2024-01-02T00:00:00Z"),
+				SchemaVersion: common.SchemaVersion,
+			},
+			want: true,
+		},
+		{
+			name: "8 hours old, but skip update",
+			args: args{
+				vuls2Cnf: config.Vuls2DictConf{
+					SkipUpdate: true,
+				},
+				now: *parse("2024-01-02T08:00:00Z"),
+			},
+			metadata: &types.Metadata{
+				LastModified:  *parse("2024-01-02T00:00:00Z"),
+				Downloaded:    parse("2024-01-02T00:00:00Z"),
+				SchemaVersion: common.SchemaVersion,
+			},
+			want: false,
+		},
+		{
+			name: "8 hours old, but download recently",
+			args: args{
+				vuls2Cnf: config.Vuls2DictConf{},
+				now:      *parse("2024-01-02T08:00:00Z"),
+			},
+			metadata: &types.Metadata{
+				LastModified:  *parse("2024-01-02T00:00:00Z"),
+				Downloaded:    parse("2024-01-02T07:30:00Z"),
+				SchemaVersion: common.SchemaVersion,
+			},
+			want: false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			d := t.TempDir()
+			tt.args.vuls2Cnf.Path = filepath.Join(d, "vuls.db")
+			if tt.metadata != nil {
+				if err := putMetadata(*tt.metadata, tt.args.vuls2Cnf.Path); err != nil {
+					t.Errorf("putMetadata err = %v", err)
+					return
+				}
+			}
+			got, err := shouldDownload(tt.args.vuls2Cnf, tt.args.now)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("shouldDownload() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("shouldDownload() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+
+}
+
+func putMetadata(metadata types.Metadata, path string) error {
+	c := common.Config{
+		Type: "boltdb",
+		Path: path,
+	}
+	dbc, err := c.New()
+	if err != nil {
+		return xerrors.Errorf("c.New(). err: %w", err)
+	}
+	if err := dbc.Open(); err != nil {
+		return xerrors.Errorf("dbc.Open(). err: %w", err)
+	}
+	defer dbc.Close()
+	if err := dbc.Initialize(); err != nil {
+		return xerrors.Errorf("dbc.Initialize(). err: %w", err)
+	}
+	if err := dbc.PutMetadata(metadata); err != nil {
+		return xerrors.Errorf("dbc.PutMetadata(). err: %w", err)
+	}
+	return nil
+}
+
+func parse(date string) *time.Time {
+	t, _ := time.Parse(time.RFC3339, date)
+	return &t
+}