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..bca385682ea --- /dev/null +++ b/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") + } + } +} 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( 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..365fd2759b0 --- /dev/null +++ b/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 { + 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 + } +}