From e6b9e25e75ef711444bda7b9cb064adcbf48bbfa Mon Sep 17 00:00:00 2001 From: Richard Li Date: Wed, 12 Feb 2025 16:12:57 -0800 Subject: [PATCH 1/5] feat(amazonq): hook up lsp payload encryption Certain LSP messages require AES-256-GCM encryption packaged as JWE --- .../services/amazonq/lsp/AmazonQLspService.kt | 12 +++- .../lsp/encryption/JwtEncryptionManager.kt | 60 +++++++++++++++++++ .../model/EncryptionInitializationRequest.kt | 20 +++++++ .../credentials/UpdateCredentialsPayload.kt | 2 +- 4 files changed, 92 insertions(+), 2 deletions(-) create mode 100644 plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManager.kt create mode 100644 plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/EncryptionInitializationRequest.kt diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLspService.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLspService.kt index 44a6896e0bb..256465ce06b 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLspService.kt +++ b/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 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() diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManager.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManager.kt new file mode 100644 index 00000000000..05de7f1470e --- /dev/null +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManager.kt @@ -0,0 +1,60 @@ +// 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 = 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") + } + } +} diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/EncryptionInitializationRequest.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/EncryptionInitializationRequest.kt new file mode 100644 index 00000000000..53748475195 --- /dev/null +++ b/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"), + } +} diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/credentials/UpdateCredentialsPayload.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/credentials/UpdateCredentialsPayload.kt index 2b5a28b6cb4..a427330c055 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/credentials/UpdateCredentialsPayload.kt +++ b/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( From 80350e05aa0e798e3ab8a41d84d559aa18c2182d Mon Sep 17 00:00:00 2001 From: Richard Li Date: Wed, 12 Feb 2025 17:55:32 -0800 Subject: [PATCH 2/5] tst --- .../lsp/encryption/JwtEncryptionManager.kt | 7 ++- .../encryption/JwtEncryptionManagerTest.kt | 56 +++++++++++++++++++ 2 files changed, 62 insertions(+), 1 deletion(-) create mode 100644 plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManagerTest.kt diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManager.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManager.kt index 05de7f1470e..bca385682ea 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManager.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManager.kt @@ -36,7 +36,12 @@ class JwtEncryptionManager(private val key: SecretKey) { fun encrypt(data: Any): String { val header = JWEHeader(JWEAlgorithm.DIR, EncryptionMethod.A256GCM) - val payload = Payload(mapper.writeValueAsBytes(data)) + val payload = if (data is String) { + Payload(data) + } else { + Payload(mapper.writeValueAsBytes(data)) + } + val jweObject = JWEObject(header, payload) jweObject.encrypt(DirectEncrypter(key)) diff --git a/plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManagerTest.kt b/plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManagerTest.kt new file mode 100644 index 00000000000..cdc96148d6c --- /dev/null +++ b/plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManagerTest.kt @@ -0,0 +1,56 @@ +// 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 org.assertj.core.api.Assertions.assertThat +import org.junit.jupiter.api.Test +import java.io.ByteArrayOutputStream +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) + assertThat(JwtEncryptionManager().encrypt(blob)) + .isNotEqualTo(JwtEncryptionManager().encrypt(blob)) + } + + @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") + assertThat(JwtEncryptionManager(key).encrypt(blob)) + .isNotEqualTo(JwtEncryptionManager(key).encrypt(blob)) + } + + @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 os = ByteArrayOutputStream() + 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") + // language=JSON + .isEqualTo(""" + |{"version":"1.0","mode":"JWT","key":"3q2-796tvu_erb7v3q2-796tvu_erb7v3q2-796tvu8"} + | + """.trimMargin() + ) + } +} From c057ab718b6e1d5a57b8a74315097687555d12fb Mon Sep 17 00:00:00 2001 From: Richard Li Date: Wed, 12 Feb 2025 18:17:25 -0800 Subject: [PATCH 3/5] detekt --- .../amazonq/lsp/encryption/JwtEncryptionManagerTest.kt | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManagerTest.kt b/plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManagerTest.kt index cdc96148d6c..ed08ac853b9 100644 --- a/plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManagerTest.kt +++ b/plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManagerTest.kt @@ -47,10 +47,11 @@ class JwtEncryptionManagerTest { // https://github.com/aws/language-server-runtimes/blob/4d7f81295dc12b59ed2e1c0ebaedb85ccb86cf76/runtimes/README.md#encryption .endsWith("\n") // language=JSON - .isEqualTo(""" + .isEqualTo( + """ |{"version":"1.0","mode":"JWT","key":"3q2-796tvu_erb7v3q2-796tvu_erb7v3q2-796tvu8"} | - """.trimMargin() + """.trimMargin() ) } } From e923b37a033d7dbdf7b67fba6ba16a05c30faf41 Mon Sep 17 00:00:00 2001 From: Richard Li Date: Thu, 13 Feb 2025 10:44:37 -0800 Subject: [PATCH 4/5] stream assert --- .../lsp/encryption/JwtEncryptionManagerTest.kt | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManagerTest.kt b/plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManagerTest.kt index ed08ac853b9..1283e9bbd4b 100644 --- a/plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManagerTest.kt +++ b/plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManagerTest.kt @@ -6,6 +6,7 @@ package software.aws.toolkits.jetbrains.services.amazonq.lsp.encryption import org.assertj.core.api.Assertions.assertThat import org.junit.jupiter.api.Test import java.io.ByteArrayOutputStream +import java.util.concurrent.atomic.AtomicBoolean import javax.crypto.spec.SecretKeySpec import kotlin.random.Random @@ -40,18 +41,26 @@ class JwtEncryptionManagerTest { val bytes = "DEADBEEF".repeat(8).hexToByteArray() // 32 bytes val key = SecretKeySpec(bytes, "HmacSHA256") - val os = ByteArrayOutputStream() + 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") - // language=JSON .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 } } From 6a12d6c3e25e208a82f7651ba47bc1aa5945b070 Mon Sep 17 00:00:00 2001 From: Richard Li Date: Thu, 13 Feb 2025 12:41:08 -0800 Subject: [PATCH 5/5] tst --- .../encryption/JwtEncryptionManagerTest.kt | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManagerTest.kt b/plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManagerTest.kt index 1283e9bbd4b..365fd2759b0 100644 --- a/plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManagerTest.kt +++ b/plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManagerTest.kt @@ -3,8 +3,10 @@ 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 @@ -14,8 +16,13 @@ class JwtEncryptionManagerTest { @Test fun `uses a different encryption key for each instance`() { val blob = Random.Default.nextBytes(256) - assertThat(JwtEncryptionManager().encrypt(blob)) - .isNotEqualTo(JwtEncryptionManager().encrypt(blob)) + val sut1 = JwtEncryptionManager() + val encrypted = sut1.encrypt(blob) + + assertThrows { + assertThat(sut1.decrypt(encrypted)) + .isNotEqualTo(JwtEncryptionManager().decrypt(encrypted)) + } } @Test @@ -24,8 +31,12 @@ class JwtEncryptionManagerTest { val blob = Random.Default.nextBytes(256) val bytes = "DEADBEEF".repeat(8).hexToByteArray() // 32 bytes val key = SecretKeySpec(bytes, "HmacSHA256") - assertThat(JwtEncryptionManager(key).encrypt(blob)) - .isNotEqualTo(JwtEncryptionManager(key).encrypt(blob)) + 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