enrichment: Optimize the way enrichments are queried (#570)

Signed-off-by: crozzy <joseph.crosland@gmail.com>
quay · Feb 21, 2022 · f91fd65 · f91fd65
1 parent 6ec7ec5
commit f91fd65
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 7 deletions.
diff --git a/enricher/cvss/cvss.go b/enricher/cvss/cvss.go
@@ -12,6 +12,7 @@ import (
 	"net/http"
 	"net/url"
 	"regexp"
+	"sort"
 	"strings"
 	"time"
 
@@ -255,6 +256,8 @@ func (e *Enricher) Enrich(ctx context.Context, g driver.EnrichmentGetter, r *cla
 	// We return any CVSS blobs for CVEs mentioned in the free-form parts of the
 	// vulnerability.
 	m := make(map[string][]json.RawMessage)
+
+	erCache := make(map[string][]driver.EnrichmentRecord)
 	for id, v := range r.Vulnerabilities {
 		t := make(map[string]struct{})
 		ctx := zlog.ContextWithValues(ctx,
@@ -278,9 +281,17 @@ func (e *Enricher) Enrich(ctx context.Context, g driver.EnrichmentGetter, r *cla
 		zlog.Debug(ctx).
 			Strs("cve", ts).
 			Msg("found CVEs")
-		rec, err := g.GetEnrichment(ctx, ts)
-		if err != nil {
-			return "", nil, err
+
+		sort.Strings(ts)
+		cveKey := strings.Join(ts, "_")
+		rec, ok := erCache[cveKey]
+		if !ok {
+			var err error
+			rec, err = g.GetEnrichment(ctx, ts)
+			if err != nil {
+				return "", nil, err
+			}
+			erCache[cveKey] = rec
 		}
 		zlog.Debug(ctx).
 			Int("count", len(rec)).

diff --git a/enricher/cvss/cvss_test.go b/enricher/cvss/cvss_test.go
@@ -289,7 +289,7 @@ func (tc parseTestcase) Run(ctx context.Context, srv *httptest.Server) func(*tes
 
 func TestEnrich(t *testing.T) {
 	t.Parallel()
-	ctx := zlog.Test(nil, t)
+	ctx := zlog.Test(context.TODO(), t)
 	feedIn, err := os.Open("testdata/feed.json")
 	if err != nil {
 		t.Fatal(err)
@@ -301,15 +301,18 @@ func TestEnrich(t *testing.T) {
 	g := &fakeGetter{itemFeed: f}
 	r := &claircore.VulnerabilityReport{
 		Vulnerabilities: map[string]*claircore.Vulnerability{
-			"-1": &claircore.Vulnerability{
+			"-1": {
 				Description: "This is a fake vulnerability that doesn't have a CVE.",
 			},
-			"1": &claircore.Vulnerability{
+			"1": {
 				Description: "This is a fake vulnerability that looks like CVE-2021-0498.",
 			},
-			"6004": &claircore.Vulnerability{
+			"6004": {
 				Description: "CVE-2020-6004 was unassigned",
 			},
+			"6005": {
+				Description: "CVE-2021-0498 duplicate",
+			},
 		},
 	}
 	e := &Enricher{}
@@ -335,6 +338,20 @@ func TestEnrich(t *testing.T) {
 			"baseScore":             7.8,
 			"baseSeverity":          "HIGH",
 		}},
+		"6005": {{
+			"version":               "3.1",
+			"vectorString":          "CVSS:3.1/AV:L/AC:L/PR:L/UI:N/S:U/C:H/I:H/A:H",
+			"attackVector":          "LOCAL",
+			"attackComplexity":      "LOW",
+			"privilegesRequired":    "LOW",
+			"userInteraction":       "NONE",
+			"scope":                 "UNCHANGED",
+			"confidentialityImpact": "HIGH",
+			"integrityImpact":       "HIGH",
+			"availabilityImpact":    "HIGH",
+			"baseScore":             7.8,
+			"baseSeverity":          "HIGH",
+		}},
 	}
 	got := map[string][]map[string]interface{}{}
 	if err := json.Unmarshal(es[0], &got); err != nil {