fix / cleanup

Author: rli

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

@@ -60,14 +60,12 @@ import software.aws.toolkits.jetbrains.services.amazonq.lsp.encryption.JwtEncryp
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.createExtendedClientMetadata
 import software.aws.toolkits.jetbrains.services.amazonq.project.EncoderServer
 import software.aws.toolkits.jetbrains.services.amazonq.project.IndexRequest
-import software.aws.toolkits.jetbrains.services.amazonq.project.IndexUpdateMode
 import software.aws.toolkits.jetbrains.services.amazonq.project.InlineBm25Chunk
-import software.aws.toolkits.jetbrains.services.amazonq.project.LspMessage
+import software.aws.toolkits.jetbrains.services.amazonq.project.LspRequest
 import software.aws.toolkits.jetbrains.services.amazonq.project.ProjectContextProvider
 import software.aws.toolkits.jetbrains.services.amazonq.project.ProjectContextProvider.Usage
 import software.aws.toolkits.jetbrains.services.amazonq.project.QueryChatRequest
 import software.aws.toolkits.jetbrains.services.amazonq.project.QueryInlineCompletionRequest
-import software.aws.toolkits.jetbrains.services.amazonq.project.RelevantDocument
 import software.aws.toolkits.jetbrains.services.amazonq.project.UpdateIndexRequest
 import software.aws.toolkits.jetbrains.services.telemetry.ClientMetadata
 import java.io.IOException
@@ -350,6 +348,8 @@ private class AmazonQServerInstance(private val project: Project, private val cs
 
 
 class EncoderServer2(private val encoderServer: EncoderServer, private val commandLine: GeneralCommandLine, private val project: Project, private val cs: CoroutineScope) : Disposable {
+    private val encryptionManager = JwtEncryptionManager()
+
     private val launcher: Launcher<EncoderServerLspInterface>
 
     val languageServer: EncoderServerLspInterface
@@ -465,8 +465,9 @@ class EncoderServer2(private val encoderServer: EncoderServer, private val comma
             )
             .wrapMessages { consumer ->
                 MessageConsumer { message ->
-                    if (message is RequestMessage && message.params is LspMessage) {
-                        message.params = encoderServer.encrypt(jacksonObjectMapper().writeValueAsString(message))
+                    if (message is RequestMessage && message.params is LspRequest) {
+                        println(message.params)
+                        message.params = encryptionManager.encrypt(jacksonObjectMapper().writeValueAsString(message.params))
                     }
                     consumer.consume(message)
                 }
@@ -479,7 +480,7 @@ class EncoderServer2(private val encoderServer: EncoderServer, private val comma
 
         initializer = cs.launch {
             // encryption info must be sent within 5s or Flare process will exit
-            launcherHandler.process.outputStream.write(encoderServer.getEncryptionRequest().toByteArray())
+            encryptionManager.writeInitializationPayload(launcherHandler.process.outputStream)
 
             val initializeResult = try {
                 withTimeout(5.seconds) {

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/project/EncoderServer.kt

@@ -3,27 +3,18 @@
 
 package software.aws.toolkits.jetbrains.services.amazonq.project
 
-import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper
 import com.intellij.execution.configurations.GeneralCommandLine
-import com.intellij.execution.process.KillableProcessHandler
 import com.intellij.openapi.Disposable
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.project.Project
-import com.intellij.openapi.util.process.ProcessCloseUtil
+import com.intellij.openapi.util.Disposer
 import com.intellij.util.io.HttpRequests
 import com.intellij.util.io.createDirectories
 import com.intellij.util.net.NetUtils
-import com.nimbusds.jose.JOSEObjectType
-import com.nimbusds.jose.JWSAlgorithm
-import com.nimbusds.jose.JWSHeader
-import com.nimbusds.jose.crypto.MACSigner
-import com.nimbusds.jwt.JWTClaimsSet
-import com.nimbusds.jwt.SignedJWT
 import kotlinx.coroutines.CoroutineScope
 import org.apache.commons.codec.digest.DigestUtils
 import software.amazon.awssdk.utils.UserHomeDirectoryUtils
 import software.aws.toolkits.core.utils.getLogger
-import software.aws.toolkits.core.utils.info
 import software.aws.toolkits.core.utils.tryDirOp
 import software.aws.toolkits.core.utils.warn
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.EncoderServer2
@@ -35,11 +26,6 @@ import java.nio.file.Files
 import java.nio.file.Path
 import java.nio.file.Paths
 import java.nio.file.attribute.PosixFilePermission
-import java.security.Key
-import java.security.SecureRandom
-import java.util.Base64
-import java.util.concurrent.atomic.AtomicInteger
-import javax.crypto.spec.SecretKeySpec
 
 class EncoderServer(val project: Project, private val cs: CoroutineScope) : Disposable {
     val cachePath = Paths.get(
@@ -49,9 +35,6 @@ class EncoderServer(val project: Project, private val cs: CoroutineScope) : Disp
     private val serverDirectoryName = "qserver-${manifestManager.currentVersion}.zip"
     val port by lazy { NetUtils.findAvailableSocketPort() }
     private val nodeRunnableName = if (manifestManager.getOs() == "windows") "node.exe" else "node"
-    private val key = generateHmacKey()
-    private var processHandler: KillableProcessHandler? = null
-    private val mapper = jacksonObjectMapper()
     lateinit var encoderServer2: EncoderServer2
 
     fun downloadArtifactsAndStartServer() {
@@ -62,39 +45,7 @@ class EncoderServer(val project: Project, private val cs: CoroutineScope) : Disp
         start()
     }
 
-    fun isNodeProcessRunning() = processHandler != null && processHandler?.process?.isAlive == true
-
-    private fun generateHmacKey(): Key {
-        val keyBytes = ByteArray(32)
-        SecureRandom().nextBytes(keyBytes)
-        return SecretKeySpec(keyBytes, "HmacSHA256")
-    }
-
-    fun encrypt(data: String): String {
-        val header = JWSHeader.Builder(JWSAlgorithm.HS256)
-            .type(JOSEObjectType.JWT)
-            .build()
-
-        val claimsSet = JWTClaimsSet.Builder()
-            .subject(Base64.getUrlEncoder().withoutPadding().encodeToString(data.toByteArray()))
-            .build()
-
-        val signedJWT = SignedJWT(header, claimsSet)
-        signedJWT.sign(MACSigner(key.encoded))
-
-        return signedJWT.serialize()
-    }
-
-    data class EncryptionRequest(
-        val version: String = "1.0",
-        val mode: String = "JWT",
-        val key: String,
-    )
-
-    fun getEncryptionRequest(): String {
-        val request = EncryptionRequest(key = Base64.getUrlEncoder().withoutPadding().encodeToString(key.encoded))
-        return mapper.writeValueAsString(request) + "\n"
-    }
+    fun isNodeProcessRunning() = encoderServer2.initializer.isCompleted
 
     private fun getCommand(): GeneralCommandLine {
         val threadCount = CodeWhispererSettings.getInstance().getProjectContextIndexThreadCount()
@@ -123,7 +74,7 @@ class EncoderServer(val project: Project, private val cs: CoroutineScope) : Disp
     }
 
     private fun close() {
-        processHandler?.process?.let { ProcessCloseUtil.close(it) }
+        Disposer.dispose(encoderServer2)
     }
 
     private fun downloadArtifactsIfNeeded() {

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/project/ProjectContextProvider.kt

@@ -6,7 +6,6 @@ package software.aws.toolkits.jetbrains.services.amazonq.project
 import com.fasterxml.jackson.annotation.JsonIgnoreProperties
 import com.fasterxml.jackson.annotation.JsonProperty
 import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper
-import com.fasterxml.jackson.module.kotlin.readValue
 import com.intellij.openapi.Disposable
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.project.BaseProjectDirectories.Companion.getBaseDirectories
@@ -31,8 +30,6 @@ import software.aws.toolkits.jetbrains.services.amazonq.SUPPLEMENTAL_CONTEXT_TIM
 import software.aws.toolkits.jetbrains.services.cwc.controller.chat.telemetry.getStartUrl
 import software.aws.toolkits.jetbrains.settings.CodeWhispererSettings
 import software.aws.toolkits.telemetry.AmazonqTelemetry
-import java.io.OutputStreamWriter
-import java.net.HttpURLConnection
 import java.util.concurrent.atomic.AtomicBoolean
 import java.util.concurrent.atomic.AtomicInteger
 import kotlin.time.measureTimedValue
@@ -212,26 +209,6 @@ class ProjectContextProvider(val project: Project, private val encoderServer: En
         )
     }
 
-    private fun setConnectionTimeout(connection: HttpURLConnection, timeout: Int) {
-        connection.connectTimeout = timeout
-        connection.readTimeout = timeout
-    }
-
-    private fun setConnectionProperties(connection: HttpURLConnection) {
-        connection.requestMethod = "POST"
-        connection.setRequestProperty("Content-Type", "text/plain")
-        connection.setRequestProperty("Accept", "text/plain")
-    }
-
-    private fun setConnectionRequest(connection: HttpURLConnection, payload: String) {
-        connection.doOutput = true
-        connection.outputStream.use { outputStream ->
-            OutputStreamWriter(outputStream).use { writer ->
-                writer.write(payload)
-            }
-        }
-    }
-
     private fun willExceedPayloadLimit(currentTotalFileSize: Long, currentFileSize: Long): Boolean {
         val maxSize = CodeWhispererSettings.getInstance().getProjectContextIndexMaxSize()
         return currentTotalFileSize.let { totalSize -> totalSize > (maxSize * 1024 * 1024 - currentFileSize) }