tst

Author: rli

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))
 

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()
+            )
+    }
+}