change sql get vulnerabilities by package id (#8540)

* change sql get vulunerabilities by package id * Adjusted logic to provide correct groupings, adjusted UTs Co-authored-by: Drew Gillies <drewgil@microsoft.com>
NuGet · Apr 23, 2021 · a2ab1c7 · a2ab1c7
1 parent 409f844
commit a2ab1c7
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 41 deletions.
diff --git a/src/NuGetGallery/Services/PackageVulnerabilitiesService.cs b/src/NuGetGallery/Services/PackageVulnerabilitiesService.cs
@@ -3,8 +3,8 @@
 
 using System;
 using System.Collections.Generic;
+using System.Collections.ObjectModel;
 using System.Data.Entity;
-using System.IO;
 using System.Linq;
 using NuGet.Services.Entities;
 
@@ -21,26 +21,15 @@ public PackageVulnerabilitiesService(IEntitiesContext entitiesContext)
 
         public IReadOnlyDictionary<int, IReadOnlyList<PackageVulnerability>> GetVulnerabilitiesById(string id)
         {
-            var result = new Dictionary<int, List<PackageVulnerability>>();
-            var packagesMatchingId = _entitiesContext.Packages
-                .Where(p => p.PackageRegistration != null && p.PackageRegistration.Id == id)
-                .Include($"{nameof(Package.VulnerablePackageRanges)}.{nameof(VulnerablePackageVersionRange.Vulnerability)}");
-            foreach (var package in packagesMatchingId)
-            {
-                if (package.VulnerablePackageRanges == null)
-                {
-                    continue;
-                }
-
-                if (package.VulnerablePackageRanges.Any())
-                {
-                    result.Add(package.Key,
-                        package.VulnerablePackageRanges.Select(vr => vr.Vulnerability).ToList());
-                }
-            }
-
-            return !result.Any() ? null :
-                result.ToDictionary(kv => kv.Key, kv => kv.Value as IReadOnlyList<PackageVulnerability>);
+            var packageKeyAndVulnerability = _entitiesContext.VulnerableRanges
+                .Include(x => x.Vulnerability)
+                .Where(x => x.PackageId == id)
+                .SelectMany(x => x.Packages.Select(p => new {PackageKey = p.Key, x.Vulnerability}))
+                .GroupBy(pv => pv.PackageKey, pv => pv.Vulnerability)
+                .ToDictionary(pv => pv.Key, pv => pv.ToList().AsReadOnly() as IReadOnlyList<PackageVulnerability>);
+
+            return !packageKeyAndVulnerability.Any() ? null
+                : new ReadOnlyDictionary<int, IReadOnlyList<PackageVulnerability>>(packageKeyAndVulnerability);
         }
 
         public bool IsPackageVulnerable(Package package)

diff --git a/tests/NuGetGallery.Facts/Services/PackageVulnerabilitiesServiceFacts.cs b/tests/NuGetGallery.Facts/Services/PackageVulnerabilitiesServiceFacts.cs
@@ -2,7 +2,6 @@
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System.Collections.Generic;
-using System.Linq;
 using NuGet.Services.Entities;
 using NuGetGallery.Framework;
 using Xunit;
@@ -22,7 +21,6 @@ public class PackageVulnerabilitiesServiceFacts : TestContainer
         private Package _packageVulnerable100;
         private Package _packageVulnerable110;
         private Package _packageVulnerable111;
-        private Package _packageVulnerable112;
 
         private Package _packageNotVulnerable;
 
@@ -31,16 +29,9 @@ public void GetsVulnerabilitiesOfPackage()
         {
             // Arrange
             SetUp();
-            var packages = new[]
-            {
-                _packageVulnerable100,
-                _packageVulnerable110,
-                _packageVulnerable111,
-                _packageVulnerable112,
-                _packageNotVulnerable
-            };
+            var vulnerableRanges = new[] {_versionRangeModerate, _versionRangeCritical};
             var context = GetFakeContext();
-            context.Packages.AddRange(packages);
+            context.VulnerableRanges.AddRange(vulnerableRanges);
             var target = Get<PackageVulnerabilitiesService>();
 
             // Act
@@ -96,12 +87,14 @@ private void SetUp()
 
             _versionRangeCritical = new VulnerablePackageVersionRange
             {
+                PackageId = "Vulnerable",
                 Vulnerability = _vulnerabilityCritical,
                 PackageVersionRange = "1.1.1",
                 FirstPatchedPackageVersion = "1.1.2"
             };
             _versionRangeModerate = new VulnerablePackageVersionRange
             {
+                PackageId = "Vulnerable",
                 Vulnerability = _vulnerabilityModerate,
                 PackageVersionRange = "<=1.1.1",
                 FirstPatchedPackageVersion = "1.1.2"
@@ -129,7 +122,7 @@ private void SetUp()
             };
             _packageVulnerable111 = new Package
             {
-                Key = 3, // simulate a different order in db - create a non-contiguous range of rows, even if the range is contiguous
+                Key = 2,
                 PackageRegistration = _registrationVulnerable,
                 Version = "1.1.1",
                 VulnerablePackageRanges = new List<VulnerablePackageVersionRange>
@@ -138,19 +131,15 @@ private void SetUp()
                     _versionRangeCritical
                 }
             };
-            _packageVulnerable112 = new Package
-            {
-                Key = 2, // simulate a different order in db  - create a non-contiguous range of rows, even if the range is contiguous
-                PackageRegistration = _registrationVulnerable,
-                Version = "1.1.2",
-                VulnerablePackageRanges = new List<VulnerablePackageVersionRange>()
-            };
             _packageNotVulnerable = new Package
             {
-                Key = 4,
+                Key = 3,
                 PackageRegistration = new PackageRegistration { Id = "NotVulnerable" },
                 VulnerablePackageRanges = new List<VulnerablePackageVersionRange>()
             };
+
+            _versionRangeCritical.Packages = new List<Package> { _packageVulnerable111 };
+            _versionRangeModerate.Packages = new List<Package> { _packageVulnerable100, _packageVulnerable110, _packageVulnerable111 };
         }
     }
 }