Skip to content

Commit 6166b30

Browse files
rlirli
committed
feat(amazonq): hook up lsp payload encryption
Certain LSP messages require AES-256-GCM encryption packaged as JWE
1 parent 4a29d0e commit 6166b30

File tree

4 files changed

+192
-16
lines changed

4 files changed

+192
-16
lines changed

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

Lines changed: 111 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -15,25 +15,40 @@ import com.intellij.openapi.Disposable
1515
import com.intellij.openapi.components.Service
1616
import com.intellij.openapi.components.service
1717
import com.intellij.openapi.project.Project
18+
import com.intellij.openapi.util.Disposer
1819
import com.intellij.openapi.util.Key
1920
import com.intellij.util.io.await
2021
import kotlinx.coroutines.CoroutineScope
22+
import kotlinx.coroutines.TimeoutCancellationException
2123
import kotlinx.coroutines.launch
24+
import kotlinx.coroutines.time.withTimeout
25+
import org.eclipse.lsp4j.ClientCapabilities
26+
import org.eclipse.lsp4j.ClientInfo
27+
import org.eclipse.lsp4j.FileOperationsWorkspaceCapabilities
2228
import org.eclipse.lsp4j.InitializeParams
2329
import org.eclipse.lsp4j.InitializedParams
30+
import org.eclipse.lsp4j.SynchronizationCapabilities
31+
import org.eclipse.lsp4j.TextDocumentClientCapabilities
32+
import org.eclipse.lsp4j.WorkspaceClientCapabilities
33+
import org.eclipse.lsp4j.WorkspaceFolder
2434
import org.eclipse.lsp4j.jsonrpc.Launcher
2535
import org.eclipse.lsp4j.launch.LSPLauncher
2636
import org.slf4j.event.Level
2737
import software.aws.toolkits.core.utils.getLogger
2838
import software.aws.toolkits.core.utils.warn
2939
import software.aws.toolkits.jetbrains.isDeveloperMode
40+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.encryption.JwtEncryptionManager
41+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.createExtendedClientMetadata
42+
import software.aws.toolkits.jetbrains.services.telemetry.ClientMetadata
3043
import java.io.IOException
3144
import java.io.OutputStreamWriter
3245
import java.io.PipedInputStream
3346
import java.io.PipedOutputStream
3447
import java.io.PrintWriter
3548
import java.io.StringWriter
49+
import java.net.URI
3650
import java.nio.charset.StandardCharsets
51+
import java.time.Duration
3752
import java.util.concurrent.Future
3853

3954
// https://github.com/redhat-developer/lsp4ij/blob/main/src/main/java/com/redhat/devtools/lsp4ij/server/LSPProcessListener.java
@@ -70,7 +85,35 @@ internal class LSPProcessListener : ProcessListener {
7085
}
7186

7287
@Service(Service.Level.PROJECT)
73-
class AmazonQLspService(project: Project, private val cs: CoroutineScope) : Disposable {
88+
class AmazonQLspService(private val project: Project, private val cs: CoroutineScope) : Disposable {
89+
private var instance: AmazonQServerInstance? = null
90+
91+
init {
92+
cs.launch {
93+
// manage lifecycle RAII-like so we can restart at arbitrary time
94+
// and suppress IDE error if server fails to start
95+
try {
96+
instance = AmazonQServerInstance(project, cs).also {
97+
Disposer.register(this@AmazonQLspService, it)
98+
}
99+
} catch (e: Exception) {
100+
LOG.warn(e) { "Failed to start LSP server" }
101+
}
102+
}
103+
}
104+
105+
override fun dispose() {
106+
}
107+
108+
companion object {
109+
private val LOG = getLogger<AmazonQLspService>()
110+
fun getInstance(project: Project) = project.service<AmazonQLspService>()
111+
}
112+
}
113+
114+
private class AmazonQServerInstance(private val project: Project, private val cs: CoroutineScope) : Disposable {
115+
private val encryptionManager = JwtEncryptionManager()
116+
74117
private val launcher: Launcher<AmazonQLanguageServer>
75118

76119
private val languageServer: AmazonQLanguageServer
@@ -80,8 +123,63 @@ class AmazonQLspService(project: Project, private val cs: CoroutineScope) : Disp
80123
private val launcherFuture: Future<Void>
81124
private val launcherHandler: KillableProcessHandler
82125

126+
private fun createClientCapabilities(): ClientCapabilities =
127+
ClientCapabilities().apply {
128+
textDocument = TextDocumentClientCapabilities().apply {
129+
// For didSaveTextDocument, other textDocument/ messages always mandatory
130+
synchronization = SynchronizationCapabilities().apply {
131+
didSave = true
132+
}
133+
}
134+
135+
workspace = WorkspaceClientCapabilities().apply {
136+
applyEdit = false
137+
138+
// For workspace folder changes
139+
workspaceFolders = true
140+
141+
// For file operations (create, delete)
142+
fileOperations = FileOperationsWorkspaceCapabilities().apply {
143+
didCreate = true
144+
didDelete = true
145+
}
146+
}
147+
}
148+
149+
// needs case handling when project's base path is null: default projects/unit tests
150+
private fun createWorkspaceFolders(): List<WorkspaceFolder> =
151+
project.basePath?.let { basePath ->
152+
listOf(
153+
WorkspaceFolder(
154+
URI("file://$basePath").toString(),
155+
project.name
156+
)
157+
)
158+
}.orEmpty() // no folders to report or workspace not folder based
159+
160+
private fun createClientInfo(): ClientInfo {
161+
val metadata = ClientMetadata.getDefault()
162+
return ClientInfo().apply {
163+
name = metadata.awsProduct.toString()
164+
version = metadata.awsVersion
165+
}
166+
}
167+
168+
private fun createInitializeParams(): InitializeParams =
169+
InitializeParams().apply {
170+
processId = ProcessHandle.current().pid().toInt()
171+
capabilities = createClientCapabilities()
172+
clientInfo = createClientInfo()
173+
workspaceFolders = createWorkspaceFolders()
174+
initializationOptions = createExtendedClientMetadata()
175+
}
176+
83177
init {
84-
val cmd = GeneralCommandLine("amazon-q-lsp")
178+
val cmd = GeneralCommandLine(
179+
"amazon-q-lsp",
180+
"--stdio",
181+
"--set-credentials-encryption-key",
182+
)
85183

86184
launcherHandler = KillableColoredProcessHandler.Silent(cmd)
87185
val inputWrapper = LSPProcessListener()
@@ -116,18 +214,17 @@ class AmazonQLspService(project: Project, private val cs: CoroutineScope) : Disp
116214
launcherFuture = launcher.startListening()
117215

118216
cs.launch {
119-
val initializeResult = languageServer.initialize(
120-
InitializeParams().apply {
121-
// does this work on windows
122-
processId = ProcessHandle.current().pid().toInt()
123-
// capabilities
124-
// client info
125-
// trace?
126-
// workspace folders?
127-
// anything else we need?
217+
// encryption info must be sent within 5s or Flare process will exit
218+
encryptionManager.writeInitializationPayload(launcherHandler.process.outputStream)
219+
220+
val initializeResult = try {
221+
withTimeout(Duration.ofSeconds(10)) {
222+
languageServer.initialize(createInitializeParams()).await()
128223
}
129-
// probably need a timeout
130-
).await()
224+
} catch (_: TimeoutCancellationException) {
225+
LOG.warn { "LSP initialization timed out" }
226+
null
227+
}
131228

132229
// then if this succeeds then we can allow the client to send requests
133230
if (initializeResult == null) {
@@ -154,7 +251,6 @@ class AmazonQLspService(project: Project, private val cs: CoroutineScope) : Disp
154251
}
155252

156253
companion object {
157-
private val LOG = getLogger<AmazonQLspService>()
158-
fun getInstance(project: Project) = project.service<AmazonQLspService>()
254+
private val LOG = getLogger<AmazonQServerInstance>()
159255
}
160256
}

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

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package software.aws.toolkits.jetbrains.services.amazonq.lsp.encryption
5+
6+
import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper
7+
import com.nimbusds.jose.EncryptionMethod
8+
import com.nimbusds.jose.JWEAlgorithm
9+
import com.nimbusds.jose.JWEHeader
10+
import com.nimbusds.jose.JWEObject
11+
import com.nimbusds.jose.Payload
12+
import com.nimbusds.jose.crypto.DirectDecrypter
13+
import com.nimbusds.jose.crypto.DirectEncrypter
14+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.EncryptionInitializationRequest
15+
import java.io.OutputStream
16+
import java.security.SecureRandom
17+
import java.util.Base64
18+
import javax.crypto.SecretKey
19+
import javax.crypto.spec.SecretKeySpec
20+
21+
class JwtEncryptionManager(private val key: SecretKey) {
22+
constructor() : this(generateHmacKey())
23+
24+
private val mapper = jacksonObjectMapper()
25+
26+
fun writeInitializationPayload(os: OutputStream) {
27+
val payload = EncryptionInitializationRequest(
28+
EncryptionInitializationRequest.Version.V1_0,
29+
EncryptionInitializationRequest.Mode.JWT,
30+
Base64.getUrlEncoder().withoutPadding().encodeToString(key.encoded)
31+
)
32+
33+
// write directly to stream because utils are closing the underlying stream
34+
os.write("${mapper.writeValueAsString(payload)}\n".toByteArray())
35+
}
36+
37+
fun encrypt(data: Any): String {
38+
val header = JWEHeader(JWEAlgorithm.DIR, EncryptionMethod.A256GCM)
39+
val payload = Payload(mapper.writeValueAsBytes(data))
40+
val jweObject = JWEObject(header, payload)
41+
jweObject.encrypt(DirectEncrypter(key))
42+
43+
return jweObject.serialize()
44+
}
45+
46+
fun decrypt(jwt: String): String {
47+
val jweObject = JWEObject.parse(jwt)
48+
jweObject.decrypt(DirectDecrypter(key))
49+
50+
return jweObject.payload.toString()
51+
}
52+
53+
private companion object {
54+
private fun generateHmacKey(): SecretKey {
55+
val keyBytes = ByteArray(32)
56+
SecureRandom().nextBytes(keyBytes)
57+
return SecretKeySpec(keyBytes, "HmacSHA256")
58+
}
59+
}
60+
}

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

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package software.aws.toolkits.jetbrains.services.amazonq.lsp.model
5+
6+
import com.fasterxml.jackson.annotation.JsonValue
7+
8+
data class EncryptionInitializationRequest(
9+
val version: Version,
10+
val mode: Mode,
11+
val key: String,
12+
) {
13+
enum class Version(@JsonValue val value: String) {
14+
V1_0("1.0"),
15+
}
16+
17+
enum class Mode(@JsonValue val value: String) {
18+
JWT("JWT"),
19+
}
20+
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ package software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentia
55

66
data class UpdateCredentialsPayload(
77
val data: String,
8-
val encrypted: String,
8+
val encrypted: Boolean,
99
)
1010

1111
data class UpdateCredentialsPayloadData(

0 commit comments

Comments
 (0)