Amazon Q : Fixing the inappropriate yellow lines shown in the editor …

…file using CodeSnippet. (#4618) * Dropping issue by checking the codeSnippet from listCodeAnalysis API
aws · Jul 6, 2024 · eccdf21 · eccdf21
1 parent 937e9b4
commit eccdf21
Show file tree

Hide file tree

Showing 5 changed files with 138 additions and 40 deletions.
diff --git a/.changes/next-release/bugfix-3da73590-7103-4720-8c6f-f13223d478cf.json b/.changes/next-release/bugfix-3da73590-7103-4720-8c6f-f13223d478cf.json
@@ -0,0 +1,4 @@
+{
+  "type" : "bugfix",
+  "description" : "Amazon Q Security Scans: Fixed unnecessary yellow lines appearing in both auto scans and project scans."
+}
diff --git a/...re/aws/toolkits/jetbrains/services/codewhisperer/codescan/CodeWhispererCodeScanManager.kt b/...re/aws/toolkits/jetbrains/services/codewhisperer/codescan/CodeWhispererCodeScanManager.kt
@@ -739,6 +739,7 @@ data class CodeWhispererCodeScanIssue(
     val severity: String,
     val recommendation: Recommendation,
     var suggestedFixes: List<SuggestedFix>,
+    val codeSnippet: List<CodeLine>,
     val issueSeverity: HighlightDisplayLevel = HighlightDisplayLevel.WARNING,
     val isInvalid: Boolean = false,
     var rangeHighlighter: RangeHighlighterEx? = null

diff --git a/...re/aws/toolkits/jetbrains/services/codewhisperer/codescan/CodeWhispererCodeScanSession.kt b/...re/aws/toolkits/jetbrains/services/codewhisperer/codescan/CodeWhispererCodeScanSession.kt
@@ -69,8 +69,6 @@ import java.time.Instant
 import java.util.Base64
 import java.util.UUID
 import kotlin.coroutines.coroutineContext
-import kotlin.math.max
-import kotlin.math.min
 
 class CodeWhispererCodeScanSession(val sessionContext: CodeScanSessionContext) {
     private val clientToken: UUID = UUID.randomUUID()
@@ -375,50 +373,69 @@ class CodeWhispererCodeScanSession(val sessionContext: CodeScanSessionContext) {
     }
 
     fun mapToCodeScanIssues(recommendations: List<String>): List<CodeWhispererCodeScanIssue> {
-        val scanRecommendations: List<CodeScanRecommendation> = recommendations.map {
-            val value: List<CodeScanRecommendation> = MAPPER.readValue(it)
-            value
-        }.flatten()
+        val scanRecommendations = recommendations.flatMap { MAPPER.readValue<List<CodeScanRecommendation>>(it) }
         if (isProjectScope()) {
             LOG.debug { "Total code scan issues returned from service: ${scanRecommendations.size}" }
         }
-        return scanRecommendations.mapNotNull {
+        return scanRecommendations.mapNotNull { recommendation ->
             val file = try {
                 LocalFileSystem.getInstance().findFileByIoFile(
-                    Path.of(sessionContext.sessionConfig.projectRoot.path, it.filePath).toFile()
+                    Path.of(sessionContext.sessionConfig.projectRoot.path, recommendation.filePath).toFile()
                 )
             } catch (e: Exception) {
-                LOG.debug { "Cannot find file at location ${it.filePath}" }
+                LOG.debug { "Cannot find file at location ${recommendation.filePath}" }
                 null
             }
-            when (file?.isDirectory) {
-                false -> {
-                    runReadAction {
-                        FileDocumentManager.getInstance().getDocument(file)
-                    }?.let { document ->
-                        val endLineInDocument = min(max(0, it.endLine - 1), document.lineCount - 1)
+
+            if (file?.isDirectory == false) {
+                runReadAction {
+                    FileDocumentManager.getInstance().getDocument(file)
+                }?.let { document ->
+
+                    val documentLines = document.getText().split("\n")
+                    val (startLine, endLine) = recommendation.run { startLine to endLine }
+                    var shouldDisplayIssue = true
+
+                    for (codeBlock in recommendation.codeSnippet) {
+                        val lineNumber = codeBlock.number - 1
+                        if (codeBlock.number in startLine..endLine) {
+                            val documentLine = documentLines.getOrNull(lineNumber)
+                            if (documentLine != codeBlock.content) {
+                                shouldDisplayIssue = false
+                                break
+                            }
+                        }
+                    }
+
+                    if (shouldDisplayIssue) {
+                        val endLineInDocument = minOf(maxOf(0, recommendation.endLine - 1), document.lineCount - 1)
                         val endCol = document.getLineEndOffset(endLineInDocument) - document.getLineStartOffset(endLineInDocument) + 1
+
                         CodeWhispererCodeScanIssue(
-                            startLine = it.startLine,
+                            startLine = recommendation.startLine,
                             startCol = 1,
-                            endLine = it.endLine,
+                            endLine = recommendation.endLine,
                             endCol = endCol,
                             file = file,
                             project = sessionContext.project,
-                            title = it.title,
-                            description = it.description,
-                            detectorId = it.detectorId,
-                            detectorName = it.detectorName,
-                            findingId = it.findingId,
-                            ruleId = it.ruleId,
-                            relatedVulnerabilities = it.relatedVulnerabilities,
-                            severity = it.severity,
-                            recommendation = it.remediation.recommendation,
-                            suggestedFixes = it.remediation.suggestedFixes
+                            title = recommendation.title,
+                            description = recommendation.description,
+                            detectorId = recommendation.detectorId,
+                            detectorName = recommendation.detectorName,
+                            findingId = recommendation.findingId,
+                            ruleId = recommendation.ruleId,
+                            relatedVulnerabilities = recommendation.relatedVulnerabilities,
+                            severity = recommendation.severity,
+                            recommendation = recommendation.remediation.recommendation,
+                            suggestedFixes = recommendation.remediation.suggestedFixes,
+                            codeSnippet = recommendation.codeSnippet
                         )
+                    } else {
+                        null
                     }
                 }
-                else -> null
+            } else {
+                null
             }
         }.onEach { issue ->
             // Add range highlighters for all the issues found.
@@ -487,7 +504,8 @@ internal data class CodeScanRecommendation(
     val ruleId: String?,
     val relatedVulnerabilities: List<String>,
     val severity: String,
-    val remediation: Remediation
+    val remediation: Remediation,
+    val codeSnippet: List<CodeLine>
 )
 
 data class Description(val text: String, val markdown: String)
@@ -498,6 +516,8 @@ data class Recommendation(val text: String, val url: String)
 
 data class SuggestedFix(val description: String, val code: String)
 
+data class CodeLine(val number: Int, val content: String)
+
 data class CodeScanSessionContext(
     val project: Project,
     val sessionConfig: CodeScanSessionConfig,

diff --git a/...e/aws/toolkits/jetbrains/services/codewhisperer/codescan/CodeWhispererCodeScanTestBase.kt b/...e/aws/toolkits/jetbrains/services/codewhisperer/codescan/CodeWhispererCodeScanTestBase.kt
@@ -77,6 +77,7 @@ open class CodeWhispererCodeScanTestBase(projectRule: CodeInsightTestFixtureRule
     internal lateinit var fakeCreateCodeScanResponseFailed: CreateCodeScanResponse
     internal lateinit var fakeCreateCodeScanResponsePending: CreateCodeScanResponse
     internal lateinit var fakeListCodeScanFindingsResponse: ListCodeScanFindingsResponse
+    internal lateinit var fakeListCodeScanFindingsResponseE2E: ListCodeScanFindingsResponse
     internal lateinit var fakeListCodeScanFindingsOutOfBoundsIndexResponse: ListCodeScanFindingsResponse
     internal lateinit var fakeGetCodeScanResponse: GetCodeScanResponse
     internal lateinit var fakeGetCodeScanResponsePending: GetCodeScanResponse
@@ -110,7 +111,22 @@ open class CodeWhispererCodeScanTestBase(projectRule: CodeInsightTestFixtureRule
         )
     }
 
-    private fun setupCodeScanFinding(filePath: Path, startLine: Int, endLine: Int) = """
+    private fun setupCodeScanFinding(
+        filePath: Path,
+        startLine: Int,
+        endLine: Int,
+        codeSnippets: List<Pair<Int, String>>
+    ): String {
+        val codeSnippetJson = codeSnippets.joinToString(",\n") { (number, content) ->
+            """
+        {
+            "number": $number,
+            "content": "$content"
+        }
+            """.trimIndent()
+        }
+
+        return """
         {
             "filePath": "${filePath.systemIndependentPath}",
             "startLine": $startLine,
@@ -124,6 +140,9 @@ open class CodeWhispererCodeScanTestBase(projectRule: CodeInsightTestFixtureRule
             "detectorName": "detectorName",
             "findingId": "findingId",
             "relatedVulnerabilities": [],
+            "codeSnippet": [
+                $codeSnippetJson
+            ],
             "severity": "severity",
             "remediation": {
                 "recommendation": {
@@ -133,19 +152,58 @@ open class CodeWhispererCodeScanTestBase(projectRule: CodeInsightTestFixtureRule
                 "suggestedFixes": []
             }
         }
-    """.trimIndent()
+        """.trimIndent()
+    }
 
     private fun setupCodeScanFindings(filePath: Path) = """
-        [
-            ${setupCodeScanFinding(filePath, 1, 2)},
-            ${setupCodeScanFinding(filePath, 1, 2)}
-        ]
+    [
+        ${setupCodeScanFinding(
+        filePath,
+        1,
+        2,
+        listOf(
+            1 to "import numpy as np",
+            2 to "               import from module1 import helper"
+        )
+    )},
+        ${setupCodeScanFinding(
+        filePath,
+        1,
+        2,
+        listOf(
+            1 to "import numpy as np",
+            2 to "               import from module1 import helper"
+        )
+    )}
+    ]
+    """
+
+    private fun setupCodeScanFindingsE2E(filePath: Path) = """
+    [
+        ${setupCodeScanFinding(
+        filePath,
+        1,
+        2,
+        listOf(
+            1 to "using Utils;",
+            2 to "using Helpers.Helper;"
+        )
+    )}
+    ]
     """
 
     private fun setupCodeScanFindingsOutOfBounds(filePath: Path) = """
-        [
-            ${setupCodeScanFinding(filePath, 99999, 99999)}
-        ]
+    [
+        ${setupCodeScanFinding(
+        filePath,
+        99999,
+        99999,
+        kotlin.collections.listOf(
+            1 to "import numpy as np",
+            2 to "               import from module1 import helper"
+        )
+    )}
+    ]
     """
 
     protected fun setupResponse(filePath: Path) {
@@ -178,6 +236,11 @@ open class CodeWhispererCodeScanTestBase(projectRule: CodeInsightTestFixtureRule
             .responseMetadata(metadata)
             .build() as ListCodeScanFindingsResponse
 
+        fakeListCodeScanFindingsResponseE2E = ListCodeScanFindingsResponse.builder()
+            .codeScanFindings(setupCodeScanFindingsE2E(filePath))
+            .responseMetadata(metadata)
+            .build() as ListCodeScanFindingsResponse
+
         fakeListCodeScanFindingsOutOfBoundsIndexResponse = ListCodeScanFindingsResponse.builder()
             .codeScanFindings(setupCodeScanFindingsOutOfBounds(filePath))
             .responseMetadata(metadata)
@@ -214,6 +277,16 @@ open class CodeWhispererCodeScanTestBase(projectRule: CodeInsightTestFixtureRule
                 "detectorName": "detectorName",
                 "findingId": "findingId",
                 "relatedVulnerabilities": [],
+                "codeSnippet": [
+                    {
+                        "number": 1,
+                        "content": "codeBlock1"
+                    },
+                    {
+                        "number": 2,
+                        "content": "codeBlock2"
+                    }
+                ],
                 "severity": "severity",
                 "remediation": {
                     "recommendation": {

diff --git a/...ws/toolkits/jetbrains/services/codewhisperer/codescan/CodeWhispererProjectCodeScanTest.kt b/...ws/toolkits/jetbrains/services/codewhisperer/codescan/CodeWhispererProjectCodeScanTest.kt
@@ -54,7 +54,7 @@ class CodeWhispererProjectCodeScanTest : CodeWhispererCodeScanTestBase(PythonCod
             onGeneric { createUploadUrl(any()) }.thenReturn(fakeCreateUploadUrlResponse)
             onGeneric { createCodeScan(any(), any()) }.thenReturn(fakeCreateCodeScanResponse)
             onGeneric { getCodeScan(any(), any()) }.thenReturn(fakeGetCodeScanResponse)
-            onGeneric { listCodeScanFindings(any(), any()) }.thenReturn(fakeListCodeScanFindingsResponse)
+            onGeneric { listCodeScanFindings(any(), any()) }.thenReturn(fakeListCodeScanFindingsResponseE2E)
         }
     }
 
@@ -104,7 +104,7 @@ class CodeWhispererProjectCodeScanTest : CodeWhispererCodeScanTestBase(PythonCod
 
     @Test
     fun `e2e happy path integration test`() {
-        assertE2ERunsSuccessfully(sessionConfigSpy, project, totalLines, 10, totalSize, 2)
+        assertE2ERunsSuccessfully(sessionConfigSpy, project, totalLines, 10, totalSize, 1)
     }
 
     private fun setupCsharpProject() {