Skip to content
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@
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
Expand Down Expand Up @@ -111,6 +112,8 @@
}

private class AmazonQServerInstance(private val project: Project, private val cs: CoroutineScope) : Disposable {
private val encryptionManager = JwtEncryptionManager()

Check warning on line 115 in plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLspService.kt

View check run for this annotation

Codecov / codecov/patch

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

Added line #L115 was not covered by tests

private val launcher: Launcher<AmazonQLanguageServer>

private val languageServer: AmazonQLanguageServer
Expand Down Expand Up @@ -172,7 +175,11 @@
}

init {
val cmd = GeneralCommandLine("amazon-q-lsp")
val cmd = GeneralCommandLine(
"amazon-q-lsp",
"--stdio",
"--set-credentials-encryption-key",

Check warning on line 181 in plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLspService.kt

View check run for this annotation

Codecov / codecov/patch

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

Added lines #L178 - L181 were not covered by tests
)

launcherHandler = KillableColoredProcessHandler.Silent(cmd)
val inputWrapper = LSPProcessListener()
Expand Down Expand Up @@ -207,6 +214,9 @@
launcherFuture = launcher.startListening()

cs.launch {
// encryption info must be sent within 5s or Flare process will exit
encryptionManager.writeInitializationPayload(launcherHandler.process.outputStream)

Check warning on line 218 in plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLspService.kt

View check run for this annotation

Codecov / codecov/patch

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

Added line #L218 was not covered by tests

val initializeResult = try {
withTimeout(Duration.ofSeconds(10)) {
languageServer.initialize(createInitializeParams()).await()
Expand Down
Original file line number Diff line number Diff line change
@@ -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())

Check warning on line 22 in plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManager.kt

View check run for this annotation

Codecov / codecov/patch

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

Added lines #L21 - L22 were not covered by tests

private val mapper = jacksonObjectMapper()

Check warning on line 24 in plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManager.kt

View check run for this annotation

Codecov / codecov/patch

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

Added line #L24 was not covered by tests

fun writeInitializationPayload(os: OutputStream) {
val payload = EncryptionInitializationRequest(
EncryptionInitializationRequest.Version.V1_0,
EncryptionInitializationRequest.Mode.JWT,
Base64.getUrlEncoder().withoutPadding().encodeToString(key.encoded)

Check warning on line 30 in plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManager.kt

View check run for this annotation

Codecov / codecov/patch

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

Added lines #L27 - L30 were not covered by tests
)

// write directly to stream because utils are closing the underlying stream
os.write("${mapper.writeValueAsString(payload)}\n".toByteArray())
}

Check warning on line 35 in plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManager.kt

View check run for this annotation

Codecov / codecov/patch

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

Added lines #L34 - L35 were not covered by tests

fun encrypt(data: Any): String {

Check warning on line 37 in plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManager.kt

View workflow job for this annotation

GitHub Actions / Qodana Community for JVM

Unused symbol

Function "encrypt" is never used

Check warning

Code scanning / QDJVMC

Unused symbol Warning

Function "encrypt" is never used
val header = JWEHeader(JWEAlgorithm.DIR, EncryptionMethod.A256GCM)

Check warning on line 38 in plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManager.kt

View check run for this annotation

Codecov / codecov/patch

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

Added line #L38 was not covered by tests
val payload = if (data is String) {
Payload(data)

Check warning on line 40 in plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManager.kt

View check run for this annotation

Codecov / codecov/patch

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

Added line #L40 was not covered by tests
} else {
Payload(mapper.writeValueAsBytes(data))

Check warning on line 42 in plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManager.kt

View check run for this annotation

Codecov / codecov/patch

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

Added line #L42 was not covered by tests
}

val jweObject = JWEObject(header, payload)
jweObject.encrypt(DirectEncrypter(key))

Check warning on line 46 in plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManager.kt

View check run for this annotation

Codecov / codecov/patch

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

Added lines #L45 - L46 were not covered by tests

return jweObject.serialize()

Check warning on line 48 in plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManager.kt

View check run for this annotation

Codecov / codecov/patch

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

Added line #L48 was not covered by tests
}

fun decrypt(jwt: String): String {

Check warning on line 51 in plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManager.kt

View workflow job for this annotation

GitHub Actions / Qodana Community for JVM

Unused symbol

Function "decrypt" is never used

Check warning

Code scanning / QDJVMC

Unused symbol Warning

Function "decrypt" is never used
val jweObject = JWEObject.parse(jwt)
jweObject.decrypt(DirectDecrypter(key))

Check warning on line 53 in plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManager.kt

View check run for this annotation

Codecov / codecov/patch

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

Added lines #L52 - L53 were not covered by tests

return jweObject.payload.toString()

Check warning on line 55 in plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManager.kt

View check run for this annotation

Codecov / codecov/patch

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

Added line #L55 was not covered by tests
}

private companion object {
private fun generateHmacKey(): SecretKey {
val keyBytes = ByteArray(32)
SecureRandom().nextBytes(keyBytes)
return SecretKeySpec(keyBytes, "HmacSHA256")

Check warning on line 62 in plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManager.kt

View check run for this annotation

Codecov / codecov/patch

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

Added lines #L60 - L62 were not covered by tests
}
}
}
Original file line number Diff line number Diff line change
@@ -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,

Check warning on line 11 in plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/EncryptionInitializationRequest.kt

View check run for this annotation

Codecov / codecov/patch

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

Added lines #L8 - L11 were not covered by tests
) {
enum class Version(@JsonValue val value: String) {
V1_0("1.0"),
}

Check warning on line 15 in plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/EncryptionInitializationRequest.kt

View check run for this annotation

Codecov / codecov/patch

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

Added lines #L13 - L15 were not covered by tests

enum class Mode(@JsonValue val value: String) {
JWT("JWT"),
}

Check warning on line 19 in plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/EncryptionInitializationRequest.kt

View check run for this annotation

Codecov / codecov/patch

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

Added lines #L17 - L19 were not covered by tests
}
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@

data class UpdateCredentialsPayload(
val data: String,
val encrypted: String,
val encrypted: Boolean,

Check warning on line 8 in plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/credentials/UpdateCredentialsPayload.kt

View check run for this annotation

Codecov / codecov/patch

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

Added line #L8 was not covered by tests
)

data class UpdateCredentialsPayloadData(
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,57 @@
// 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()
)
}
}
Loading