From 9ede32ec85193b4c07b64578e10a18bc66ac124b Mon Sep 17 00:00:00 2001
From: Bastian Doetsch <bastian.doetsch@snyk.io>
Date: Sun, 13 Oct 2024 09:04:21 -0400
Subject: [PATCH] fix: update root node presentation (#623)

---
 .../SnykToolWindowSnykScanListenerLS.kt       | 64 +++++++++++++++----
 .../snyk/common/lsp/LanguageServerWrapper.kt  |  5 +-
 2 files changed, 54 insertions(+), 15 deletions(-)

diff --git a/src/main/kotlin/io/snyk/plugin/ui/toolwindow/SnykToolWindowSnykScanListenerLS.kt b/src/main/kotlin/io/snyk/plugin/ui/toolwindow/SnykToolWindowSnykScanListenerLS.kt
index 0534239b..517c0a6c 100644
--- a/src/main/kotlin/io/snyk/plugin/ui/toolwindow/SnykToolWindowSnykScanListenerLS.kt
+++ b/src/main/kotlin/io/snyk/plugin/ui/toolwindow/SnykToolWindowSnykScanListenerLS.kt
@@ -169,18 +169,46 @@ class SnykToolWindowSnykScanListenerLS(
         )
     }
 
-    private fun displayResults(snykResults: Map<SnykFile, List<ScanIssue>>, enabledInSettings: Boolean, filterTree: Boolean, rootNode: DefaultMutableTreeNode) {
+    private fun displayResults(
+        snykResults: Map<SnykFile, List<ScanIssue>>,
+        enabledInSettings: Boolean,
+        filterTree: Boolean,
+        rootNode: DefaultMutableTreeNode,
+        issueType: String
+    ) {
         if (disposed) return
         if (getSnykCachedResults(project)?.currentIacError != null) return
 
-        displayIssues(
-            enabledInSettings = enabledInSettings,
-            filterTree = filterTree,
-            snykResults = snykResults,
-            rootNode = rootNode,
-            iacResultsCount = snykResults.values.flatten().distinct().size,
-            fixableIssuesCount = snykResults.values.flatten().count { it.additionalData.isUpgradable }
-        )
+        val flattenedResults = snykResults.values.flatten()
+
+        when (issueType) {
+            ScanIssue.OPEN_SOURCE -> {
+                val ossResultsCount =
+                    flattenedResults.filter { it.filterableIssueType == ScanIssue.OPEN_SOURCE }.distinct().size
+                displayIssues(
+                    enabledInSettings = enabledInSettings,
+                    filterTree = filterTree,
+                    snykResults = snykResults,
+                    rootNode = rootNode,
+                    ossResultsCount = ossResultsCount,
+                    fixableIssuesCount = flattenedResults.count { it.additionalData.isUpgradable }
+                )
+            }
+
+            ScanIssue.INFRASTRUCTURE_AS_CODE -> {
+                val iacResultsCount =
+                    flattenedResults.filter { it.filterableIssueType == ScanIssue.INFRASTRUCTURE_AS_CODE }
+                        .distinct().size
+                displayIssues(
+                    enabledInSettings = enabledInSettings,
+                    filterTree = filterTree,
+                    snykResults = snykResults,
+                    rootNode = rootNode,
+                    iacResultsCount = iacResultsCount,
+                    fixableIssuesCount = flattenedResults.count { it.additionalData.isUpgradable }
+                )
+            }
+        }
     }
 
     fun displayOssResults(snykResults: Map<SnykFile, List<ScanIssue>>) {
@@ -189,7 +217,13 @@ class SnykToolWindowSnykScanListenerLS(
 
         val settings = pluginSettings()
 
-        displayResults(snykResults, settings.ossScanEnable, settings.treeFiltering.ossResults, this.rootOssIssuesTreeNode)
+        displayResults(
+            snykResults,
+            settings.ossScanEnable,
+            settings.treeFiltering.ossResults,
+            this.rootOssIssuesTreeNode,
+            ScanIssue.OPEN_SOURCE
+        )
     }
 
     fun displayIacResults(snykResults: Map<SnykFile, List<ScanIssue>>) {
@@ -197,7 +231,13 @@ class SnykToolWindowSnykScanListenerLS(
         if (getSnykCachedResults(project)?.currentIacError != null) return
 
         val settings = pluginSettings()
-        displayResults(snykResults, settings.iacScanEnabled, settings.treeFiltering.iacResults, this.rootIacIssuesTreeNode)
+        displayResults(
+            snykResults,
+            settings.iacScanEnabled,
+            settings.treeFiltering.iacResults,
+            this.rootIacIssuesTreeNode,
+            ScanIssue.INFRASTRUCTURE_AS_CODE,
+        )
     }
 
     private fun displayIssues(
@@ -305,7 +345,7 @@ class SnykToolWindowSnykScanListenerLS(
             } else {
                 "ies"
             }
-                text = "✋ $issuesCount vulnerabilit$plural found by Snyk"
+            text = "✋ $issuesCount vulnerabilit$plural found by Snyk"
             if (pluginSettings().isGlobalIgnoresFeatureEnabled) {
                 text += ", $ignoredIssuesCount ignored"
             }
diff --git a/src/main/kotlin/snyk/common/lsp/LanguageServerWrapper.kt b/src/main/kotlin/snyk/common/lsp/LanguageServerWrapper.kt
index 0a326b76..fa4e235c 100644
--- a/src/main/kotlin/snyk/common/lsp/LanguageServerWrapper.kt
+++ b/src/main/kotlin/snyk/common/lsp/LanguageServerWrapper.kt
@@ -348,7 +348,6 @@ class LanguageServerWrapper(
     fun sendScanCommand(project: Project) {
         if (notAuthenticated()) return
         DumbService.getInstance(project).runWhenSmart {
-            refreshFeatureFlags()
             getTrustedContentRoots(project).forEach {
                 sendFolderScanCommand(it.path, project)
             }
@@ -356,11 +355,11 @@ class LanguageServerWrapper(
     }
 
     fun refreshFeatureFlags() {
-        runInBackground("Snyk: refreshing feature flags...") {
+        runAsync {
             // this check should be async, as refresh is called from initialization and notAuthenticated is triggering
             // initialization. So, to make it wait patiently for its turn, it needs to be checked and executed in a
             // different thread.
-            if (notAuthenticated()) return@runInBackground
+            if (notAuthenticated()) return@runAsync
             pluginSettings().isGlobalIgnoresFeatureEnabled = getFeatureFlagStatus("snykCodeConsistentIgnores")
         }
     }