added: display last db update whenever trivy server is started

go.mod

@@ -26,6 +26,7 @@ require (
 	github.com/kylelemons/godebug v1.1.0
 	github.com/olekukonko/tablewriter v0.0.2-0.20190607075207-195002e6e56a
 	github.com/open-policy-agent/opa v0.21.1
+	github.com/prometheus/client_golang v1.0.0
 	github.com/spf13/afero v1.2.2
 	github.com/stretchr/testify v1.6.1
 	github.com/testcontainers/testcontainers-go v0.3.1

go.sum

@@ -117,6 +117,7 @@ github.com/aws/aws-sdk-go v1.16.26/go.mod h1:KmX6BPdI08NWTb3/sm4ZGu5ShLoqVDhKgpi
 github.com/aws/aws-sdk-go v1.27.1 h1:MXnqY6SlWySaZAqNnXThOvjRFdiiOuKtC6i7baFdNdU=
 github.com/aws/aws-sdk-go v1.27.1/go.mod h1:KmX6BPdI08NWTb3/sm4ZGu5ShLoqVDhKgpiN924inxo=
 github.com/beorn7/perks v0.0.0-20180321164747-3a771d992973/go.mod h1:Dwedo/Wpr24TaqPxmxbtue+5NUziq4I4S80YR8gNf3Q=
+github.com/beorn7/perks v1.0.0 h1:HWo1m869IqiPhD389kmkxeTalrjNbbJTC8LXupb+sl0=
 github.com/beorn7/perks v1.0.0/go.mod h1:KWe93zE9D1o94FZ5RNwFwVgaQK1VOXiVxmqh+CedLV8=
 github.com/bgentry/speakeasy v0.1.0/go.mod h1:+zsyZBPWlz7T6j88CTgSN5bM796AkVf0kBD4zp0CCIs=
 github.com/blang/semver v3.5.0+incompatible/go.mod h1:kRBLl5iJ+tD4TcOOxsy/0fnwebNt5EWlYSAyrTnjyyk=
@@ -449,6 +450,7 @@ github.com/mattn/go-runewidth v0.0.4/go.mod h1:LwmH8dsx7+W8Uxz3IHJYH5QSwggIsqBzp
 github.com/mattn/go-runewidth v0.0.6/go.mod h1:H031xJmbD/WCDINGzjvQ9THkh0rPKHF+m2gUSrubnMI=
 github.com/mattn/go-runewidth v0.0.9 h1:Lm995f3rfxdpd6TSmuVCHVb/QhupuXlYr8sCI/QdE+0=
 github.com/mattn/go-runewidth v0.0.9/go.mod h1:H031xJmbD/WCDINGzjvQ9THkh0rPKHF+m2gUSrubnMI=
+github.com/matttproud/golang_protobuf_extensions v1.0.1 h1:4hp9jkHxhMHkqkrB3Ix0jegS5sx/RkqARlsWZ6pIwiU=
 github.com/matttproud/golang_protobuf_extensions v1.0.1/go.mod h1:D8He9yQNgCq6Z5Ld7szi9bcBfOoFv/3dc6xSMkL2PC0=
 github.com/maxbrunsfeld/counterfeiter/v6 v6.2.2/go.mod h1:eD9eIE7cdwcMi9rYluz88Jz2VyhSmden33/aXg4oVIY=
 github.com/mitchellh/go-homedir v1.1.0 h1:lukF9ziXFxDFPkA1vsr5zpc1XuPDn/wFntq5mG+4E0Y=
@@ -513,13 +515,18 @@ github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZN
 github.com/pquerna/cachecontrol v0.0.0-20171018203845-0dec1b30a021/go.mod h1:prYjPmNq4d1NPVmpShWobRqXY3q7Vp+80DqgxxUrUIA=
 github.com/prometheus/client_golang v0.0.0-20181025174421-f30f42803563/go.mod h1:7SWBe2y4D6OKWSNQJUaRYU/AaXPKyh/dDVn+NZz0KFw=
 github.com/prometheus/client_golang v0.9.1/go.mod h1:7SWBe2y4D6OKWSNQJUaRYU/AaXPKyh/dDVn+NZz0KFw=
+github.com/prometheus/client_golang v1.0.0 h1:vrDKnkGzuGvhNAL56c7DBz29ZL+KxnoR0x7enabFceM=
 github.com/prometheus/client_golang v1.0.0/go.mod h1:db9x61etRT2tGnBNRi70OPL5FsnadC4Ky3P0J6CfImo=
+github.com/prometheus/client_golang v1.9.0 h1:Rrch9mh17XcxvEu9D9DEpb4isxjGBtcevQjKvxPRQIU=
 github.com/prometheus/client_model v0.0.0-20180712105110-5c3871d89910/go.mod h1:MbSGuTsp3dbXC40dX6PRTWyKYBIrTGTE9sqQNg2J8bo=
 github.com/prometheus/client_model v0.0.0-20190129233127-fd36f4220a90/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=
+github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4 h1:gQz4mCbXsO+nc9n1hCxHcGA3Zx3Eo+UHZoInFGUIXNM=
 github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=
 github.com/prometheus/common v0.0.0-20181020173914-7e9e6cabbd39/go.mod h1:daVV7qP5qjZbuso7PdcryaAu0sAZbrN9i7WWcTMWvro=
+github.com/prometheus/common v0.4.1 h1:K0MGApIoQvMw27RTdJkPbr3JZ7DNbtxQNyi5STVM6Kw=
 github.com/prometheus/common v0.4.1/go.mod h1:TNfzLD0ON7rHzMJeJkieUDPYmFC7Snx/y86RQel1bk4=
 github.com/prometheus/procfs v0.0.0-20181005140218-185b4288413d/go.mod h1:c3At6R/oaqEKCNdg8wHV1ftS6bRYblBhIjjI8uT2IGk=
+github.com/prometheus/procfs v0.0.2 h1:6LJUbpNm42llc4HRCuvApCSWB/WfhuNo9K98Q9sNGfs=
 github.com/prometheus/procfs v0.0.2/go.mod h1:TjEm7ze935MbeOT/UhFTIMYKhuLP4wbCsTZCD3I8kEA=
 github.com/rcrowley/go-metrics v0.0.0-20181016184325-3113b8401b8a h1:9ZKAASQSHhDYGoxY8uLVpewe1GDZ2vu2Tr/vTdVAkFQ=
 github.com/rcrowley/go-metrics v0.0.0-20181016184325-3113b8401b8a/go.mod h1:bCqnVzQkZxMG4s8nGwiZ5l3QUCyqpo9Y+/ZMZ9VjZe4=

internal/server/extendedconfig/extended_config.go

@@ -0,0 +1,34 @@
+package extendedconfig
+
+import (
+	"github.com/prometheus/client_golang/prometheus"
+
+	"github.com/aquasecurity/trivy/internal/server/config"
+)
+
+// ExtendedConfig represents prometheus-related configuration alongside the provided config.Config
+type ExtendedConfig struct {
+	Config          config.Config
+	MetricsRegistry *prometheus.Registry
+	GaugeMetric     *prometheus.GaugeVec
+}
+
+// New bootstraps an ExtendedConfig object out of a config.Config object for further operation
+func New(c config.Config) ExtendedConfig {
+	return ExtendedConfig{
+		Config: c,
+	}
+}
+
+// Init populates the relevant prometheus-related configuration in an ExtendedConfig object
+func (ec *ExtendedConfig) Init() {
+	ec.MetricsRegistry = prometheus.NewRegistry()
+	ec.GaugeMetric = prometheus.NewGaugeVec(
+		prometheus.GaugeOpts{
+			Name: "trivy",
+			Help: "Gauge Metrics associated with trivy - Last DB Update, Last DB Update Attempt ...",
+		},
+		[]string{"action"},
+	)
+	ec.MetricsRegistry.MustRegister(ec.GaugeMetric)
+}

internal/server/run.go

@@ -7,6 +7,7 @@ import (
 	"github.com/aquasecurity/trivy-db/pkg/db"
 	"github.com/aquasecurity/trivy/internal/operation"
 	"github.com/aquasecurity/trivy/internal/server/config"
+	"github.com/aquasecurity/trivy/internal/server/extendedconfig"
 	"github.com/aquasecurity/trivy/pkg/log"
 	"github.com/aquasecurity/trivy/pkg/rpc/server"
 	"github.com/aquasecurity/trivy/pkg/utils"
@@ -52,6 +53,8 @@ func run(c config.Config) (err error) {
 	if err = db.Init(c.CacheDir); err != nil {
 		return xerrors.Errorf("error in vulnerability DB initialize: %w", err)
 	}
+	ec := extendedconfig.New(c)
+	ec.Init()
 
-	return server.ListenAndServe(c, cache)
+	return server.ListenAndServe(ec, cache)
 }

pkg/rpc/server/listen.go

@@ -8,13 +8,17 @@ import (
 	"sync"
 	"time"
 
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promhttp"
+	"github.com/spf13/afero"
+
 	"github.com/google/wire"
 	"github.com/twitchtv/twirp"
 	"golang.org/x/xerrors"
 
 	"github.com/aquasecurity/fanal/cache"
 	"github.com/aquasecurity/trivy-db/pkg/db"
-	"github.com/aquasecurity/trivy/internal/server/config"
+	"github.com/aquasecurity/trivy/internal/server/extendedconfig"
 	dbFile "github.com/aquasecurity/trivy/pkg/db"
 	"github.com/aquasecurity/trivy/pkg/log"
 	"github.com/aquasecurity/trivy/pkg/utils"
@@ -31,28 +35,32 @@ var DBWorkerSuperSet = wire.NewSet(
 )
 
 // ListenAndServe starts Trivy server
-func ListenAndServe(c config.Config, serverCache cache.Cache) error {
+func ListenAndServe(ec extendedconfig.ExtendedConfig, serverCache cache.Cache) error {
+	c := ec.Config
 	requestWg := &sync.WaitGroup{}
 	dbUpdateWg := &sync.WaitGroup{}
 
 	go func() {
 		worker := initializeDBWorker(c.CacheDir, true)
+		if err := repopulateMetricGauge(ec.GaugeMetric, c.CacheDir); err != nil {
+			log.Logger.Errorf("%+v\n", err)
+		}
 		ctx := context.Background()
 		for {
 			time.Sleep(1 * time.Hour)
-			if err := worker.update(ctx, c.AppVersion, c.CacheDir, dbUpdateWg, requestWg); err != nil {
+			if err := worker.update(ctx, c.AppVersion, c.CacheDir, dbUpdateWg, requestWg, ec.GaugeMetric); err != nil {
 				log.Logger.Errorf("%+v\n", err)
 			}
 		}
 	}()
 
-	mux := newServeMux(serverCache, dbUpdateWg, requestWg, c.Token, c.TokenHeader)
+	mux := newServeMux(ec, serverCache, dbUpdateWg, requestWg, c.Token, c.TokenHeader)
 	log.Logger.Infof("Listening %s...", c.Listen)
 
 	return http.ListenAndServe(c.Listen, mux)
 }
 
-func newServeMux(serverCache cache.Cache, dbUpdateWg, requestWg *sync.WaitGroup, token, tokenHeader string) *http.ServeMux {
+func newServeMux(ec extendedconfig.ExtendedConfig, serverCache cache.Cache, dbUpdateWg, requestWg *sync.WaitGroup, token, tokenHeader string) *http.ServeMux {
 	withWaitGroup := func(base http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			// Stop processing requests during DB update
@@ -83,6 +91,9 @@ func newServeMux(serverCache cache.Cache, dbUpdateWg, requestWg *sync.WaitGroup,
 	libHandler := rpcDetector.NewLibDetectorServer(initializeLibServer(), nil)
 	mux.Handle(rpcDetector.LibDetectorPathPrefix, withToken(withWaitGroup(libHandler), token, tokenHeader))
 
+	promHandler := promhttp.HandlerFor(ec.MetricsRegistry, promhttp.HandlerOpts{Timeout: 10 * time.Second})
+	mux.Handle("/metrics", withToken(withWaitGroup(promHandler), ec.Config.Token, ec.Config.TokenHeader))
+
 	mux.HandleFunc("/healthz", func(rw http.ResponseWriter, r *http.Request) {
 		if _, err := rw.Write([]byte("ok")); err != nil {
 			log.Logger.Errorf("health check error: %s", err)
@@ -111,7 +122,10 @@ func newDBWorker(dbClient dbFile.Operation) dbWorker {
 }
 
 func (w dbWorker) update(ctx context.Context, appVersion, cacheDir string,
-	dbUpdateWg, requestWg *sync.WaitGroup) error {
+	dbUpdateWg, requestWg *sync.WaitGroup, gaugeMetric *prometheus.GaugeVec) error {
+	if err := updateLastDBUpdatePrometheus(gaugeMetric, float64(time.Now().Unix()), true); err != nil {
+		return err
+	}
 	log.Logger.Debug("Check for DB update...")
 	needsUpdate, err := w.dbClient.NeedsUpdate(appVersion, false, false)
 	if err != nil {
@@ -124,6 +138,9 @@ func (w dbWorker) update(ctx context.Context, appVersion, cacheDir string,
 	if err = w.hotUpdate(ctx, cacheDir, dbUpdateWg, requestWg); err != nil {
 		return xerrors.Errorf("failed DB hot update")
 	}
+	if err = repopulateMetricGauge(gaugeMetric, cacheDir); err != nil {
+		return xerrors.Errorf("error occurred while updating the prometheus gauge vec: %w", err)
+	}
 	return nil
 }
 
@@ -164,3 +181,30 @@ func (w dbWorker) hotUpdate(ctx context.Context, cacheDir string, dbUpdateWg, re
 
 	return nil
 }
+
+func repopulateMetricGauge(gauge *prometheus.GaugeVec, cacheDir string) error {
+	m := dbFile.NewMetadata(afero.NewOsFs(), cacheDir)
+	metadata, err := m.Get()
+	if err != nil {
+		return xerrors.Errorf("error populating the metrics at prometheus endpoint: %w", err)
+	}
+	if err = updateLastDBUpdatePrometheus(gauge, float64(metadata.UpdatedAt.Unix()), true); err != nil {
+		return err
+	}
+	if err = updateLastDBUpdatePrometheus(gauge, float64(metadata.UpdatedAt.Unix()), false); err != nil {
+		return err
+	}
+	return nil
+}
+
+func updateLastDBUpdatePrometheus(gauge *prometheus.GaugeVec, time float64, onlyDBAttempt bool) error {
+	if gauge == nil {
+		return xerrors.Errorf("prometheus gauge found to be nil while making last db update")
+	}
+	if onlyDBAttempt {
+		gauge.With(prometheus.Labels{"action": "last_db_update_attempt"}).Set(time)
+	} else {
+		gauge.With(prometheus.Labels{"action": "last_db_update"}).Set(time)
+	}
+	return nil
+}

pkg/rpc/server/listen_test.go

@@ -12,6 +12,10 @@ import (
 	"testing"
 	"time"
 
+	"github.com/aquasecurity/trivy/internal/config"
+	c "github.com/aquasecurity/trivy/internal/server/config"
+	"github.com/aquasecurity/trivy/internal/server/extendedconfig"
+	"github.com/prometheus/client_golang/prometheus"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
@@ -52,7 +56,8 @@ func Test_dbWorker_update(t *testing.T) {
 	}
 
 	type args struct {
-		appVersion string
+		appVersion  string
+		nilGaugeVec bool
 	}
 	tests := []struct {
 		name        string
@@ -109,6 +114,15 @@ func Test_dbWorker_update(t *testing.T) {
 			args:    args{appVersion: "1"},
 			wantErr: "failed DB hot update",
 		},
+		{
+			name: "Nil GaugeVec returns an error",
+			needsUpdate: needsUpdate{
+				input:  needsUpdateInput{appVersion: "1", skip: false},
+				output: needsUpdateOutput{needsUpdate: true},
+			},
+			args:    args{appVersion: "1", nilGaugeVec: true},
+			wantErr: "prometheus gauge found to be nil while making last db update",
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
@@ -135,14 +149,29 @@ func Test_dbWorker_update(t *testing.T) {
 						require.NoError(t, os.MkdirAll(filepath.Dir(dbPath), 0777), tt.name)
 						err = ioutil.WriteFile(dbPath, content, 0444)
 						require.NoError(t, err, tt.name)
+						dbDir := filepath.Join(cacheDir, "db")
+						err = db.Config{}.StoreMetadata(tt.want, dbDir)
+						require.NoError(t, err, tt.name)
 					}).Return(tt.download.err)
 			}
 
 			w := newDBWorker(mockDBClient)
-
+			var gaugeVec *prometheus.GaugeVec
+			if tt.args.nilGaugeVec {
+				gaugeVec = nil
+			} else {
+				gaugeVec = prometheus.NewGaugeVec(
+					prometheus.GaugeOpts{
+						Name: "trivy",
+						Help: "Gauge Metrics associated with trivy - Last DB Update, Last DB Update Attempt ...",
+					},
+					[]string{"action"},
+				)
+				prometheus.NewRegistry().MustRegister(gaugeVec)
+			}
 			var dbUpdateWg, requestWg sync.WaitGroup
 			err = w.update(context.Background(), tt.args.appVersion, cacheDir,
-				&dbUpdateWg, &requestWg)
+				&dbUpdateWg, &requestWg, gaugeVec)
 			if tt.wantErr != "" {
 				require.NotNil(t, err, tt.name)
 				assert.Contains(t, err.Error(), tt.wantErr, tt.name)
@@ -171,11 +200,13 @@ func Test_newServeMux(t *testing.T) {
 		tokenHeader string
 	}
 	tests := []struct {
-		name   string
-		args   args
-		path   string
-		header http.Header
-		want   int
+		name         string
+		globalConfig config.GlobalConfig
+		dbConfig     config.DBConfig
+		args         args
+		path         string
+		header       http.Header
+		want         int
 	}{
 		{
 			name: "health check",
@@ -185,6 +216,9 @@ func Test_newServeMux(t *testing.T) {
 		{
 			name: "cache endpoint",
 			path: path.Join(rpcCache.CachePathPrefix, "MissingBlobs"),
+			dbConfig: config.DBConfig{
+				Reset: true,
+			},
 			header: http.Header{
 				"Content-Type": []string{"application/protobuf"},
 			},
@@ -206,6 +240,9 @@ func Test_newServeMux(t *testing.T) {
 		{
 			name: "sad path: no handler",
 			path: "/sad",
+			dbConfig: config.DBConfig{
+				Reset: false,
+			},
 			header: http.Header{
 				"Content-Type": []string{"application/protobuf"},
 			},
@@ -228,11 +265,20 @@ func Test_newServeMux(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			dbUpdateWg, requestWg := &sync.WaitGroup{}, &sync.WaitGroup{}
 
-			c, err := cache.NewFSCache(t.TempDir())
+			internalConfig := &c.Config{
+				DBConfig: tt.dbConfig,
+			}
+
+			err := internalConfig.Init()
+
+			fsCache, err := cache.NewFSCache(t.TempDir())
 			require.NoError(t, err)
 
+			ec := extendedconfig.New(*internalConfig)
+			ec.Init()
+
 			ts := httptest.NewServer(newServeMux(
-				c, dbUpdateWg, requestWg, tt.args.token, tt.args.tokenHeader),
+				ec, fsCache, dbUpdateWg, requestWg, tt.args.token, tt.args.tokenHeader),
 			)
 			defer ts.Close()