perf: use a map to track which advisories should be checked for which…

… packages (#216)

* perf: use a map to track which advisories should be checked for which packages

* fix: ensure that vulnerabilities are deduplicated

* refactor: record the number of loaded vulnerabilities

* fix: don't load advisories that don't affect any packages

fixtures/configs-extra-dbs/db/osv-1.json

@@ -0,0 +1,17 @@
+{
+  "id": "OSV-1",
+  "affected": [
+    {
+      "package": {
+        "ecosystem": "npm",
+        "name": "request"
+      }
+    },
+    {
+      "package": {
+        "ecosystem": "npm",
+        "name": "@cypress/request"
+      }
+    }
+  ]
+}

fixtures/configs-extra-dbs/db/osv-2.json

@@ -0,0 +1,3 @@
+{
+  "id": "OSV-2"
+}

fixtures/configs-extra-dbs/db/osv-3.json

@@ -0,0 +1,3 @@
+{
+  "id": "OSV-3"
+}

main.go

@@ -170,7 +170,7 @@ func describeDB(db database.DB) string {
 			color.YellowString("%d", tt.BatchSize),
 		)
 	case *database.ZipDB:
-		count := len(tt.Vulnerabilities(true))
+		count := tt.VulnerabilitiesCount
 
 		return fmt.Sprintf(
 			"%s %s, including withdrawn - last updated %s",
@@ -179,7 +179,7 @@ func describeDB(db database.DB) string {
 			tt.UpdatedAt,
 		)
 	case *database.DirDB:
-		count := len(tt.Vulnerabilities(true))
+		count := tt.VulnerabilitiesCount
 
 		return fmt.Sprintf(
 			"%s %s, including withdrawn",

main_test.go

@@ -1036,12 +1036,12 @@ func TestRun_Configs(t *testing.T) {
 			wantStdout: `
 				Loaded the following OSV databases:
 					api#https://example.com/v1 (using batches of 1000)
-					dir#file:/fixtures/configs-extra-dbs (0 vulnerabilities, including withdrawn)
+					dir#file:/fixtures/configs-extra-dbs (3 vulnerabilities, including withdrawn)
 					zip#https://example.com/osvs/all
 				fixtures/configs-extra-dbs/yarn.lock: found 0 packages
 					Using config at fixtures/configs-extra-dbs/.osv-detector.yaml (0 ignores)
 					Using db api#https://example.com/v1 (using batches of 1000)
-					Using db dir#file:/fixtures/configs-extra-dbs (0 vulnerabilities, including withdrawn)
+					Using db dir#file:/fixtures/configs-extra-dbs (3 vulnerabilities, including withdrawn)
 
 					no known vulnerabilities found
 			`,

pkg/database/dir.go

@@ -29,7 +29,7 @@ var ErrDirPathWrongProtocol = errors.New("directory path must start with \"file:
 // load walks the filesystem starting with the working directory within the local path,
 // loading all OSVs found along the way.
 func (db *DirDB) load() error {
-	db.vulnerabilities = []OSV{}
+	db.vulnerabilities = make(map[string][]OSV)
 
 	if !strings.HasPrefix(db.LocalPath, "file:") {
 		return ErrDirPathWrongProtocol
@@ -78,7 +78,7 @@ func (db *DirDB) load() error {
 			return nil
 		}
 
-		db.vulnerabilities = append(db.vulnerabilities, pa)
+		db.addVulnerability(pa)
 
 		return nil
 	})

pkg/database/dir_test.go

@@ -9,7 +9,17 @@ import (
 func TestNewDirDB(t *testing.T) {
 	t.Parallel()
 
-	osvs := []database.OSV{{ID: "OSV-1"}, {ID: "OSV-2"}, {ID: "GHSA-1234"}}
+	osvs := []database.OSV{
+		withDefaultAffected("OSV-1"),
+		withDefaultAffected("OSV-2"),
+		{
+			ID: "GHSA-1234",
+			Affected: []database.Affected{
+				{Package: database.Package{Ecosystem: "npm", Name: "request"}},
+				{Package: database.Package{Ecosystem: "npm", Name: "@cypress/request"}},
+			},
+		},
+	}
 
 	db, err := database.NewDirDB(database.Config{URL: "file:/fixtures/db"}, false)
 
@@ -69,7 +79,7 @@ func TestNewDirDB_DoesNotExist(t *testing.T) {
 func TestNewDirDB_WorkingDirectory(t *testing.T) {
 	t.Parallel()
 
-	osvs := []database.OSV{{ID: "OSV-1"}}
+	osvs := []database.OSV{withDefaultAffected("OSV-1")}
 
 	db, err := database.NewDirDB(database.Config{URL: "file:/fixtures/db", WorkingDirectory: "nested-1"}, false)
 

pkg/database/fixtures/db/file.json

@@ -1,3 +1,17 @@
 {
-  "id": "GHSA-1234"
+  "id": "GHSA-1234",
+  "affected": [
+    {
+      "package": {
+        "ecosystem": "npm",
+        "name": "request"
+      }
+    },
+    {
+      "package": {
+        "ecosystem": "npm",
+        "name": "@cypress/request"
+      }
+    }
+  ]
 }

pkg/database/fixtures/db/nested-1/osv-1.json

@@ -1,3 +1,12 @@
 {
-  "id": "OSV-1"
+  "id": "OSV-1",
+  "affected": [
+    {
+      "package": {
+        "name": "mine",
+        "ecosystem": "PyPi"
+      },
+      "versions": []
+    }
+  ]
 }

pkg/database/fixtures/db/nested-2/osv-2.json

@@ -1,3 +1,12 @@
 {
-  "id": "OSV-2"
+  "id": "OSV-2",
+  "affected": [
+    {
+      "package": {
+        "name": "mine",
+        "ecosystem": "PyPi"
+      },
+      "versions": []
+    }
+  ]
 }

pkg/database/mem-check.go

@@ -7,19 +7,39 @@ import (
 // an OSV database that lives in-memory, and can be used by other structs
 // that handle loading the vulnerabilities from where ever
 type memDB struct {
-	vulnerabilities []OSV
+	vulnerabilities      map[string][]OSV
+	VulnerabilitiesCount int
 }
 
-func (db *memDB) Vulnerabilities(includeWithdrawn bool) []OSV {
-	if includeWithdrawn {
-		return db.vulnerabilities
+func (db *memDB) addVulnerability(osv OSV) {
+	db.VulnerabilitiesCount++
+
+	for _, affected := range osv.Affected {
+		hash := string(affected.Package.Ecosystem) + "-" + affected.Package.NormalizedName()
+		vulns := db.vulnerabilities[hash]
+
+		if vulns == nil {
+			vulns = []OSV{}
+		}
+
+		db.vulnerabilities[hash] = append(vulns, osv)
 	}
+}
 
+func (db *memDB) Vulnerabilities(includeWithdrawn bool) []OSV {
 	var vulnerabilities []OSV
+	ids := make(map[string]struct{})
 
-	for _, vulnerability := range db.vulnerabilities {
-		if vulnerability.Withdrawn == nil {
-			vulnerabilities = append(vulnerabilities, vulnerability)
+	for _, vulns := range db.vulnerabilities {
+		for _, vulnerability := range vulns {
+			if _, ok := ids[vulnerability.ID]; ok {
+				continue
+			}
+
+			if (vulnerability.Withdrawn == nil) || includeWithdrawn {
+				vulnerabilities = append(vulnerabilities, vulnerability)
+				ids[vulnerability.ID] = struct{}{}
+			}
 		}
 	}
 
@@ -29,9 +49,13 @@ func (db *memDB) Vulnerabilities(includeWithdrawn bool) []OSV {
 func (db *memDB) VulnerabilitiesAffectingPackage(pkg internal.PackageDetails) Vulnerabilities {
 	var vulnerabilities Vulnerabilities
 
-	for _, vulnerability := range db.Vulnerabilities(false) {
-		if vulnerability.IsAffected(pkg) && !vulnerabilities.Includes(vulnerability) {
-			vulnerabilities = append(vulnerabilities, vulnerability)
+	hash := string(pkg.Ecosystem) + "-" + pkg.Name
+
+	if vulns, ok := db.vulnerabilities[hash]; ok {
+		for _, vulnerability := range vulns {
+			if vulnerability.Withdrawn == nil && vulnerability.IsAffected(pkg) && !vulnerabilities.Includes(vulnerability) {
+				vulnerabilities = append(vulnerabilities, vulnerability)
+			}
 		}
 	}
 

pkg/database/zip.go

@@ -152,7 +152,7 @@ func (db *ZipDB) loadZipFile(zipFile *zip.File) {
 		return
 	}
 
-	db.vulnerabilities = append(db.vulnerabilities, osv)
+	db.addVulnerability(osv)
 }
 
 // load fetches a zip archive of the OSV database and loads known vulnerabilities
@@ -162,7 +162,7 @@ func (db *ZipDB) loadZipFile(zipFile *zip.File) {
 // so that a new version of the archive is only downloaded if it has been
 // modified, per HTTP caching standards.
 func (db *ZipDB) load() error {
-	db.vulnerabilities = []OSV{}
+	db.vulnerabilities = make(map[string][]OSV)
 
 	body, err := db.fetchZip()