Encrypt and decrypt chat lsp communication (#55)

* Initialize encryption server and implement LspEncryptionManager

* Send chat request with encryption and decryption

* Clean up Encryption classes

* Fix checkstyle errors

* Clean up Connection Provider

* Clean LspJsonWebTokenHandler

* Update class name to LspJsonWebToken

* Decrypt partial result notification

* Remove unused

* Remove singleton on LspEncryptionManager

* Revert "Remove singleton on LspEncryptionManager"

This reverts commit 9b6ea01.

* Update name of message to encryptedMessage

* Encrypt quick actions

* Fix encryption for chat and quick actions

* Respond to pr comments

* Fix checkstyle errors

* Fix typos and copyright

plugin/META-INF/MANIFEST.MF

@@ -93,6 +93,7 @@ Bundle-Classpath: .,
  target/dependency/netty-transport-4.1.108.Final.jar,
  target/dependency/netty-transport-classes-epoll-4.1.108.Final.jar,
  target/dependency/netty-transport-native-unix-common-4.1.108.Final.jar,
+ target/dependency/nimbus-jose-jwt-9.11.jar,
  target/dependency/org.apache.aries.spifly.dynamic.bundle-1.3.7.jar,
  target/dependency/org.apache.felix.scr-2.2.10.jar,
  target/dependency/org.eclipse.lsp4j-0.23.1.jar,

plugin/pom.xml

@@ -97,6 +97,11 @@
             <artifactId>junit-jupiter</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+		    <groupId>com.nimbusds</groupId>
+		    <artifactId>nimbus-jose-jwt</artifactId>
+		    <version>9.11</version>
+		</dependency>
     </dependencies>
 
     <build>

plugin/src/software/aws/toolkits/eclipse/amazonq/chat/ChatCommunicationManager.java

@@ -13,12 +13,15 @@
 import software.aws.toolkits.eclipse.amazonq.chat.models.ChatResult;
 import software.aws.toolkits.eclipse.amazonq.chat.models.ChatUIInboundCommand;
 import software.aws.toolkits.eclipse.amazonq.chat.models.ChatUIInboundCommandName;
+import software.aws.toolkits.eclipse.amazonq.chat.models.EncryptedChatParams;
+import software.aws.toolkits.eclipse.amazonq.chat.models.EncryptedQuickActionParams;
 import software.aws.toolkits.eclipse.amazonq.chat.models.GenericTabParams;
 import software.aws.toolkits.eclipse.amazonq.chat.models.QuickActionParams;
 import software.aws.toolkits.eclipse.amazonq.exception.AmazonQPluginException;
+import software.aws.toolkits.eclipse.amazonq.lsp.encryption.LspEncryptionManager;
 import software.aws.toolkits.eclipse.amazonq.util.JsonHandler;
 import software.aws.toolkits.eclipse.amazonq.util.PluginLogger;
-import software.aws.toolkits.eclipse.amazonq.util.ProgressNotficationUtils;
+import software.aws.toolkits.eclipse.amazonq.util.ProgressNotificationUtils;
 import software.aws.toolkits.eclipse.amazonq.views.ChatUiRequestListener;
 import software.aws.toolkits.eclipse.amazonq.views.model.Command;
 
@@ -35,12 +38,14 @@ public final class ChatCommunicationManager {
     private final JsonHandler jsonHandler;
     private final CompletableFuture<ChatMessageProvider> chatMessageProvider;
     private final ChatPartialResultMap chatPartialResultMap;
+    private final LspEncryptionManager lspEncryptionManager;
     private ChatUiRequestListener chatUiRequestListener;
 
     private ChatCommunicationManager() {
         this.jsonHandler = new JsonHandler();
         this.chatMessageProvider = ChatMessageProvider.createAsync();
         this.chatPartialResultMap = new ChatPartialResultMap();
+        this.lspEncryptionManager = LspEncryptionManager.getInstance();
     }
 
     public static synchronized ChatCommunicationManager getInstance() {
@@ -56,18 +61,28 @@ public CompletableFuture<ChatResult> sendMessageToChatServer(final Command comma
                 switch (command) {
                     case CHAT_SEND_PROMPT:
                         ChatRequestParams chatRequestParams = jsonHandler.convertObject(params, ChatRequestParams.class);
-                        return sendChatRequest(chatRequestParams.getTabId(), token -> {
-                            chatRequestParams.setPartialResultToken(token);
+                        return sendEncryptedChatMessage(chatRequestParams.getTabId(), token -> {
+                            String encryptedChatResult = lspEncryptionManager.encrypt(chatRequestParams);
 
-                            return chatMessageProvider.sendChatPrompt(chatRequestParams);
+                            EncryptedChatParams encryptedChatRequestParams = new EncryptedChatParams(
+                                encryptedChatResult,
+                                token
+                            );
+
+                            return chatMessageProvider.sendChatPrompt(chatRequestParams.getTabId(), encryptedChatRequestParams);
                         });
                     case CHAT_QUICK_ACTION:
                         QuickActionParams quickActionParams = jsonHandler.convertObject(params, QuickActionParams.class);
-                        return sendChatRequest(quickActionParams.getTabId(), token -> {
-                            quickActionParams.setPartialResultToken(token);
+                        return sendEncryptedChatMessage(quickActionParams.getTabId(), token -> {
+                            String encryptedChatResult = lspEncryptionManager.encrypt(quickActionParams);
+
+                            EncryptedQuickActionParams encryptedQuickActionParams = new EncryptedQuickActionParams(
+                                encryptedChatResult,
+                                token
+                            );
 
-                        return chatMessageProvider.sendQuickAction(quickActionParams);
-                    });
+                            return chatMessageProvider.sendQuickAction(encryptedQuickActionParams);
+                        });
                     case CHAT_READY:
                         chatMessageProvider.sendChatReady();
                         return CompletableFuture.completedFuture(null);
@@ -93,8 +108,8 @@ public CompletableFuture<ChatResult> sendMessageToChatServer(final Command comma
         });
     }
 
-    private CompletableFuture<ChatResult> sendChatRequest(final String tabId,
-            final Function<String, CompletableFuture<ChatResult>> action) {
+    private CompletableFuture<ChatResult> sendEncryptedChatMessage(final String tabId,
+            final Function<String, CompletableFuture<String>> action) {
         // Retrieving the chat result is expected to be a long-running process with
         // intermittent progress notifications being sent
         // from the LSP server. The progress notifications provide a token and a partial
@@ -103,10 +118,14 @@ private CompletableFuture<ChatResult> sendChatRequest(final String tabId,
         // a message for the UI.
         String partialResultToken = addPartialChatMessage(tabId);
 
-        return action.apply(partialResultToken).thenApply(result -> {
+        return action.apply(partialResultToken).thenApply(encryptedChatResult -> {
             // The mapping entry no longer needs to be maintained once the final result is
             // retrieved.
             removePartialChatMessage(partialResultToken);
+
+            String serializedData = lspEncryptionManager.decrypt(encryptedChatResult);
+            ChatResult result = jsonHandler.deserialize(serializedData, ChatResult.class);
+
             // show chat response in Chat UI
             ChatUIInboundCommand chatUIInboundCommand = new ChatUIInboundCommand(
                     ChatUIInboundCommandName.ChatPrompt.toString(), tabId, result, false);
@@ -140,7 +159,7 @@ public void sendMessageToChatUI(final ChatUIInboundCommand command) {
      * - Sends a partial chat prompt message to the webview.
      */
     public void handlePartialResultProgressNotification(final ProgressParams params) {
-        String token = ProgressNotficationUtils.getToken(params);
+        String token = ProgressNotificationUtils.getToken(params);
         String tabId = getPartialChatMessage(token);
 
         if (tabId == null || tabId.isEmpty()) {
@@ -152,7 +171,9 @@ public void handlePartialResultProgressNotification(final ProgressParams params)
             throw new AmazonQPluginException("Error occurred while handling partial result notification: expected Object value");
         }
 
-        ChatResult partialChatResult = ProgressNotficationUtils.getObject(params, ChatResult.class);
+        String encryptedPartialChatResult = ProgressNotificationUtils.getObject(params, String.class);
+        String serializedData = lspEncryptionManager.decrypt(encryptedPartialChatResult);
+        ChatResult partialChatResult = jsonHandler.deserialize(serializedData, ChatResult.class);
 
         // Check to ensure the body has content in order to keep displaying the spinner while loading
         if (partialChatResult.body() == null || partialChatResult.body().length() == 0) {

plugin/src/software/aws/toolkits/eclipse/amazonq/chat/ChatMessage.java

@@ -4,10 +4,9 @@
 
 import java.util.concurrent.CompletableFuture;
 
-import software.aws.toolkits.eclipse.amazonq.chat.models.ChatRequestParams;
-import software.aws.toolkits.eclipse.amazonq.chat.models.ChatResult;
+import software.aws.toolkits.eclipse.amazonq.chat.models.EncryptedChatParams;
+import software.aws.toolkits.eclipse.amazonq.chat.models.EncryptedQuickActionParams;
 import software.aws.toolkits.eclipse.amazonq.chat.models.GenericTabParams;
-import software.aws.toolkits.eclipse.amazonq.chat.models.QuickActionParams;
 import software.aws.toolkits.eclipse.amazonq.lsp.AmazonQLspServer;
 
 public final class ChatMessage {
@@ -17,11 +16,13 @@ public ChatMessage(final AmazonQLspServer amazonQLspServer) {
         this.amazonQLspServer = amazonQLspServer;
     }
 
-    public CompletableFuture<ChatResult> sendChatPrompt(final ChatRequestParams chatRequestParams) {
-        return amazonQLspServer.sendChatPrompt(chatRequestParams);
+    // Returns a ChatResult as an encrypted message {@link LspEncryptionManager#decrypt()}
+    public CompletableFuture<String> sendChatPrompt(final EncryptedChatParams params) {
+        return amazonQLspServer.sendChatPrompt(params);
     }
 
-    public CompletableFuture<ChatResult> sendQuickAction(final QuickActionParams params) {
+    // Returns a ChatResult as an encrypted message {@link LspEncryptionManager#decrypt()}
+    public CompletableFuture<String> sendQuickAction(final EncryptedQuickActionParams params) {
         return amazonQLspServer.sendQuickAction(params);
     }
 

plugin/src/software/aws/toolkits/eclipse/amazonq/chat/ChatMessageProvider.java

@@ -6,10 +6,9 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
 
-import software.aws.toolkits.eclipse.amazonq.chat.models.ChatRequestParams;
-import software.aws.toolkits.eclipse.amazonq.chat.models.ChatResult;
+import software.aws.toolkits.eclipse.amazonq.chat.models.EncryptedChatParams;
+import software.aws.toolkits.eclipse.amazonq.chat.models.EncryptedQuickActionParams;
 import software.aws.toolkits.eclipse.amazonq.chat.models.GenericTabParams;
-import software.aws.toolkits.eclipse.amazonq.chat.models.QuickActionParams;
 import software.aws.toolkits.eclipse.amazonq.lsp.AmazonQLspServer;
 import software.aws.toolkits.eclipse.amazonq.providers.LspProvider;
 
@@ -18,7 +17,7 @@ public final class ChatMessageProvider {
     private final AmazonQLspServer amazonQLspServer;
     // Map of in-flight requests per tab Ids
     // TODO ECLIPSE-349: Handle disposing resources of this class including this map
-    private Map<String, CompletableFuture<ChatResult>> inflightRequestByTabId = new ConcurrentHashMap<String, CompletableFuture<ChatResult>>();
+    private Map<String, CompletableFuture<String>> inflightRequestByTabId = new ConcurrentHashMap<String, CompletableFuture<String>>();
 
     public static CompletableFuture<ChatMessageProvider> createAsync() {
         return LspProvider.getAmazonQServer()
@@ -29,23 +28,23 @@ private ChatMessageProvider(final AmazonQLspServer amazonQLspServer) {
         this.amazonQLspServer = amazonQLspServer;
     }
 
-    public CompletableFuture<ChatResult> sendChatPrompt(final ChatRequestParams chatRequestParams) {
+    public CompletableFuture<String> sendChatPrompt(final String tabId, final EncryptedChatParams encryptedChatRequestParams) {
         ChatMessage chatMessage = new ChatMessage(amazonQLspServer);
 
-        var response = chatMessage.sendChatPrompt(chatRequestParams);
+        var response = chatMessage.sendChatPrompt(encryptedChatRequestParams);
         // We assume there is only one outgoing request per tab because the input is
         // blocked when there is an outgoing request
-        inflightRequestByTabId.put(chatRequestParams.getTabId(), response);
+        inflightRequestByTabId.put(tabId, response);
         response.whenComplete((result, exception) -> {
             // stop tracking in-flight requests once response is received
-            inflightRequestByTabId.remove(chatRequestParams.getTabId());
+            inflightRequestByTabId.remove(tabId);
         });
         return response;
     }
 
-    public CompletableFuture<ChatResult> sendQuickAction(final QuickActionParams quickActionParams) {
+    public CompletableFuture<String> sendQuickAction(final EncryptedQuickActionParams encryptedQuickActionParams) {
         ChatMessage chatMessage = new ChatMessage(amazonQLspServer);
-        return chatMessage.sendQuickAction(quickActionParams);
+        return chatMessage.sendQuickAction(encryptedQuickActionParams);
     }
 
     public void sendChatReady() {

plugin/src/software/aws/toolkits/eclipse/amazonq/chat/models/EncryptedChatParams.java

@@ -0,0 +1,10 @@
+// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+
+package software.aws.toolkits.eclipse.amazonq.chat.models;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+public record EncryptedChatParams(
+        @JsonProperty("message") String message, // Message as encrypted jwt
+        @JsonProperty("partialResultToken") String partialResultToken) {
+}

plugin/src/software/aws/toolkits/eclipse/amazonq/chat/models/EncryptedQuickActionParams.java

@@ -0,0 +1,10 @@
+// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+
+package software.aws.toolkits.eclipse.amazonq.chat.models;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+public record EncryptedQuickActionParams(
+        @JsonProperty("message") String message, // Message as encrypted jwt
+        @JsonProperty("partialResultToken") String partialResultToken) {
+}

plugin/src/software/aws/toolkits/eclipse/amazonq/lsp/AmazonQLspServer.java

@@ -10,10 +10,9 @@
 import org.eclipse.lsp4j.jsonrpc.services.JsonNotification;
 import org.eclipse.lsp4j.services.LanguageServer;
 
-import software.aws.toolkits.eclipse.amazonq.chat.models.ChatRequestParams;
-import software.aws.toolkits.eclipse.amazonq.chat.models.ChatResult;
+import software.aws.toolkits.eclipse.amazonq.chat.models.EncryptedChatParams;
+import software.aws.toolkits.eclipse.amazonq.chat.models.EncryptedQuickActionParams;
 import software.aws.toolkits.eclipse.amazonq.chat.models.GenericTabParams;
-import software.aws.toolkits.eclipse.amazonq.chat.models.QuickActionParams;
 import software.aws.toolkits.eclipse.amazonq.lsp.model.GetConfigurationFromServerParams;
 import software.aws.toolkits.eclipse.amazonq.lsp.model.InlineCompletionParams;
 import software.aws.toolkits.eclipse.amazonq.lsp.model.InlineCompletionResponse;
@@ -26,10 +25,10 @@ public interface AmazonQLspServer extends LanguageServer {
     CompletableFuture<InlineCompletionResponse> inlineCompletionWithReferences(InlineCompletionParams params);
 
     @JsonRequest("aws/chat/sendChatPrompt")
-    CompletableFuture<ChatResult> sendChatPrompt(ChatRequestParams chatRequestParams);
+    CompletableFuture<String> sendChatPrompt(EncryptedChatParams encryptedChatRequestParams);
 
     @JsonRequest("aws/chat/sendChatQuickAction")
-    CompletableFuture<ChatResult> sendQuickAction(QuickActionParams quickActionParams);
+    CompletableFuture<String> sendQuickAction(EncryptedQuickActionParams encryptedQuickActionParams);
 
     @JsonNotification("aws/chat/tabAdd")
     void tabAdd(GenericTabParams params);

plugin/src/software/aws/toolkits/eclipse/amazonq/lsp/connection/QLspConnectionProvider.java

@@ -4,12 +4,15 @@
 package software.aws.toolkits.eclipse.amazonq.lsp.connection;
 
 import java.io.IOException;
+import java.io.OutputStream;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 
+import software.aws.toolkits.eclipse.amazonq.lsp.encryption.LspEncryptionManager;
 import software.aws.toolkits.eclipse.amazonq.lsp.manager.LspManager;
 import software.aws.toolkits.eclipse.amazonq.providers.LspManagerProvider;
+import software.aws.toolkits.eclipse.amazonq.util.PluginLogger;
 
 public class QLspConnectionProvider extends AbstractLspConnectionProvider {
 
@@ -22,6 +25,7 @@ public QLspConnectionProvider() throws IOException {
         commands.add("--nolazy");
         commands.add("--inspect=5599");
         commands.add("--stdio");
+        commands.add("--set-credentials-encryption-key");
         setCommands(commands);
     }
 
@@ -31,4 +35,19 @@ protected final void addEnvironmentVariables(final Map<String, String> env) {
         env.put("ENABLE_TOKEN_PROVIDER", "true");
     }
 
+    @Override
+    public final void start() throws IOException {
+        super.start();
+
+        PluginLogger.info("Initializing encrypted communication with Amazon Q Lsp Server");
+
+        try {
+            LspEncryptionManager lspEncryption = LspEncryptionManager.getInstance();
+            OutputStream serverStdIn = getOutputStream();
+
+            lspEncryption.initializeEncrypedCommunication(serverStdIn);
+        } catch (Exception e) {
+            PluginLogger.error("Error occured while initializing encrypted communication with Amazon Q Lsp Server", e);
+        }
+    }
 }

plugin/src/software/aws/toolkits/eclipse/amazonq/lsp/encryption/LspEncryptionKey.java

@@ -0,0 +1,40 @@
+// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+
+package software.aws.toolkits.eclipse.amazonq.lsp.encryption;
+
+import java.util.Base64;
+
+import javax.crypto.KeyGenerator;
+import javax.crypto.SecretKey;
+
+import software.aws.toolkits.eclipse.amazonq.exception.AmazonQPluginException;
+
+public final class LspEncryptionKey {
+    private SecretKey key;
+
+    public LspEncryptionKey() {
+        this.key = generateKey();
+    }
+
+    public SecretKey getKey() {
+        return key;
+    }
+
+    public String getKeyAsBase64() {
+        return base64Encode(key);
+    }
+
+    private String base64Encode(final SecretKey key) {
+        return Base64.getEncoder().encodeToString(key.getEncoded());
+    }
+
+    public static SecretKey generateKey() {
+        try {
+            KeyGenerator keyGen = KeyGenerator.getInstance("AES");
+            keyGen.init(256);
+            return keyGen.generateKey();
+        } catch (Exception e) {
+            throw new AmazonQPluginException("Error occurred while generating lsp encryption key", e);
+        }
+    }
+}