diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageServer.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageServer.kt index 6a40867a7e0..e8281f6b56a 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageServer.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageServer.kt @@ -4,6 +4,7 @@ package software.aws.toolkits.jetbrains.services.amazonq.lsp import org.eclipse.lsp4j.jsonrpc.messages.ResponseMessage +import org.eclipse.lsp4j.jsonrpc.services.JsonNotification import org.eclipse.lsp4j.jsonrpc.services.JsonRequest import org.eclipse.lsp4j.services.LanguageServer import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.UpdateCredentialsPayload @@ -16,4 +17,7 @@ import java.util.concurrent.CompletableFuture interface AmazonQLanguageServer : LanguageServer { @JsonRequest("aws/credentials/token/update") fun updateTokenCredentials(payload: UpdateCredentialsPayload): CompletableFuture + + @JsonNotification("aws/credentials/token/delete") + fun deleteTokenCredentials(): CompletableFuture } 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 28445578c67..49141741de0 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 @@ -154,28 +154,26 @@ class AmazonQLspService(private val project: Project, private val cs: CoroutineS instance = start() } - suspend fun execute(runnable: suspend (AmazonQLanguageServer) -> Unit) { + suspend fun execute(runnable: suspend (AmazonQLanguageServer) -> T): T { val lsp = withTimeout(10.seconds) { val holder = mutex.withLock { instance }.await() holder.initializer.join() holder.languageServer } - - runnable(lsp) + return runnable(lsp) } - fun executeSync(runnable: suspend (AmazonQLanguageServer) -> Unit) { + fun executeSync(runnable: suspend (AmazonQLanguageServer) -> T): T = runBlocking(cs.coroutineContext) { execute(runnable) } - } companion object { private val LOG = getLogger() fun getInstance(project: Project) = project.service() - fun executeIfRunning(project: Project, runnable: (AmazonQLanguageServer) -> Unit) = + fun executeIfRunning(project: Project, runnable: (AmazonQLanguageServer) -> T): T? = project.serviceIfCreated()?.executeSync(runnable) } } diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/AuthCredentialsService.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/AuthCredentialsService.kt new file mode 100644 index 00000000000..a38c8da4bbc --- /dev/null +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/AuthCredentialsService.kt @@ -0,0 +1,12 @@ +// 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.auth + +import org.eclipse.lsp4j.jsonrpc.messages.ResponseMessage +import java.util.concurrent.CompletableFuture + +interface AuthCredentialsService { + fun updateTokenCredentials(accessToken: String, encrypted: Boolean): CompletableFuture + fun deleteTokenCredentials(): CompletableFuture +} diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsService.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsService.kt new file mode 100644 index 00000000000..6fdc3bf1ef8 --- /dev/null +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsService.kt @@ -0,0 +1,51 @@ +// 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.auth + +import com.intellij.openapi.project.Project +import org.eclipse.lsp4j.jsonrpc.messages.ResponseMessage +import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService +import software.aws.toolkits.jetbrains.services.amazonq.lsp.encryption.JwtEncryptionManager +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.BearerCredentials +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.UpdateCredentialsPayload +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.UpdateCredentialsPayloadData +import java.util.concurrent.CompletableFuture + +class DefaultAuthCredentialsService( + private val project: Project, + private val encryptionManager: JwtEncryptionManager, +) : AuthCredentialsService { + + override fun updateTokenCredentials(accessToken: String, encrypted: Boolean): CompletableFuture { + val token = if (encrypted) { + encryptionManager.decrypt(accessToken) + } else { + accessToken + } + + val payload = createUpdateCredentialsPayload(token) + + return AmazonQLspService.executeIfRunning(project) { server -> + server.updateTokenCredentials(payload) + } ?: (CompletableFuture.failedFuture(IllegalStateException("LSP Server not running"))) + } + + override fun deleteTokenCredentials(): CompletableFuture = + CompletableFuture().also { completableFuture -> + AmazonQLspService.executeIfRunning(project) { server -> + server.deleteTokenCredentials() + completableFuture.complete(null) + } ?: completableFuture.completeExceptionally(IllegalStateException("LSP Server not running")) + } + + private fun createUpdateCredentialsPayload(token: String): UpdateCredentialsPayload = + UpdateCredentialsPayload( + data = encryptionManager.encrypt( + UpdateCredentialsPayloadData( + BearerCredentials(token) + ) + ), + encrypted = true + ) +} diff --git a/plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsServiceTest.kt b/plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsServiceTest.kt new file mode 100644 index 00000000000..2c4c003ba8c --- /dev/null +++ b/plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsServiceTest.kt @@ -0,0 +1,93 @@ +// 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.auth + +import com.intellij.openapi.components.serviceIfCreated +import com.intellij.openapi.project.Project +import io.mockk.every +import io.mockk.mockk +import io.mockk.verify +import org.eclipse.lsp4j.jsonrpc.messages.ResponseMessage +import org.junit.Before +import org.junit.Test +import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLanguageServer +import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService +import software.aws.toolkits.jetbrains.services.amazonq.lsp.encryption.JwtEncryptionManager +import java.util.concurrent.CompletableFuture + +class DefaultAuthCredentialsServiceTest { + private lateinit var project: Project + private lateinit var mockLanguageServer: AmazonQLanguageServer + private lateinit var mockEncryptionManager: JwtEncryptionManager + private lateinit var sut: DefaultAuthCredentialsService + + @Before + fun setUp() { + project = mockk() + mockLanguageServer = mockk() + mockEncryptionManager = mockk() + every { mockEncryptionManager.encrypt(any()) } returns "mock-encrypted-data" + + // Mock the service methods on Project + val mockLspService = mockk() + every { project.getService(AmazonQLspService::class.java) } returns mockLspService + every { project.serviceIfCreated() } returns mockLspService + + // Mock the LSP service's executeSync method as a suspend function + every { + mockLspService.executeSync>(any()) + } coAnswers { + val func = firstArg CompletableFuture>() + func.invoke(mockLanguageServer) + } + + sut = DefaultAuthCredentialsService(project, this.mockEncryptionManager) + } + + @Test + fun `test updateTokenCredentials unencrypted success`() { + val token = "unencryptedToken" + val isEncrypted = false + + every { + mockLanguageServer.updateTokenCredentials(any()) + } returns CompletableFuture.completedFuture(ResponseMessage()) + + sut.updateTokenCredentials(token, isEncrypted) + + verify(exactly = 0) { + mockEncryptionManager.decrypt(any()) + } + verify(exactly = 1) { + mockLanguageServer.updateTokenCredentials(any()) + } + } + + @Test + fun `test updateTokenCredentials encrypted success`() { + val encryptedToken = "encryptedToken" + val decryptedToken = "decryptedToken" + val isEncrypted = true + + every { mockEncryptionManager.decrypt(encryptedToken) } returns decryptedToken + every { mockEncryptionManager.encrypt(any()) } returns "mock-encrypted-data" + every { + mockLanguageServer.updateTokenCredentials(any()) + } returns CompletableFuture.completedFuture(ResponseMessage()) + + sut.updateTokenCredentials(encryptedToken, isEncrypted) + + verify(exactly = 1) { mockEncryptionManager.decrypt(encryptedToken) } + verify(exactly = 1) { mockLanguageServer.updateTokenCredentials(any()) } + } + + @Test + fun `test deleteTokenCredentials success`() { + every { mockLanguageServer.deleteTokenCredentials() } returns CompletableFuture.completedFuture(Unit) + + sut.deleteTokenCredentials() + + verify(exactly = 1) { mockLanguageServer.deleteTokenCredentials() } + } +}