stackrox · RTann · Oct 18, 2024 · Oct 7, 2024
@@ -34,7 +34,6 @@ import (
 	"github.com/stackrox/rox/scanner/enricher/fixedby"
 	"github.com/stackrox/rox/scanner/enricher/nvd"
 	"github.com/stackrox/rox/scanner/internal/httputil"
-	"github.com/stackrox/rox/scanner/matcher/updater/distribution"
 	"github.com/stackrox/rox/scanner/matcher/updater/vuln"
 )
 
@@ -83,8 +82,7 @@ type matcherImpl struct {
 	metadataStore postgres.MatcherMetadataStore
 	pool          *pgxpool.Pool
 
-	vulnUpdater   *vuln.Updater
-	distroUpdater *distribution.Updater
+	vulnUpdater *vuln.Updater
 }
 
 // NewMatcher creates a new matcher.
@@ -177,18 +175,6 @@ func NewMatcher(ctx context.Context, cfg config.MatcherConfig) (Matcher, error)
 		return nil, fmt.Errorf("creating vuln updater: %w", err)
 	}
 
-	distroUpdater, err := distribution.New(ctx, store, vulnUpdater.Initialized)
-	if err != nil {
-		return nil, fmt.Errorf("creating known-distribution updater: %w", err)
-	}
-
-	// Start the known-distributions updater.
-	go func() {
-		if err := distroUpdater.Start(); err != nil {
-			zlog.Error(ctx).Err(err).Msg("known-distributions updater failed")
-		}
-	}()
-
 	// Start the vulnerability updater.
 	go func() {
 		if err := vulnUpdater.Start(); err != nil {
@@ -202,8 +188,7 @@ func NewMatcher(ctx context.Context, cfg config.MatcherConfig) (Matcher, error)
 		metadataStore: metadataStore,
 		pool:          pool,
 
-		vulnUpdater:   vulnUpdater,
-		distroUpdater: distroUpdater,
+		vulnUpdater: vulnUpdater,
 	}, nil
 }
 
@@ -218,13 +203,13 @@ func (m *matcherImpl) GetLastVulnerabilityUpdate(ctx context.Context) (time.Time
 }
 
 func (m *matcherImpl) GetKnownDistributions(_ context.Context) []claircore.Distribution {
-	return m.distroUpdater.Known()
+	return m.vulnUpdater.KnownDistributions()
 }
 
 // Close closes the matcher.
 func (m *matcherImpl) Close(ctx context.Context) error {
 	ctx = zlog.ContextWithValues(ctx, "component", "scanner/backend/matcher.Close")
-	err := errors.Join(m.distroUpdater.Stop(), m.vulnUpdater.Stop(), m.libVuln.Close(ctx))
+	err := errors.Join(m.vulnUpdater.Stop(), m.libVuln.Close(ctx))
 	m.pool.Close()
 	return err
 }

@@ -0,0 +1,63 @@
+package vuln
+
+import (
+	"context"
+	"slices"
+
+	"github.com/quay/claircore"
+	"github.com/quay/zlog"
+	"github.com/stackrox/rox/pkg/sync"
+	"github.com/stackrox/rox/scanner/datastore/postgres"
+)
+
+// distManager manages updates to known-distributions.
+type distManager struct {
+	store postgres.MatcherStore
+
+	// mutex protects access to known.
+	mutex sync.RWMutex
+	// known tracks all known distributions.
+	known []claircore.Distribution
+}
+
+func newDistManager(store postgres.MatcherStore) *distManager {
+	return &distManager{
+		store: store,
+	}
+}
+
+// update updates the currently known distributions, on-demand.
+//
+// It is *not* safe to call update concurrently.
+func (dm *distManager) update(ctx context.Context) (err error) {
+	ctx = zlog.ContextWithValues(ctx, "component", "updater/vuln/dists.update")
+	zlog.Info(ctx).Msg("updating vuln distributions")
+	defer func() {
+		zlog.Info(ctx).Bool("updated", err == nil).Msg("done updating vuln distributions")
+	}()
+
+	var dists []claircore.Distribution
+	dists, err = dm.store.Distributions(ctx)
+	if err != nil {
+		return
+	}
+
+	func() {
+		dm.mutex.Lock()
+		defer dm.mutex.Unlock()
+
+		dm.known = nil // Hint to GC.
+		dm.known = dists
+	}()
+
+	return
+}
+
+// get returns the currently known distributions.
+func (dm *distManager) get() []claircore.Distribution {
+	dm.mutex.RLock()
+	defer dm.mutex.RUnlock()
+
+	known := slices.Clone(dm.known)
+	return known
+}
@@ -1,20 +1,20 @@
-package distribution
+package vuln
 
 import (
 	"context"
+	"errors"
 	"testing"
 
-	"github.com/pkg/errors"
 	"github.com/quay/claircore"
 	"github.com/stackrox/rox/scanner/datastore/postgres/mocks"
 	"github.com/stretchr/testify/assert"
 	"go.uber.org/mock/gomock"
 )
 
-func TestUpdate(t *testing.T) {
+func TestDistManager(t *testing.T) {
 	ctx := context.Background()
 	store := mocks.NewMockMatcherStore(gomock.NewController(t))
-	u := &Updater{
+	m := &distManager{
 		store: store,
 	}
 
@@ -51,15 +51,18 @@ func TestUpdate(t *testing.T) {
 		},
 	}
 
+	// Get nothing.
+	assert.Nil(t, m.get())
+
 	// Successful fetch.
 	store.EXPECT().Distributions(gomock.Any()).Return(dists, nil)
-	err := u.update(ctx)
+	err := m.update(ctx)
 	assert.NoError(t, err)
-	assert.ElementsMatch(t, dists, u.Known())
+	assert.ElementsMatch(t, dists, m.get())
 
 	// Unsuccessful should return same dists as before.
 	store.EXPECT().Distributions(gomock.Any()).Return(nil, errors.New("error"))
-	err = u.update(ctx)
+	err = m.update(ctx)
 	assert.Error(t, err)
-	assert.ElementsMatch(t, dists, u.Known())
+	assert.ElementsMatch(t, dists, m.get())
 }
@@ -20,7 +20,6 @@ import (
 	"github.com/google/uuid"
 	"github.com/klauspost/compress/zstd"
 	"github.com/quay/claircore"
-	"github.com/quay/claircore/datastore"
 	"github.com/quay/claircore/libvuln"
 	"github.com/quay/claircore/libvuln/driver"
 	"github.com/quay/claircore/libvuln/updates"
@@ -64,7 +63,7 @@ var (
 // Store, Locker, MetadataStore, and URL are required.
 // The rest are optional.
 type Opts struct {
-	Store         datastore.MatcherStore
+	Store         postgres.MatcherStore
 	Locker        *ctxlock.Locker
 	MetadataStore postgres.MatcherMetadataStore
 
@@ -87,7 +86,7 @@ type Updater struct {
 	ctx    context.Context
 	cancel context.CancelFunc
 
-	store         datastore.MatcherStore
+	store         postgres.MatcherStore
 	locker        updates.LockSource
 	metadataStore postgres.MatcherMetadataStore
 
@@ -110,6 +109,8 @@ type Updater struct {
 
 	retryDelay time.Duration
 	retryMax   int
+
+	distManager *distManager
 }
 
 // New creates a new Updater based on the given options.
@@ -138,6 +139,8 @@ func New(ctx context.Context, opts Opts) (*Updater, error) {
 
 		retryDelay: opts.RetryDelay,
 		retryMax:   opts.RetryMax,
+
+		distManager: newDistManager(opts.Store),
 	}
 	u.importFunc = func(ctx context.Context, reader io.Reader) error {
 		return u.Import(ctx, reader)
@@ -302,6 +305,10 @@ func (u *Updater) Start() error {
 		go u.runGCFullPeriodic()
 	}
 
+	if err := u.distManager.update(ctx); err != nil {
+		zlog.Warn(ctx).Err(err).Msg("failed to initialize known-distributions")
+	}
+
 	// Start immediately, all matchers will compete to update each vulnerability
 	// bundle if multi-bundle mode is on, or the single bundle.
 	timer := time.NewTimer(0)
@@ -345,6 +352,10 @@ func (u *Updater) Initialized(ctx context.Context) bool {
 	return true
 }
 
+func (u *Updater) KnownDistributions() []claircore.Distribution {
+	return u.distManager.get()
+}
+
 // Update runs the full vulnerability update process.
 //
 // Note: periodic full GC will not be started.
@@ -434,6 +445,10 @@ func (u *Updater) runSingleBundleUpdate(ctx context.Context) (bool, error) {
 		return false, err
 	}
 
+	if err := u.distManager.update(ctx); err != nil {
+		return false, fmt.Errorf("updating known-distributions: %w", err)
+	}
+
 	if u.initialized.CompareAndSwap(false, true) {
 		zlog.Info(ctx).Msg("finished initial updater run: setting updater to initialized")
 	}
@@ -505,6 +520,11 @@ func (u *Updater) runMultiBundleUpdate(ctx context.Context) (bool, error) {
 		return false, fmt.Errorf("cleaning vuln updates: %w", err)
 	}
 
+	err = u.distManager.update(ctx)
+	if err != nil {
+		return false, fmt.Errorf("updating known-distributions: %w", err)
+	}
+
 	_ = u.Initialized(ctx)
 
 	return true, nil