datastore: drop non-matching vulnerabilities while scanning Get()

datastore/postgres/get.go

@@ -36,24 +36,39 @@ var (
 		},
 		[]string{"query"},
 	)
+	getVulnerabilitiesCandidateRows = promauto.NewCounter(
+		prometheus.CounterOpts{
+			Namespace: "claircore",
+			Subsystem: "vulnstore",
+			Name:      "getvulnerabilities_candidate_rows_total",
+			Help:      "Total number of vulnerability rows read by Get.",
+		},
+	)
+	getVulnerabilitiesAssociations = promauto.NewCounter(
+		prometheus.CounterOpts{
+			Namespace: "claircore",
+			Subsystem: "vulnstore",
+			Name:      "getvulnerabilities_accepted_associations_total",
+			Help:      "Total number of vulnerability associations retained by Get.",
+		},
+	)
 )
 
-type recordQuery struct {
-	record *claircore.IndexRecord
-	query  string
-}
-
-// Get implements vulnstore.Vulnerability.
+// Get implements [datastore.Vulnerability].
+//
+// If opts.Vulnerable is set, it is called for each candidate row while the
+// transaction and query batch are still open.
 func (s *MatcherStore) Get(ctx context.Context, records []*claircore.IndexRecord, opts datastore.GetOpts) (map[string][]*claircore.Vulnerability, error) {
 	tx, err := s.pool.Begin(ctx)
 	if err != nil {
 		return nil, err
 	}
 	defer tx.Rollback(ctx)
-	// start a batch
+	// Queue each distinct SELECT once. Records that render the same SQL are
+	// kept on that statement so each scanned row can be offered to all of them.
 	batch := &pgx.Batch{}
-	resCache := map[string][]*claircore.Vulnerability{}
-	rqs := []*recordQuery{}
+	bySQL := make(map[string][]*claircore.IndexRecord)
+	var queries []string
 	for _, record := range records {
 		query, err := buildGetQuery(record, &opts)
 		if err != nil {
@@ -63,13 +78,11 @@ func (s *MatcherStore) Get(ctx context.Context, records []*claircore.IndexRecord
 				"record", record)
 			continue
 		}
-		rqs = append(rqs, &recordQuery{query: query, record: record})
-		if _, ok := resCache[query]; ok {
-			continue
+		if _, ok := bySQL[query]; !ok {
+			queries = append(queries, query)
+			batch.Queue(query)
 		}
-		// queue the select query
-		batch.Queue(query)
-		resCache[query] = nil
+		bySQL[query] = append(bySQL[query], record)
 	}
 	// send the batch
 	start := time.Now()
@@ -82,22 +95,16 @@ func (s *MatcherStore) Get(ctx context.Context, records []*claircore.IndexRecord
 	results := make(map[string][]*claircore.Vulnerability)
 	vulnSet := make(map[string]map[string]struct{})
 	// intern shares decoded Vulnerability objects by vuln.id across queries
-	// in this Get. Later hits skip Scan so overlapping source-name rows are
-	// not fully decoded again.
+	// in this Get. A row is interned only after Vulnerable accepts it, so
+	// later hits skip Scan and rejected rows are not retained.
 	intern := map[int64]*claircore.Vulnerability{}
-	for _, rq := range rqs {
-		rid := rq.record.Package.ID
-		vulns, ok := resCache[rq.query]
-		if !ok {
-			return nil, fmt.Errorf("unexpected vulnerability query: %s", rq.query)
-		}
-		if vulns != nil { // We already have results we don't need to go back to the DB.
-			for _, v := range vulns {
-				addVuln(results, vulnSet, rid, v)
-			}
-			continue
-		}
-		results[rid] = []*claircore.Vulnerability{}
+	var candidates, accepted int
+	defer func() {
+		getVulnerabilitiesCandidateRows.Add(float64(candidates))
+		getVulnerabilitiesAssociations.Add(float64(accepted))
+	}()
+	for _, query := range queries {
+		recs := bySQL[query]
 		err := func() error {
 			rows, err := res.Query()
 			if err != nil {
@@ -106,56 +113,69 @@ func (s *MatcherStore) Get(ctx context.Context, records []*claircore.IndexRecord
 			}
 			defer rows.Close()
 			for rows.Next() {
+				candidates++
 				id, err := peekInt8(rows)
 				if err != nil {
 					res.Close()
 					return fmt.Errorf("failed to read vulnerability id: %w", err)
 				}
-				if v, ok := intern[id]; ok {
-					addVuln(results, vulnSet, rid, v)
-					continue
-				}
-				v := &claircore.Vulnerability{
-					Package: &claircore.Package{},
-					Dist:    &claircore.Distribution{},
-					Repo:    &claircore.Repository{},
+				v, ok := intern[id]
+				if !ok {
+					v = &claircore.Vulnerability{
+						Package: &claircore.Package{},
+						Dist:    &claircore.Distribution{},
+						Repo:    &claircore.Repository{},
+					}
+					err = rows.Scan(
+						&id,
+						&v.Name,
+						&v.Description,
+						&v.Issued,
+						&v.Links,
+						&v.Severity,
+						&v.NormalizedSeverity,
+						&v.Package.Name,
+						&v.Package.Version,
+						&v.Package.Module,
+						&v.Package.Arch,
+						&v.Package.Kind,
+						&v.Dist.DID,
+						&v.Dist.Name,
+						&v.Dist.Version,
+						&v.Dist.VersionCodeName,
+						&v.Dist.VersionID,
+						&v.Dist.Arch,
+						&v.Dist.CPE,
+						&v.Dist.PrettyName,
+						&v.ArchOperation,
+						&v.Repo.Name,
+						&v.Repo.Key,
+						&v.Repo.URI,
+						&v.FixedInVersion,
+						&v.Updater,
+						&v.Invert,
+					)
+					if err != nil {
+						res.Close()
+						return fmt.Errorf("failed to scan vulnerability: %w", err)
+					}
+					v.ID = strconv.FormatInt(id, 10)
 				}
-				err = rows.Scan(
-					&id,
-					&v.Name,
-					&v.Description,
-					&v.Issued,
-					&v.Links,
-					&v.Severity,
-					&v.NormalizedSeverity,
-					&v.Package.Name,
-					&v.Package.Version,
-					&v.Package.Module,
-					&v.Package.Arch,
-					&v.Package.Kind,
-					&v.Dist.DID,
-					&v.Dist.Name,
-					&v.Dist.Version,
-					&v.Dist.VersionCodeName,
-					&v.Dist.VersionID,
-					&v.Dist.Arch,
-					&v.Dist.CPE,
-					&v.Dist.PrettyName,
-					&v.ArchOperation,
-					&v.Repo.Name,
-					&v.Repo.Key,
-					&v.Repo.URI,
-					&v.FixedInVersion,
-					&v.Updater,
-					&v.Invert,
-				)
-				if err != nil {
-					res.Close()
-					return fmt.Errorf("failed to scan vulnerability: %w", err)
+				for _, record := range recs {
+					if opts.Vulnerable != nil {
+						ok, err := opts.Vulnerable(ctx, record, v)
+						if err != nil {
+							res.Close()
+							return err
+						}
+						if !ok {
+							continue
+						}
+					}
+					intern[id] = v
+					addVuln(results, vulnSet, record.Package.ID, v)
+					accepted++
 				}
-				v.ID = strconv.FormatInt(id, 10)
-				intern[id] = v
-				addVuln(results, vulnSet, rid, v)
 			}
 			if err := rows.Err(); err != nil {
 				res.Close()
@@ -166,7 +186,6 @@ func (s *MatcherStore) Get(ctx context.Context, records []*claircore.IndexRecord
 		if err != nil {
 			return nil, err
 		}
-		resCache[rq.query] = results[rid]
 	}
 	if err := res.Close(); err != nil {
 		return nil, fmt.Errorf("some weird batch error: %v", err)

datastore/postgres/get_test.go

@@ -1,7 +1,9 @@
 package postgres
 
 import (
+	"context"
 	"encoding/binary"
+	"errors"
 	"slices"
 	"testing"
 
@@ -122,6 +124,153 @@ func TestGetInternsOverlappingRows(t *testing.T) {
 	}
 }
 
+func TestGetVulnerable(t *testing.T) {
+	integration.NeedDB(t)
+	ctx := test.Logging(t)
+
+	pool := pgtest.TestMatcherDB(ctx, t)
+	store := NewMatcherStore(pool)
+
+	srcKind := types.SourcePackage
+	binKind := types.BinaryPackage
+	_, err := store.UpdateVulnerabilities(ctx, "test-updater", driver.Fingerprint(uuid.New().String()), []*claircore.Vulnerability{
+		{Updater: "test-updater", Name: "CVE-YES", Package: &claircore.Package{Name: "bash", Kind: binKind}},
+		{Updater: "test-updater", Name: "CVE-NO", Package: &claircore.Package{Name: "bash", Kind: binKind}},
+		{Updater: "test-updater", Name: "CVE-SRC", Package: &claircore.Package{Name: "kernel", Kind: srcKind}},
+		{Updater: "test-updater", Name: "CVE-CORE", Package: &claircore.Package{Name: "kernel-core", Kind: binKind}},
+		{Updater: "test-updater", Name: "CVE-MOD", Package: &claircore.Package{Name: "kernel-modules", Kind: binKind}},
+		{Updater: "test-updater", Name: "CVE-EXTRA", Package: &claircore.Package{Name: "kernel-extra", Kind: binKind}},
+	})
+	if err != nil {
+		t.Fatalf("update: %v", err)
+	}
+
+	bash := &claircore.Package{ID: "bash", Name: "bash", Kind: binKind, Source: &claircore.Package{}}
+	t.Run("nil", func(t *testing.T) {
+		res, err := store.Get(ctx, []*claircore.IndexRecord{{Package: bash}}, datastore.GetOpts{})
+		if err != nil {
+			t.Fatal(err)
+		}
+		if diff := cmp.Diff([]string{"CVE-NO", "CVE-YES"}, vulnNames(res["bash"])); diff != "" {
+			t.Fatal(diff)
+		}
+	})
+	t.Run("drops", func(t *testing.T) {
+		res, err := store.Get(ctx, []*claircore.IndexRecord{{Package: bash}}, datastore.GetOpts{
+			Vulnerable: func(context.Context, *claircore.IndexRecord, *claircore.Vulnerability) (bool, error) {
+				return false, nil
+			},
+		})
+		if err != nil {
+			t.Fatal(err)
+		}
+		if len(res["bash"]) != 0 {
+			t.Fatalf("kept %s", vulnNames(res["bash"]))
+		}
+	})
+	t.Run("either record", func(t *testing.T) {
+		res, err := store.Get(ctx, []*claircore.IndexRecord{
+			{Package: bash, Repository: &claircore.Repository{Name: "no"}},
+			{Package: bash, Repository: &claircore.Repository{Name: "yes"}},
+		}, datastore.GetOpts{
+			Vulnerable: func(_ context.Context, record *claircore.IndexRecord, vuln *claircore.Vulnerability) (bool, error) {
+				return record.Repository.Name == "yes" && vuln.Name == "CVE-YES", nil
+			},
+		})
+		if err != nil {
+			t.Fatal(err)
+		}
+		if diff := cmp.Diff([]string{"CVE-YES"}, vulnNames(res["bash"])); diff != "" {
+			t.Fatal(diff)
+		}
+		if len(res["bash"]) != 1 {
+			t.Fatalf("got %d rows", len(res["bash"]))
+		}
+	})
+	t.Run("both records once", func(t *testing.T) {
+		res, err := store.Get(ctx, []*claircore.IndexRecord{
+			{Package: bash, Repository: &claircore.Repository{Name: "a"}},
+			{Package: bash, Repository: &claircore.Repository{Name: "b"}},
+		}, datastore.GetOpts{
+			Vulnerable: func(_ context.Context, _ *claircore.IndexRecord, vuln *claircore.Vulnerability) (bool, error) {
+				return vuln.Name == "CVE-YES", nil
+			},
+		})
+		if err != nil {
+			t.Fatal(err)
+		}
+		if diff := cmp.Diff([]string{"CVE-YES"}, vulnNames(res["bash"])); diff != "" {
+			t.Fatal(diff)
+		}
+		if len(res["bash"]) != 1 {
+			t.Fatalf("got %d rows", len(res["bash"]))
+		}
+	})
+	t.Run("two statements", func(t *testing.T) {
+		a := &claircore.Package{ID: "pkg", Name: "twostmt", Kind: binKind, Module: "a", Source: &claircore.Package{}}
+		b := &claircore.Package{ID: "pkg", Name: "twostmt", Kind: binKind, Module: "b", Source: &claircore.Package{}}
+		_, err := store.UpdateVulnerabilities(ctx, "test-updater-mod", driver.Fingerprint(uuid.New().String()), []*claircore.Vulnerability{
+			{Updater: "test-updater-mod", Name: "CVE-MOD-A", Package: &claircore.Package{Name: "twostmt", Kind: binKind, Module: "a"}},
+			{Updater: "test-updater-mod", Name: "CVE-MOD-B", Package: &claircore.Package{Name: "twostmt", Kind: binKind, Module: "b"}},
+		})
+		if err != nil {
+			t.Fatal(err)
+		}
+		res, err := store.Get(ctx, []*claircore.IndexRecord{{Package: a}, {Package: b}}, datastore.GetOpts{
+			Matchers: []driver.MatchConstraint{driver.PackageModule},
+		})
+		if err != nil {
+			t.Fatal(err)
+		}
+		if diff := cmp.Diff([]string{"CVE-MOD-A", "CVE-MOD-B"}, vulnNames(res["pkg"])); diff != "" {
+			t.Fatal(diff)
+		}
+	})
+	t.Run("error", func(t *testing.T) {
+		_, err := store.Get(ctx, []*claircore.IndexRecord{{Package: bash}}, datastore.GetOpts{
+			Vulnerable: func(context.Context, *claircore.IndexRecord, *claircore.Vulnerability) (bool, error) {
+				return false, errors.New("vulnerable failed")
+			},
+		})
+		if err == nil {
+			t.Fatal("expected Vulnerable error")
+		}
+	})
+
+	src := &claircore.Package{Name: "kernel", Kind: srcKind}
+	t.Run("shared row stays interned", func(t *testing.T) {
+		res, err := store.Get(ctx, []*claircore.IndexRecord{
+			{Package: &claircore.Package{ID: "core", Name: "kernel-core", Kind: binKind, Source: src}},
+			{Package: &claircore.Package{ID: "mod", Name: "kernel-modules", Kind: binKind, Source: src}},
+			{Package: &claircore.Package{ID: "extra", Name: "kernel-extra", Kind: binKind, Source: src}},
+		}, datastore.GetOpts{
+			Vulnerable: func(_ context.Context, record *claircore.IndexRecord, vuln *claircore.Vulnerability) (bool, error) {
+				if record.Package.Name == "kernel-modules" {
+					return vuln.Name == "CVE-MOD", nil
+				}
+				return vuln.Name != "CVE-MOD", nil
+			},
+		})
+		if err != nil {
+			t.Fatal(err)
+		}
+		if diff := cmp.Diff([]string{"CVE-CORE", "CVE-SRC"}, vulnNames(res["core"])); diff != "" {
+			t.Fatal(diff)
+		}
+		if diff := cmp.Diff([]string{"CVE-MOD"}, vulnNames(res["mod"])); diff != "" {
+			t.Fatal(diff)
+		}
+		if diff := cmp.Diff([]string{"CVE-EXTRA", "CVE-SRC"}, vulnNames(res["extra"])); diff != "" {
+			t.Fatal(diff)
+		}
+		srcCore := vulnByName(res["core"], "CVE-SRC")
+		srcExtra := vulnByName(res["extra"], "CVE-SRC")
+		if srcCore == nil || srcExtra == nil || srcCore != srcExtra {
+			t.Fatal("expected kernel-core and kernel-extra to share CVE-SRC")
+		}
+	})
+}
+
 func vulnNames(vs []*claircore.Vulnerability) []string {
 	out := make([]string, 0, len(vs))
 	for _, v := range vs {

datastore/vulnerability.go

@@ -19,6 +19,10 @@ type GetOpts struct {
 	// VersionFiltering enables filtering based on the normalized versions in
 	// the database.
 	VersionFiltering bool
+	// Vulnerable reports whether a scanned vulnerability affects this record.
+	// A nil Vulnerable retains every row. When several records share a package
+	// ID, the row is retained if any of them accepts it.
+	Vulnerable func(ctx context.Context, record *claircore.IndexRecord, vuln *claircore.Vulnerability) (bool, error)
 }
 
 // Vulnerability is the interface for querying stored Vulnerabilities.