Add the concept of a decorator. Add a vulnerability decorator

src/Common/Common.Test/Common.Test/PackageDependencyGraphTests.cs

@@ -13,7 +13,7 @@ public async Task FromAssetsFile_WithLargeGraph_ParsesGraphCorrectly()
             var assetsFileText = TestHelpers.GetResource("Common.Test.compiler.resources.nuget.common.assets.json", GetType());
 
             var assetsFile = new LockFileFormat().Parse(assetsFileText, Path.GetTempPath());
-            var graphs = await PackageDependencyGraph.GenerateAllDependencyGraphsFromAssetsFileAsync(assetsFile, new GraphOptions(checkVulnerabilities: false, generateProjectsOnly: false));
+            var graphs = await PackageDependencyGraph.GenerateAllDependencyGraphsFromAssetsFileAsync(assetsFile, new GraphOptions(generateProjectsOnly: false));
             graphs.Should().HaveCount(2);
 
             var graph = graphs.First().Value;
@@ -296,7 +296,7 @@ public async Task GenerateAllDependencyGraphsFromAssetsFile_WithMultipleFramewor
             var assetsFileText = TestHelpers.GetResource("Common.Test.compiler.resources.multitargeted.assets.json", GetType());
 
             var assetsFile = new LockFileFormat().Parse(assetsFileText, Path.GetTempPath());
-            Dictionary<string, PackageDependencyGraph> graphs = await PackageDependencyGraph.GenerateAllDependencyGraphsFromAssetsFileAsync(assetsFile, new GraphOptions(checkVulnerabilities: false, generateProjectsOnly: false));
+            Dictionary<string, PackageDependencyGraph> graphs = await PackageDependencyGraph.GenerateAllDependencyGraphsFromAssetsFileAsync(assetsFile, new GraphOptions(generateProjectsOnly: false));
 
             var net472Graph = graphs["net472"];
             net472Graph.Node.Identity.Id.Should().Be("TestProject");
@@ -330,7 +330,7 @@ public async Task FromAssetsFile_WithProjectReferenceWithAPackageId_ParsesGraphC
             var assetsFile = new LockFileFormat().Parse(assetsFileText, Path.GetTempPath());
             var dgSpec = DependencyGraphSpec.Load(tempFile.FilePath);
 
-            var graphs = await PackageDependencyGraph.GenerateAllDependencyGraphsFromAssetsFileAsync(assetsFile, dgSpec, new GraphOptions(checkVulnerabilities: false, generateProjectsOnly: false));
+            var graphs = await PackageDependencyGraph.GenerateAllDependencyGraphsFromAssetsFileAsync(assetsFile, dgSpec, new GraphOptions(generateProjectsOnly: false), new(), CancellationToken.None);
             graphs.Should().HaveCount(1);
             var graph = graphs.Single().Value;
 
@@ -357,7 +357,7 @@ private PackageDependencyGraph GetOnlyDependencyGraph(string resourceName)
             var assetsFileText = TestHelpers.GetResource(resourceName, GetType());
 
             var assetsFile = new LockFileFormat().Parse(assetsFileText, Path.GetTempPath());
-            var graphs = PackageDependencyGraph.GenerateAllDependencyGraphsFromAssetsFileAsync(assetsFile, new GraphOptions(checkVulnerabilities: false, generateProjectsOnly: false)).Result;
+            var graphs = PackageDependencyGraph.GenerateAllDependencyGraphsFromAssetsFileAsync(assetsFile, new GraphOptions(generateProjectsOnly: false)).Result;
             graphs.Should().HaveCount(1);
             var graph = graphs.Single().Value;
             return graph;

src/Common/GraphOptions.cs

@@ -2,13 +2,11 @@
 {
     public class GraphOptions
     {
-        public GraphOptions(bool checkVulnerabilities, bool generateProjectsOnly)
+        public GraphOptions(bool generateProjectsOnly)
         {
-            CheckVulnerabilities = checkVulnerabilities;
             GenerateProjectsOnly = generateProjectsOnly;
         }
 
-        public bool CheckVulnerabilities { get; }
         public bool GenerateProjectsOnly { get; }
     }
 }

src/Common/IPackageDependencyNodeDecorator.cs

@@ -0,0 +1,7 @@
+namespace Common
+{
+    public interface IPackageDependencyNodeDecorator
+    {
+        Task DecorateAsync(PackageDependencyNode dependencyNode, CancellationToken cancellationToken);
+    }
+}

src/Common/PackageDependencyGraph.cs

@@ -1,10 +1,8 @@
-using NuGet.Configuration;
-using NuGet.Packaging.Core;
+using NuGet.Packaging.Core;
 using NuGet.ProjectModel;
-using NuGet.Protocol;
-using NuGet.Protocol.Core.Types;
 using NuGet.Versioning;
 using System.Diagnostics;
+using System.Threading;
 
 namespace Common
 {
@@ -22,11 +20,19 @@ public PackageDependencyGraph(Node<DependencyNodeIdentity, VersionRange> node) :
         /// <returns></returns>
         /// <exception cref="InvalidOperationException">If the assets file is not valid</exception>
         /// <exception cref="ArgumentNullException">If the assets file is null</exception>
-        public static async Task<Dictionary<string, PackageDependencyGraph>> GenerateAllDependencyGraphsFromAssetsFileAsync(LockFile assetsFile, DependencyGraphSpec dependencyGraphSpec, GraphOptions graphOptions)
+        public static async Task<Dictionary<string, PackageDependencyGraph>> GenerateAllDependencyGraphsFromAssetsFileAsync(
+            LockFile assetsFile,
+            DependencyGraphSpec dependencyGraphSpec,
+            GraphOptions graphOptions,
+            List<IPackageDependencyNodeDecorator> decorators,
+            CancellationToken cancellationToken)
         {
             ArgumentNullException.ThrowIfNull(assetsFile);
-            DependencyNodeIdentity projectIdentity = new(assetsFile.PackageSpec.Name, assetsFile.PackageSpec.Version, DependencyType.Project);
+            ArgumentNullException.ThrowIfNull(dependencyGraphSpec);
+            ArgumentNullException.ThrowIfNull(graphOptions);
+            ArgumentNullException.ThrowIfNull(decorators);
 
+            DependencyNodeIdentity projectIdentity = new(assetsFile.PackageSpec.Name, assetsFile.PackageSpec.Version, DependencyType.Project);
             List<LockFileTarget> frameworks = assetsFile.Targets.Where(e => string.IsNullOrEmpty(e.RuntimeIdentifier)).ToList();
 
             if (frameworks.Count == 0)
@@ -46,7 +52,7 @@ public static async Task<Dictionary<string, PackageDependencyGraph>> GenerateAll
 
             foreach (var framework in frameworks)
             {
-                PackageDependencyGraph dependencyGraph = await GenerateGraphForAGivenFramework(projectIdentity, framework, assetsFile.PackageSpec, projectPathToProjectNameMap, graphOptions.CheckVulnerabilities, graphOptions.GenerateProjectsOnly);
+                PackageDependencyGraph dependencyGraph = await GenerateGraphForAGivenFramework(projectIdentity, framework, assetsFile.PackageSpec, projectPathToProjectNameMap, graphOptions.GenerateProjectsOnly, decorators, cancellationToken);
                 TargetFrameworkInformation alias = assetsFile.PackageSpec.GetTargetFramework(framework.TargetFramework);
                 aliasToDependencyGraph.Add(alias.TargetAlias, dependencyGraph);
             }
@@ -77,7 +83,7 @@ public static async Task<Dictionary<string, PackageDependencyGraph>> GenerateAll
 
             foreach (var framework in frameworks)
             {
-                var dependenyGraph = await GenerateGraphForAGivenFramework(projectIdentity, framework, assetsFile.PackageSpec, new(), graphOptions.CheckVulnerabilities, projectsOnly: false);
+                var dependenyGraph = await GenerateGraphForAGivenFramework(projectIdentity, framework, assetsFile.PackageSpec, new(), projectsOnly: false, new(), CancellationToken.None);
                 var alias = assetsFile.PackageSpec.GetTargetFramework(framework.TargetFramework);
                 aliasToDependencyGraph.Add(alias.TargetAlias, dependenyGraph);
             }
@@ -90,8 +96,9 @@ private static async Task<PackageDependencyGraph> GenerateGraphForAGivenFramewor
             LockFileTarget framework,
             PackageSpec packageSpec,
             Dictionary<string, string> projectPathToProjectNameMap,
-            bool checkVulnerabilities,
-            bool projectsOnly)
+            bool projectsOnly,
+            List<IPackageDependencyNodeDecorator> decorators,
+            CancellationToken cancellationToken)
         {
             ArgumentNullException.ThrowIfNull(projectIdentity);
             ArgumentNullException.ThrowIfNull(framework);
@@ -101,11 +108,6 @@ private static async Task<PackageDependencyGraph> GenerateGraphForAGivenFramewor
 
             Dictionary<string, PackageDependencyNode> packageIdToNode = GenerateNodesForAllPackagesInGraph(framework, projectsOnly);
 
-            if (checkVulnerabilities)
-            {
-                await VulnerabilityHelpers.GenerateVulnerabilitiesAsync(packageSpec, packageIdToNode, CancellationToken.None);
-            }
-
             packageIdToNode.Add(graph.Node.Identity.Id, (PackageDependencyNode)graph.Node);
 
             // Populate Node to Node edges
@@ -156,6 +158,14 @@ private static async Task<PackageDependencyGraph> GenerateGraphForAGivenFramewor
                 node.ParentNodes.Add((graph.Node, versionRange));
             }
 
+            foreach ((string _, PackageDependencyNode packageIdNode) in packageIdToNode)
+            {
+                foreach (var decorator in decorators)
+                {
+                    await decorator.DecorateAsync(packageIdNode, cancellationToken);
+                }
+            }
+
             return graph;
 
             static Dictionary<string, PackageDependencyNode> GenerateNodesForAllPackagesInGraph(LockFileTarget framework, bool projectsOnly)

src/Common/VulnerabilityChecker.cs → src/Common/VulnerabilityInfoDecorator.cs

@@ -1,32 +1,53 @@
-using NuGet.Protocol.Core.Types;
-using NuGet.Protocol.Model;
+using NuGet.Packaging.Core;
+using NuGet.Protocol.Core.Types;
 using NuGet.Protocol;
-using NuGet.Common;
-using Logging;
+using NuGet.Protocol.Model;
 using Microsoft.Extensions.Logging;
-using NuGet.Packaging.Core;
+using Logging;
 
 namespace Common
 {
-    internal class VulnerabilityChecker
+    public class VulnerabilityInfoDecorator : IPackageDependencyNodeDecorator
     {
-        private readonly Microsoft.Extensions.Logging.ILogger _logger;
+        private readonly ILogger _logger;
         private readonly List<SourceRepository> _sourceRepositories;
         private readonly SourceCacheContext _sourceCacheContext;
-        private IReadOnlyList<IReadOnlyDictionary<string, IReadOnlyList<PackageVulnerabilityInfo>>> _vulnerabilityData;
+        private IReadOnlyList<IReadOnlyDictionary<string, IReadOnlyList<PackageVulnerabilityInfo>>>? _vulnerabilityData;
         private bool _vulnerabilityDataChecked;
 
-        public VulnerabilityChecker(List<SourceRepository> sourceRepositories, SourceCacheContext sourceCacheContext)
+        public VulnerabilityInfoDecorator(List<SourceRepository> sourceRepositories, SourceCacheContext sourceCacheContext)
         {
-            _logger = AppLogger.Logger;
             _sourceRepositories = sourceRepositories ?? throw new ArgumentNullException(nameof(sourceRepositories));
             _sourceCacheContext = sourceCacheContext ?? throw new ArgumentNullException(nameof(sourceCacheContext));
+            _logger = AppLogger.Logger;
+        }
+
+        public async Task DecorateAsync(PackageDependencyNode dependencyNode, CancellationToken cancellationToken)
+        {
+            if (!_vulnerabilityDataChecked)
+            {
+                IReadOnlyList<IReadOnlyDictionary<string, IReadOnlyList<PackageVulnerabilityInfo>>>? loadedVulnerabilityData = await LoadVulnerabilitiesAsync(_sourceRepositories, _sourceCacheContext, cancellationToken);
+                if (loadedVulnerabilityData is not null)
+                {
+                    _vulnerabilityData = loadedVulnerabilityData;
+                }
+                _vulnerabilityDataChecked = true;
+            }
+
+            if (_vulnerabilityData != null)
+            {
+                var packageIdentity = (PackageIdentity)dependencyNode.Identity;
+                if (IsPackageVulnerable(_vulnerabilityData, packageIdentity))
+                {
+                    dependencyNode.Identity.Vulnerable = true;
+                }
+            }
         }
 
         private async Task<IReadOnlyList<IReadOnlyDictionary<string, IReadOnlyList<PackageVulnerabilityInfo>>>?> LoadVulnerabilitiesAsync(
-            List<SourceRepository> sourceRepositories,
-            SourceCacheContext sourceCacheContext,
-            CancellationToken cancellationToken)
+        List<SourceRepository> sourceRepositories,
+        SourceCacheContext sourceCacheContext,
+        CancellationToken cancellationToken)
         {
             GetVulnerabilityInfoResult? allVulnerabilityData = await GetAllVulnerabilityDataAsync(sourceRepositories, sourceCacheContext, cancellationToken);
 
@@ -58,26 +79,11 @@ static bool IsAnyVulnerabilityDataFound(IReadOnlyList<IReadOnlyDictionary<string
             }
         }
 
-
-        public async Task<bool> IsPackageVulnerableAsync(
-            PackageIdentity packageIdentity,
-            CancellationToken cancellationToken)
+        internal static bool IsPackageVulnerable(
+            IReadOnlyList<IReadOnlyDictionary<string, IReadOnlyList<PackageVulnerabilityInfo>>> vulnerabilityData,
+            PackageIdentity packageIdentity)
         {
-            if (!_vulnerabilityDataChecked)
-            {
-                IReadOnlyList<IReadOnlyDictionary<string, IReadOnlyList<PackageVulnerabilityInfo>>>? loadedVulnerabilityData = await LoadVulnerabilitiesAsync(_sourceRepositories, _sourceCacheContext, cancellationToken);
-                if (loadedVulnerabilityData is not null)
-                {
-                    _vulnerabilityData = loadedVulnerabilityData;
-                }
-                _vulnerabilityDataChecked = true;
-            }
-
-            if (_vulnerabilityData == null)
-            {
-                return false;
-            }
-            foreach (var singleDataFile in _vulnerabilityData)
+            foreach (var singleDataFile in vulnerabilityData)
             {
                 if (singleDataFile.TryGetValue(packageIdentity.Id, out var vulnerabilityList))
                 {
@@ -157,7 +163,7 @@ public async Task<bool> IsPackageVulnerableAsync(
                 {
                     return null;
                 }
-                return await vulnerabilityInfoResource.GetVulnerabilityInfoAsync(cacheContext, NullLogger.Instance, CancellationToken.None);
+                return await vulnerabilityInfoResource.GetVulnerabilityInfoAsync(cacheContext, NuGet.Common.NullLogger.Instance, CancellationToken.None);
             }
         }
     }

src/DependencyVisualizerTool/DependencyVisualizerTool.Test/PackageDependencyVisualizerToolTests.cs

@@ -62,7 +62,7 @@ public async Task TransGraphToDGMLXDocument_multitargeted_CreateDGMLCorrectly()
             var assetsFileText = TestHelpers.GetResource("DependencyVisualizerTool.Test.compiler.resources.multitargeted.assets.json", GetType());
 
             var assetsFile = new LockFileFormat().Parse(assetsFileText, Path.GetTempPath());
-            var graphs = await PackageDependencyGraph.GenerateAllDependencyGraphsFromAssetsFileAsync(assetsFile, new GraphOptions(checkVulnerabilities: false, generateProjectsOnly: false));
+            var graphs = await PackageDependencyGraph.GenerateAllDependencyGraphsFromAssetsFileAsync(assetsFile, new GraphOptions(generateProjectsOnly: false));
             graphs.Should().HaveCount(2);
             var graph = graphs.First().Value;
 
@@ -79,7 +79,7 @@ public async Task TransGraphToDGMLXDocument_nugetcommon_CreateDGMLCorrectly()
             var assetsFileText = TestHelpers.GetResource("DependencyVisualizerTool.Test.compiler.resources.nuget.common.assets.json", GetType());
 
             var assetsFile = new LockFileFormat().Parse(assetsFileText, Path.GetTempPath());
-            var graphs = await PackageDependencyGraph.GenerateAllDependencyGraphsFromAssetsFileAsync(assetsFile, new GraphOptions(checkVulnerabilities: false, generateProjectsOnly: false));
+            var graphs = await PackageDependencyGraph.GenerateAllDependencyGraphsFromAssetsFileAsync(assetsFile, new GraphOptions(generateProjectsOnly: false));
             graphs.Should().HaveCount(2);
             var graph = graphs.First().Value;
 
@@ -149,7 +149,7 @@ private async Task<PackageDependencyGraph> GetOnlyDependencyGraphAsync(string re
             var assetsFileText = TestHelpers.GetResource(resourceName, GetType());
 
             var assetsFile = new LockFileFormat().Parse(assetsFileText, Path.GetTempPath());
-            var graphs = await PackageDependencyGraph.GenerateAllDependencyGraphsFromAssetsFileAsync(assetsFile, new GraphOptions(checkVulnerabilities: false, generateProjectsOnly: false));
+            var graphs = await PackageDependencyGraph.GenerateAllDependencyGraphsFromAssetsFileAsync(assetsFile, new GraphOptions(generateProjectsOnly: false));
             graphs.Should().HaveCount(1);
             var graph = graphs.Single().Value;
             return graph;