feat(amazonq): hook up lsp payload encryption (aws#5370)

Author: rli

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLspService.kt

@@ -37,6 +37,7 @@ import org.slf4j.event.Level
 import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.core.utils.warn
 import software.aws.toolkits.jetbrains.isDeveloperMode
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.encryption.JwtEncryptionManager
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.createExtendedClientMetadata
 import software.aws.toolkits.jetbrains.services.telemetry.ClientMetadata
 import java.io.IOException
@@ -111,6 +112,8 @@ class AmazonQLspService(private val project: Project, private val cs: CoroutineS
 }
 
 private class AmazonQServerInstance(private val project: Project, private val cs: CoroutineScope) : Disposable {
+    private val encryptionManager = JwtEncryptionManager()
+
     private val launcher: Launcher<AmazonQLanguageServer>
 
     private val languageServer: AmazonQLanguageServer
@@ -172,7 +175,11 @@ private class AmazonQServerInstance(private val project: Project, private val cs
         }
 
     init {
-        val cmd = GeneralCommandLine("amazon-q-lsp")
+        val cmd = GeneralCommandLine(
+            "amazon-q-lsp",
+            "--stdio",
+            "--set-credentials-encryption-key",
+        )
 
         launcherHandler = KillableColoredProcessHandler.Silent(cmd)
         val inputWrapper = LSPProcessListener()
@@ -207,6 +214,9 @@ private class AmazonQServerInstance(private val project: Project, private val cs
         launcherFuture = launcher.startListening()
 
         cs.launch {
+            // encryption info must be sent within 5s or Flare process will exit
+            encryptionManager.writeInitializationPayload(launcherHandler.process.outputStream)
+
             val initializeResult = try {
                 withTimeout(Duration.ofSeconds(10)) {
                     languageServer.initialize(createInitializeParams()).await()

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManager.kt

@@ -0,0 +1,65 @@
+// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package software.aws.toolkits.jetbrains.services.amazonq.lsp.encryption
+
+import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper
+import com.nimbusds.jose.EncryptionMethod
+import com.nimbusds.jose.JWEAlgorithm
+import com.nimbusds.jose.JWEHeader
+import com.nimbusds.jose.JWEObject
+import com.nimbusds.jose.Payload
+import com.nimbusds.jose.crypto.DirectDecrypter
+import com.nimbusds.jose.crypto.DirectEncrypter
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.EncryptionInitializationRequest
+import java.io.OutputStream
+import java.security.SecureRandom
+import java.util.Base64
+import javax.crypto.SecretKey
+import javax.crypto.spec.SecretKeySpec
+
+class JwtEncryptionManager(private val key: SecretKey) {
+    constructor() : this(generateHmacKey())
+
+    private val mapper = jacksonObjectMapper()
+
+    fun writeInitializationPayload(os: OutputStream) {
+        val payload = EncryptionInitializationRequest(
+            EncryptionInitializationRequest.Version.V1_0,
+            EncryptionInitializationRequest.Mode.JWT,
+            Base64.getUrlEncoder().withoutPadding().encodeToString(key.encoded)
+        )
+
+        // write directly to stream because utils are closing the underlying stream
+        os.write("${mapper.writeValueAsString(payload)}\n".toByteArray())
+    }
+
+    fun encrypt(data: Any): String {
+        val header = JWEHeader(JWEAlgorithm.DIR, EncryptionMethod.A256GCM)
+        val payload = if (data is String) {
+            Payload(data)
+        } else {
+            Payload(mapper.writeValueAsBytes(data))
+        }
+
+        val jweObject = JWEObject(header, payload)
+        jweObject.encrypt(DirectEncrypter(key))
+
+        return jweObject.serialize()
+    }
+
+    fun decrypt(jwt: String): String {
+        val jweObject = JWEObject.parse(jwt)
+        jweObject.decrypt(DirectDecrypter(key))
+
+        return jweObject.payload.toString()
+    }
+
+    private companion object {
+        private fun generateHmacKey(): SecretKey {
+            val keyBytes = ByteArray(32)
+            SecureRandom().nextBytes(keyBytes)
+            return SecretKeySpec(keyBytes, "HmacSHA256")
+        }
+    }
+}

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/EncryptionInitializationRequest.kt

@@ -0,0 +1,20 @@
+// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package software.aws.toolkits.jetbrains.services.amazonq.lsp.model
+
+import com.fasterxml.jackson.annotation.JsonValue
+
+data class EncryptionInitializationRequest(
+    val version: Version,
+    val mode: Mode,
+    val key: String,
+) {
+    enum class Version(@JsonValue val value: String) {
+        V1_0("1.0"),
+    }
+
+    enum class Mode(@JsonValue val value: String) {
+        JWT("JWT"),
+    }
+}

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/credentials/UpdateCredentialsPayload.kt

@@ -5,7 +5,7 @@ package software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentia
 
 data class UpdateCredentialsPayload(
     val data: String,
-    val encrypted: String,
+    val encrypted: Boolean,
 )
 
 data class UpdateCredentialsPayloadData(

plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManagerTest.kt

@@ -0,0 +1,77 @@
+// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package software.aws.toolkits.jetbrains.services.amazonq.lsp.encryption
+
+import com.nimbusds.jose.JOSEException
+import org.assertj.core.api.Assertions.assertThat
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.assertThrows
+import java.io.ByteArrayOutputStream
+import java.util.concurrent.atomic.AtomicBoolean
+import javax.crypto.spec.SecretKeySpec
+import kotlin.random.Random
+
+class JwtEncryptionManagerTest {
+    @Test
+    fun `uses a different encryption key for each instance`() {
+        val blob = Random.Default.nextBytes(256)
+        val sut1 = JwtEncryptionManager()
+        val encrypted = sut1.encrypt(blob)
+
+        assertThrows<JOSEException> {
+            assertThat(sut1.decrypt(encrypted))
+                .isNotEqualTo(JwtEncryptionManager().decrypt(encrypted))
+        }
+    }
+
+    @Test
+    @OptIn(ExperimentalStdlibApi::class)
+    fun `encryption is stable with static key`() {
+        val blob = Random.Default.nextBytes(256)
+        val bytes = "DEADBEEF".repeat(8).hexToByteArray() // 32 bytes
+        val key = SecretKeySpec(bytes, "HmacSHA256")
+        val sut1 = JwtEncryptionManager(key)
+        val encrypted = sut1.encrypt(blob)
+
+        // each encrypt() call will use a different IV so we can't just directly compare
+        assertThat(sut1.decrypt(encrypted))
+            .isEqualTo(JwtEncryptionManager(key).decrypt(encrypted))
+    }
+
+    @Test
+    fun `encryption can be round-tripped`() {
+        val sut = JwtEncryptionManager()
+        val blob = "DEADBEEF".repeat(8)
+        assertThat(sut.decrypt(sut.encrypt(blob))).isEqualTo(blob)
+    }
+
+    @Test
+    @OptIn(ExperimentalStdlibApi::class)
+    fun writeInitializationPayload() {
+        val bytes = "DEADBEEF".repeat(8).hexToByteArray() // 32 bytes
+        val key = SecretKeySpec(bytes, "HmacSHA256")
+
+        val closed = AtomicBoolean(false)
+        val os = object : ByteArrayOutputStream() {
+            override fun close() {
+                closed.set(true)
+            }
+        }
+        JwtEncryptionManager(key).writeInitializationPayload(os)
+        assertThat(os.toString())
+            // Flare requires encryption ends with new line
+            // https://github.com/aws/language-server-runtimes/blob/4d7f81295dc12b59ed2e1c0ebaedb85ccb86cf76/runtimes/README.md#encryption
+            .endsWith("\n")
+            .isEqualTo(
+                // language=JSON
+                """
+            |{"version":"1.0","mode":"JWT","key":"3q2-796tvu_erb7v3q2-796tvu_erb7v3q2-796tvu8"}
+            |
+                """.trimMargin()
+            )
+
+        // writeInitializationPayload should not close the stream
+        assertThat(closed.get()).isFalse
+    }
+}