From 2e2d0fb983e942f56c0b18d3b09969d3e5a4bb9f Mon Sep 17 00:00:00 2001
From: Victor Lyuboslavsky <victor.lyuboslavsky@gmail.com>
Date: Fri, 21 Jun 2024 14:14:56 -0500
Subject: [PATCH] Removed duplicate `os_versions` results in
 /api/latest/fleet/vulnerabilities/:cve endpoint (#19912)

#19819
Removed duplicate `os_versions` results in
/api/latest/fleet/vulnerabilities/:cve endpoint

Could not manually test since I do not have an Intel mac. We will need
to QA on dogfood.

# Checklist for submitter

If some of the following don't apply, delete the relevant line.

<!-- Note that API documentation changes are now addressed by the
product design team. -->

- [x] Changes file added for user-visible changes in `changes/`,
`orbit/changes/` or `ee/fleetd-chrome/changes`.
See [Changes
files](https://fleetdm.com/docs/contributing/committing-changes#changes-files)
for more information.
- [x] Added/updated tests
- [ ] Manual QA for all new/changed functionality
---
 changes/19819-duplicate-os-versions           |  1 +
 server/datastore/mysql/vulnerabilities.go     | 24 +++++++++++--
 .../datastore/mysql/vulnerabilities_test.go   | 34 ++++++++++++++++---
 3 files changed, 53 insertions(+), 6 deletions(-)
 create mode 100644 changes/19819-duplicate-os-versions

diff --git a/changes/19819-duplicate-os-versions b/changes/19819-duplicate-os-versions
new file mode 100644
index 0000000000..d6147dfd14
--- /dev/null
+++ b/changes/19819-duplicate-os-versions
@@ -0,0 +1 @@
+- Removed duplicate `os_versions` results in /api/latest/fleet/vulnerabilities/:cve endpoint
diff --git a/server/datastore/mysql/vulnerabilities.go b/server/datastore/mysql/vulnerabilities.go
index f9d57fc822..6208fd6d92 100644
--- a/server/datastore/mysql/vulnerabilities.go
+++ b/server/datastore/mysql/vulnerabilities.go
@@ -119,10 +119,11 @@ func (ds *Datastore) OSVersionsByCVE(ctx context.Context, cve string, teamID *ui
 
 	updatedAt = osvs.CountsUpdatedAt
 
-	var osVersionWithResolved []struct {
+	type osVersionWithResolvedType struct {
 		OSVersionID     uint    `db:"os_version_id"`
 		ResolvedVersion *string `db:"resolved_in_version"`
 	}
+	var osVersionWithResolved []osVersionWithResolvedType
 
 	selectStmt := `
 		SELECT os.os_version_id, osv.resolved_in_version
@@ -138,8 +139,27 @@ func (ds *Datastore) OSVersionsByCVE(ctx context.Context, cve string, teamID *ui
 		return vos, updatedAt, ctxerr.Wrap(ctx, err, "fetching OS version and resolved version by CVE")
 	}
 
+	// Remove duplicates, which may occur since the same OS can be installed on multiple architectures (amd64, arm64, etc.)
+	type osVersionKey struct {
+		OSVersionID     uint
+		ResolvedVersion string
+	}
+	seen := make(map[osVersionKey]struct{}, len(osVersionWithResolved))
+	verResolvedDedup := make([]osVersionWithResolvedType, 0, len(osVersionWithResolved))
+	for _, id := range osVersionWithResolved {
+		var resolved string
+		if id.ResolvedVersion != nil {
+			resolved = *id.ResolvedVersion
+		}
+		key := osVersionKey{OSVersionID: id.OSVersionID, ResolvedVersion: resolved}
+		if _, ok := seen[key]; !ok {
+			verResolvedDedup = append(verResolvedDedup, id)
+			seen[key] = struct{}{}
+		}
+	}
+
 	for _, osv := range osvs.OSVersions {
-		for _, id := range osVersionWithResolved {
+		for _, id := range verResolvedDedup {
 			if osv.OSVersionID == id.OSVersionID {
 				vos = append(vos, &fleet.VulnerableOS{
 					OSVersion:         osv,
diff --git a/server/datastore/mysql/vulnerabilities_test.go b/server/datastore/mysql/vulnerabilities_test.go
index fca02a6bf7..e4571b28a7 100644
--- a/server/datastore/mysql/vulnerabilities_test.go
+++ b/server/datastore/mysql/vulnerabilities_test.go
@@ -892,10 +892,14 @@ func seedVulnerabilities(t *testing.T, ds *Datastore) {
 
 	// update 15 hosts to windows
 	for i := 0; i < 10; i++ {
+		arch := "arm64"
+		if i%2 == 0 {
+			arch = "x86_64"
+		}
 		err := ds.UpdateHostOperatingSystem(context.Background(), hostids[i], fleet.OperatingSystem{
 			Name:     "Microsoft Windows 11 Enterprise 22H2",
 			Version:  "10.0.22621.2715",
-			Arch:     "x86_64",
+			Arch:     arch,
 			Platform: "windows",
 		})
 		require.NoError(t, err)
@@ -903,10 +907,14 @@ func seedVulnerabilities(t *testing.T, ds *Datastore) {
 
 	// update 5 hosts to macOS
 	for i := 10; i < 15; i++ {
+		arch := "arm64"
+		if i%2 == 0 {
+			arch = "x86_64"
+		}
 		err := ds.UpdateHostOperatingSystem(context.Background(), hostids[i], fleet.OperatingSystem{
 			Name:     "macOS",
 			Version:  "14.1.2",
-			Arch:     "arm64",
+			Arch:     arch,
 			Platform: "darwin",
 		})
 		require.NoError(t, err)
@@ -988,7 +996,7 @@ func seedVulnerabilities(t *testing.T, ds *Datastore) {
 
 	osVulns := []fleet.OSVulnerability{
 		{
-			OSID:              1,
+			OSID:              1, // windows x86_64
 			CVE:               "CVE-2020-1238",
 			ResolvedInVersion: ptr.String("1.0.0"),
 		},
@@ -998,13 +1006,31 @@ func seedVulnerabilities(t *testing.T, ds *Datastore) {
 			ResolvedInVersion: ptr.String("1.0.1"),
 		},
 		{
-			OSID: 2,
+			OSID:              2, // windows arm64
+			CVE:               "CVE-2020-1238",
+			ResolvedInVersion: ptr.String("1.0.0"),
+		},
+		{
+			OSID:              2,
+			CVE:               "CVE-2020-1239",
+			ResolvedInVersion: ptr.String("1.0.1"),
+		},
+		{
+			OSID: 2, // macos x86_64
 			CVE:  "CVE-2020-1240",
 		},
 		{
 			OSID: 2,
 			CVE:  "CVE-2020-1241",
 		},
+		{
+			OSID: 3, // macos arm64
+			CVE:  "CVE-2020-1240",
+		},
+		{
+			OSID: 3,
+			CVE:  "CVE-2020-1241",
+		},
 	}
 
 	mockTime := time.Date(2020, 1, 1, 0, 0, 0, 0, time.UTC)