Amazon Q[ACR]: fixing the line number on document change

.changes/next-release/bugfix-6862d9d5-751a-40c6-a907-53478af6991c.json

@@ -0,0 +1,4 @@
+{
+  "type" : "bugfix",
+  "description" : "Security Scan: Improved accuracy when applying security fixes"
+}

plugins/toolkit/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/codescan/CodeWhispererCodeScanManager.kt

@@ -74,6 +74,7 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.util.runIfIdcConne
 import software.aws.toolkits.jetbrains.utils.isQConnected
 import software.aws.toolkits.jetbrains.utils.isQExpired
 import software.aws.toolkits.jetbrains.utils.isRunningOnRemoteBackend
+import software.aws.toolkits.jetbrains.utils.offsetSuggestedFix
 import software.aws.toolkits.resources.message
 import software.aws.toolkits.telemetry.Result
 import java.time.Duration
@@ -457,6 +458,18 @@ class CodeWhispererCodeScanManager(val project: Project) {
         }
     }
 
+    fun updateScanNodesForOffSet(file: VirtualFile, lineOffset: Int, editedTextRange: TextRange) {
+        val document = FileDocumentManager.getInstance().getDocument(file) ?: return
+        scanNodesLookup[file]?.forEach { node ->
+            val issue = node.userObject as CodeWhispererCodeScanIssue
+            if (document.getLineNumber(editedTextRange.startOffset) <= issue.startLine) {
+                issue.startLine = issue.startLine + lineOffset
+                issue.endLine = issue.endLine + lineOffset
+                issue.suggestedFixes = issue.suggestedFixes.map { fix -> offsetSuggestedFix(fix, lineOffset) }
+            }
+        }
+    }
+
     private fun CodeWhispererCodeScanIssue.copyRange(newRange: TextRange): CodeWhispererCodeScanIssue {
         val newStartLine = document.getLineNumber(newRange.startOffset)
         val newStartCol = newRange.startOffset - document.getLineStartOffset(newStartLine)
@@ -497,7 +510,7 @@ class CodeWhispererCodeScanManager(val project: Project) {
             val editorFactory = EditorFactory.getInstance()
             editorFactory.eventMulticaster.addDocumentListener(documentListener, project)
             editorFactory.addEditorFactoryListener(fileListener, project)
-            EditorFactory.getInstance().eventMulticaster.addEditorMouseMotionListener(
+            editorFactory.eventMulticaster.addEditorMouseMotionListener(
                 editorMouseListener,
                 codeScanIssuesContent
             )
@@ -672,9 +685,9 @@ class CodeWhispererCodeScanManager(val project: Project) {
 data class CodeWhispererCodeScanIssue(
     val project: Project,
     val file: VirtualFile,
-    val startLine: Int,
+    var startLine: Int,
     val startCol: Int,
-    val endLine: Int,
+    var endLine: Int,
     val endCol: Int,
     val title: @InspectionMessage String,
     val description: Description,
@@ -685,7 +698,7 @@ data class CodeWhispererCodeScanIssue(
     val relatedVulnerabilities: List<String>,
     val severity: String,
     val recommendation: Recommendation,
-    val suggestedFixes: List<SuggestedFix>,
+    var suggestedFixes: List<SuggestedFix>,
     val issueSeverity: HighlightDisplayLevel = HighlightDisplayLevel.WARNING,
     val isInvalid: Boolean = false,
     var rangeHighlighter: RangeHighlighterEx? = null

plugins/toolkit/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/codescan/listeners/CodeWhispererCodeScanDocumentListener.kt

@@ -26,8 +26,17 @@ internal class CodeWhispererCodeScanDocumentListener(val project: Project) : Doc
 
         val fileEditorManager = FileEditorManager.getInstance(project)
         val activeEditor = fileEditorManager.selectedEditor
+        val deletedLineCount = event.oldFragment.toString().count { it == '\n' }
+        val insertedLineCount = event.newFragment.toString().count { it == '\n' }
+
+        val lineOffset = when {
+            deletedLineCount == 0 && insertedLineCount != 0 -> insertedLineCount
+            deletedLineCount != 0 && insertedLineCount == 0 -> -deletedLineCount
+            else -> 0
+        }
 
         val editedTextRange = TextRange.create(event.offset, event.offset + event.oldLength)
+        scanManager.updateScanNodesForOffSet(file, lineOffset, editedTextRange)
         val nodes = scanManager.getOverlappingScanNodes(file, editedTextRange)
         nodes.forEach {
             val issue = it.userObject as CodeWhispererCodeScanIssue

plugins/toolkit/jetbrains-core/src/software/aws/toolkits/jetbrains/utils/TextUtils.kt

@@ -12,6 +12,7 @@ import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.text.StringUtil
 import com.intellij.psi.PsiFileFactory
 import com.intellij.psi.codeStyle.CodeStyleManager
+import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.SuggestedFix
 
 fun formatText(project: Project, language: Language, content: String): String {
     var result = content
@@ -42,3 +43,18 @@ fun applyPatch(patch: String, fileContent: String, filePath: String): String? {
     val unifiedPatch = generateUnifiedPatch(patch, filePath)
     return PlainSimplePatchApplier.apply(fileContent, unifiedPatch.hunks)
 }
+
+fun offsetSuggestedFix(suggestedFix: SuggestedFix, lines: Int): SuggestedFix {
+    val updatedCode = suggestedFix.code.replace(
+        Regex("""(@@ -)(\d+)(,\d+ \+)(\d+)(,\d+ @@)""")
+    ) { result ->
+        val prefix = result.groupValues[1]
+        val startLine = result.groupValues[2].toInt() + lines
+        val middle = result.groupValues[3]
+        val endLine = result.groupValues[4].toInt() + lines
+        val suffix = result.groupValues[5]
+        "$prefix$startLine$middle$endLine$suffix"
+    }
+
+    return suggestedFix.copy(code = updatedCode)
+}

plugins/toolkit/jetbrains-core/tst/software/aws/toolkits/jetbrains/utils/TextUtilsTest.kt

@@ -8,9 +8,11 @@ import com.intellij.testFramework.ProjectRule
 import com.intellij.testFramework.runInEdtAndWait
 import org.assertj.core.api.Assertions.assertThat
 import org.intellij.lang.annotations.Language
+import org.junit.Assert.assertEquals
 import org.junit.Rule
 import org.junit.Test
 import software.aws.toolkits.core.utils.convertMarkdownToHTML
+import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.SuggestedFix
 
 class TextUtilsTest {
     @Rule
@@ -137,4 +139,108 @@ class TextUtilsTest {
         val inputPatchLines = inputPatch.split("\n")
         hunk.lines.forEachIndexed { index, patchLine -> assertThat(inputPatchLines[index + 1].substring(1)).isEqualTo(patchLine.text) }
     }
+
+    @Test
+    fun offsetSuggestedFixUpdateLineNumbersWithInsertion() {
+        val suggestedFix = SuggestedFix(
+            code = """
+            @@ -1,3 +1,4 @@
+             fun main() {
+            +    val greeting = "Hello, Suggested Fix is Here!"
+                 println("Hello, Suggested Fix is Here!")
+             }
+            """.trimIndent(),
+            description = "Add a variable for the greeting"
+        )
+
+        val expectedCode = """
+            @@ -2,3 +2,4 @@
+             fun main() {
+            +    val greeting = "Hello, Suggested Fix is Here!"
+                 println("Hello, Suggested Fix is Here!")
+             }
+        """.trimIndent()
+
+        val result = offsetSuggestedFix(suggestedFix, 1)
+        assertEquals(expectedCode, result.code)
+    }
+
+    @Test
+    fun offsetSuggestedFixUpdateMultipleLineNumbersWithInsertion() {
+        val suggestedFix = SuggestedFix(
+            code = """
+            @@ -1,3 +1,5 @@
+             fun main() {
+            +    val greeting = "Hello, Suggested Fix is Here!"
+                 println("Hello, Suggested Fix is Here!"))
+            +    println("Hello, Welcome to Amazon Q")
+             }
+            """.trimIndent(),
+            description = "Add a variable for the greeting with multiple lines"
+        )
+
+        val expectedCode = """
+            @@ -4,3 +4,5 @@
+             fun main() {
+            +    val greeting = "Hello, Suggested Fix is Here!"
+                 println("Hello, Suggested Fix is Here!"))
+            +    println("Hello, Welcome to Amazon Q")
+             }
+        """.trimIndent()
+
+        val result = offsetSuggestedFix(suggestedFix, 3)
+        assertEquals(expectedCode, result.code)
+    }
+
+    @Test
+    fun offsetSuggestedFixUpdateLineNumbersWithDeletion() {
+        val suggestedFix = SuggestedFix(
+            code = """
+            @@ -24,3 +24,4 @@
+             fun main() {
+            +    val greeting = "Hello, Suggested Fix is Here!"
+                 println("Hello, Suggested Fix is Here!")
+             }
+            """.trimIndent(),
+            description = "Add a variable for the greeting"
+        )
+
+        val expectedCode = """
+            @@ -19,3 +19,4 @@
+             fun main() {
+            +    val greeting = "Hello, Suggested Fix is Here!"
+                 println("Hello, Suggested Fix is Here!")
+             }
+        """.trimIndent()
+
+        val result = offsetSuggestedFix(suggestedFix, -5)
+        assertEquals(expectedCode, result.code)
+    }
+
+    @Test
+    fun offsetSuggestedFixUpdateMultipleLineNumbersWithDeletion() {
+        val suggestedFix = SuggestedFix(
+            code = """
+            @@ -10,3 +10,5 @@
+             fun main() {
+            +    val greeting = "Hello, Suggested Fix is Here!"
+                 println("Hello, Suggested Fix is Here!"))
+            +    println("Hello, Welcome to Amazon Q")
+             }
+            """.trimIndent(),
+            description = "Add a variable for the greeting with multiple lines"
+        )
+
+        val expectedCode = """
+            @@ -8,3 +8,5 @@
+             fun main() {
+            +    val greeting = "Hello, Suggested Fix is Here!"
+                 println("Hello, Suggested Fix is Here!"))
+            +    println("Hello, Welcome to Amazon Q")
+             }
+        """.trimIndent()
+
+        val result = offsetSuggestedFix(suggestedFix, -2)
+        assertEquals(expectedCode, result.code)
+    }
 }