Infrastructure/zot
0
Fork 0
mirror of https://github.com/project-zot/zot.git synced 2024-12-30 22:34:13 -05:00
Code Issues Activity

fix: trivydb update now uses task scheduler (#1204)

Browse source 
Signed-off-by: Ana-Roberta Lisca <ana.kagome@yahoo.com>
This commit is contained in:
Lisca Ana-Roberta 2023-03-02 19:43:54 +02:00 committed by GitHub
parent 5a2fb4108d
commit 6bbf730061
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 107 additions and 23 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
pkg/api/controller.go
View file

@ -634,7 +634,7 @@ func (c *Controller) StartBackgroundTasks(reloadCtx context.Context) {
// Enable extensions if extension config is provided for DefaultStore
if c.Config != nil && c.Config.Extensions != nil {
ext.EnableMetricsExtension(c.Config, c.Log, c.Config.Storage.RootDirectory)
ext.EnableSearchExtension(c.Config, c.StoreController, c.RepoDB, c.CveInfo, c.Log)
ext.EnableSearchExtension(c.Config, c.StoreController, c.RepoDB, taskScheduler, c.CveInfo, c.Log)
}
if c.Config.Storage.SubPaths != nil {

108
pkg/extensions/extension_search.go
View file

@ -5,6 +5,7 @@ package extensions
import (
"net/http"
"sync"
"time"
gqlHandler "github.com/99designs/gqlgen/graphql/handler"
@ -18,10 +19,20 @@ import (
"zotregistry.io/zot/pkg/extensions/search/gql_generated"
"zotregistry.io/zot/pkg/log"
"zotregistry.io/zot/pkg/meta/repodb"
"zotregistry.io/zot/pkg/scheduler"
"zotregistry.io/zot/pkg/storage"
)
type CveInfo cveinfo.CveInfo
type (
CveInfo cveinfo.CveInfo
state int
)
const (
pending state = iota
running
done
)
func GetCVEInfo(config *config.Config, storeController storage.StoreController,
repoDB repodb.RepoDB, log log.Logger,
@ -40,7 +51,7 @@ func GetCVEInfo(config *config.Config, storeController storage.StoreController,
}
func EnableSearchExtension(config *config.Config, storeController storage.StoreController,
repoDB repodb.RepoDB, cveInfo CveInfo, log log.Logger,
repoDB repodb.RepoDB, taskScheduler *scheduler.Scheduler, cveInfo CveInfo, log log.Logger,
) {
if config.Extensions.Search != nil && *config.Extensions.Search.Enable && config.Extensions.Search.CVE != nil {
defaultUpdateInterval, _ := time.ParseDuration("2h")
@ -51,30 +62,93 @@ func EnableSearchExtension(config *config.Config, storeController storage.StoreC
log.Warn().Msg("CVE update interval set to too-short interval < 2h, changing update duration to 2 hours and continuing.") //nolint:lll // gofumpt conflicts with lll
}
go func() {
err := downloadTrivyDB(cveInfo, log, config.Extensions.Search.CVE.UpdateInterval)
if err != nil {
log.Error().Err(err).Msg("error while downloading TrivyDB")
}
}()
updateInterval := config.Extensions.Search.CVE.UpdateInterval
downloadTrivyDB(updateInterval, taskScheduler, cveInfo, log)
} else {
log.Info().Msg("CVE config not provided, skipping CVE update")
}
}
func downloadTrivyDB(cveInfo CveInfo, log log.Logger, updateInterval time.Duration) error {
for {
log.Info().Msg("updating the CVE database")
func downloadTrivyDB(interval time.Duration, sch *scheduler.Scheduler, cveInfo CveInfo, log log.Logger) {
generator := &trivyTaskGenerator{interval, cveInfo, log, pending, 0, time.Now(), &sync.Mutex{}}
err := cveInfo.UpdateDB()
if err != nil {
return err
}
sch.SubmitGenerator(generator, interval, scheduler.HighPriority)
}
log.Info().Str("DB update completed, next update scheduled after", updateInterval.String()).Msg("")
type trivyTaskGenerator struct {
interval time.Duration
cveInfo CveInfo
log log.Logger
status state
waitTime time.Duration
lastTaskTime time.Time
lock *sync.Mutex
}
time.Sleep(updateInterval)
func (gen *trivyTaskGenerator) GenerateTask() (scheduler.Task, error) {
var newTask scheduler.Task
gen.lock.Lock()
if gen.status != running && time.Since(gen.lastTaskTime) >= gen.waitTime {
newTask = newTrivyTask(gen.interval, gen.cveInfo, gen, gen.log)
gen.status = running
}
gen.lock.Unlock()
return newTask, nil
}
func (gen *trivyTaskGenerator) IsDone() bool {
gen.lock.Lock()
status := gen.status
gen.lock.Unlock()
return status == done
}
func (gen *trivyTaskGenerator) Reset() {
gen.lock.Lock()
gen.status = pending
gen.waitTime = 0
gen.lock.Unlock()
}
type trivyTask struct {
interval time.Duration
cveInfo cveinfo.CveInfo
generator *trivyTaskGenerator
log log.Logger
}
func newTrivyTask(interval time.Duration, cveInfo cveinfo.CveInfo,
generator *trivyTaskGenerator, log log.Logger,
) *trivyTask {
return &trivyTask{interval, cveInfo, generator, log}
}
func (trivyT *trivyTask) DoWork() error {
trivyT.log.Info().Msg("updating the CVE database")
err := trivyT.cveInfo.UpdateDB()
if err != nil {
trivyT.generator.lock.Lock()
trivyT.generator.status = pending
trivyT.generator.waitTime += time.Second
trivyT.generator.lastTaskTime = time.Now()
trivyT.generator.lock.Unlock()
return err
}
trivyT.generator.lock.Lock()
trivyT.generator.lastTaskTime = time.Now()
trivyT.generator.status = done
trivyT.generator.lock.Unlock()
trivyT.log.Info().Str("DB update completed, next update scheduled after", trivyT.interval.String()).Msg("")
return nil
}
func addSearchSecurityHeaders(h http.Handler) http.HandlerFunc { //nolint:varnamelen

3
pkg/extensions/extension_search_disabled.go
View file

@ -10,6 +10,7 @@ import (
"zotregistry.io/zot/pkg/api/config"
"zotregistry.io/zot/pkg/log"
"zotregistry.io/zot/pkg/meta/repodb"
"zotregistry.io/zot/pkg/scheduler"
"zotregistry.io/zot/pkg/storage"
)
@ -23,7 +24,7 @@ func GetCVEInfo(config *config.Config, storeController storage.StoreController,
// EnableSearchExtension ...
func EnableSearchExtension(config *config.Config, storeController storage.StoreController,
repoDB repodb.RepoDB, cveInfo CveInfo, log log.Logger,
repoDB repodb.RepoDB, scheduler *scheduler.Scheduler, cveInfo CveInfo, log log.Logger,
) {
log.Warn().Msg("skipping enabling search extension because given zot binary doesn't include this feature," +
"please build a binary that does so")

17
pkg/extensions/search/cve/cve_test.go
View file

@ -509,11 +509,20 @@ func TestCVESearch(t *testing.T) {
ctrlManager.StartAndWait(port)
// trivy db download fail
err = os.Mkdir(dbDir+"/_trivy", 0o000)
So(err, ShouldBeNil)
found, err := ReadLogFileAndSearchString(logPath, "Error downloading Trivy DB to destination dir", 180*time.Second)
So(err, ShouldBeNil)
So(found, ShouldBeTrue)
err = os.Chmod(dbDir+"/_trivy", 0o755)
So(err, ShouldBeNil)
// Wait for trivy db to download
_, err = ReadLogFileAndSearchString(logPath, "DB update completed, next update scheduled", 90*time.Second)
if err != nil {
panic(err)
}
found, err = ReadLogFileAndSearchString(logPath, "DB update completed, next update scheduled", 180*time.Second)
So(err, ShouldBeNil)
So(found, ShouldBeTrue)
defer ctrlManager.StopServer()