Pass through auth token to Q server without decrypting

plugin/src/software/aws/toolkits/eclipse/amazonq/lsp/auth/AuthCredentialsService.java

@@ -7,7 +7,9 @@
 
 import org.eclipse.lsp4j.jsonrpc.messages.ResponseMessage;
 
+import software.aws.toolkits.eclipse.amazonq.lsp.model.UpdateCredentialsPayload;
+
 public interface AuthCredentialsService {
-    CompletableFuture<ResponseMessage> updateTokenCredentials(String accessToken, boolean isEncrypted);
+    CompletableFuture<ResponseMessage> updateTokenCredentials(UpdateCredentialsPayload params);
     CompletableFuture<Void> deleteTokenCredentials();
 }

plugin/src/software/aws/toolkits/eclipse/amazonq/lsp/auth/AuthTokenService.java

@@ -5,13 +5,13 @@
 
 import java.util.concurrent.CompletableFuture;
 
+import software.aws.toolkits.eclipse.amazonq.lsp.auth.model.GetSsoTokenResult;
 import software.aws.toolkits.eclipse.amazonq.lsp.auth.model.InvalidateSsoTokenParams;
 import software.aws.toolkits.eclipse.amazonq.lsp.auth.model.InvalidateSsoTokenResult;
 import software.aws.toolkits.eclipse.amazonq.lsp.auth.model.LoginParams;
 import software.aws.toolkits.eclipse.amazonq.lsp.auth.model.LoginType;
-import software.aws.toolkits.eclipse.amazonq.lsp.auth.model.SsoToken;
 
 public interface AuthTokenService {
-    CompletableFuture<SsoToken> getSsoToken(LoginType loginType, LoginParams loginParams, boolean loginOnInvalidToken);
+    CompletableFuture<GetSsoTokenResult> getSsoToken(LoginType loginType, LoginParams loginParams, boolean loginOnInvalidToken);
     CompletableFuture<InvalidateSsoTokenResult> invalidateSsoToken(InvalidateSsoTokenParams invalidateSsoTokenParams);
 }

plugin/src/software/aws/toolkits/eclipse/amazonq/lsp/auth/DefaultAuthCredentialsService.java

@@ -9,35 +9,25 @@
 import org.eclipse.lsp4j.jsonrpc.messages.ResponseMessage;
 
 import software.aws.toolkits.eclipse.amazonq.exception.AmazonQPluginException;
-import software.aws.toolkits.eclipse.amazonq.lsp.encryption.LspEncryptionManager;
-import software.aws.toolkits.eclipse.amazonq.lsp.model.BearerCredentials;
 import software.aws.toolkits.eclipse.amazonq.lsp.model.UpdateCredentialsPayload;
-import software.aws.toolkits.eclipse.amazonq.lsp.model.UpdateCredentialsPayloadData;
 import software.aws.toolkits.eclipse.amazonq.plugin.Activator;
 import software.aws.toolkits.eclipse.amazonq.providers.LspProvider;
 
 public final class DefaultAuthCredentialsService implements AuthCredentialsService {
     private LspProvider lspProvider;
-    private LspEncryptionManager encryptionManager;
 
     private DefaultAuthCredentialsService(final Builder builder) {
         this.lspProvider = Objects.requireNonNull(builder.lspProvider, "lspProvider must not be null");
-        this.encryptionManager = Objects.requireNonNull(builder.encryptionManager, "encryptionManager must not be null");
     }
 
     public static Builder builder() {
         return new Builder();
     }
 
     @Override
-    public CompletableFuture<ResponseMessage> updateTokenCredentials(final String accessToken, final boolean isEncrypted) {
-        String token = accessToken;
-        if (isEncrypted) {
-            token = decryptSsoToken(accessToken);
-        }
-        UpdateCredentialsPayload payload = createUpdateCredentialsPayload(token);
+    public CompletableFuture<ResponseMessage> updateTokenCredentials(final UpdateCredentialsPayload params) {
         return lspProvider.getAmazonQServer()
-                .thenCompose(server -> server.updateTokenCredentials(payload))
+                .thenCompose(server -> server.updateTokenCredentials(params))
                 .exceptionally(throwable -> {
                     throw new AmazonQPluginException("Failed to update token credentials", throwable);
                 });
@@ -52,32 +42,13 @@ public CompletableFuture<Void> deleteTokenCredentials() {
                 });
     }
 
-    private String decryptSsoToken(final String encryptedSsoToken) {
-        String decryptedToken = encryptionManager.decrypt(encryptedSsoToken);
-        return decryptedToken.substring(1, decryptedToken.length() - 1); // Remove extra quotes surrounding token
-    }
-
-    private UpdateCredentialsPayload createUpdateCredentialsPayload(final String ssoToken) {
-        BearerCredentials credentials = new BearerCredentials();
-        credentials.setToken(ssoToken);
-
-        UpdateCredentialsPayloadData data = new UpdateCredentialsPayloadData(credentials);
-        String encryptedData = encryptionManager.encrypt(data);
-        return new UpdateCredentialsPayload(encryptedData, true);
-    }
-
     public static class Builder {
         private LspProvider lspProvider;
-        private LspEncryptionManager encryptionManager;
 
         public final Builder withLspProvider(final LspProvider lspProvider) {
             this.lspProvider = lspProvider;
             return this;
         }
-        public final Builder withEncryptionManager(final LspEncryptionManager encryptionManager) {
-            this.encryptionManager = encryptionManager;
-            return this;
-        }
 
         public final DefaultAuthCredentialsService build() {
             if (lspProvider == null) {

plugin/src/software/aws/toolkits/eclipse/amazonq/lsp/auth/DefaultAuthTokenService.java

@@ -13,6 +13,7 @@
 import software.aws.toolkits.eclipse.amazonq.exception.AmazonQPluginException;
 import software.aws.toolkits.eclipse.amazonq.lsp.auth.model.GetSsoTokenOptions;
 import software.aws.toolkits.eclipse.amazonq.lsp.auth.model.GetSsoTokenParams;
+import software.aws.toolkits.eclipse.amazonq.lsp.auth.model.GetSsoTokenResult;
 import software.aws.toolkits.eclipse.amazonq.lsp.auth.model.GetSsoTokenSource;
 import software.aws.toolkits.eclipse.amazonq.lsp.auth.model.InvalidateSsoTokenParams;
 import software.aws.toolkits.eclipse.amazonq.lsp.auth.model.InvalidateSsoTokenResult;
@@ -22,7 +23,6 @@
 import software.aws.toolkits.eclipse.amazonq.lsp.auth.model.ProfileSettings;
 import software.aws.toolkits.eclipse.amazonq.lsp.auth.model.SsoSession;
 import software.aws.toolkits.eclipse.amazonq.lsp.auth.model.SsoSessionSettings;
-import software.aws.toolkits.eclipse.amazonq.lsp.auth.model.SsoToken;
 import software.aws.toolkits.eclipse.amazonq.lsp.auth.model.UpdateProfileOptions;
 import software.aws.toolkits.eclipse.amazonq.lsp.auth.model.UpdateProfileParams;
 import software.aws.toolkits.eclipse.amazonq.plugin.Activator;
@@ -41,7 +41,7 @@ public static Builder builder() {
     }
 
     @Override
-    public CompletableFuture<SsoToken> getSsoToken(final LoginType loginType, final LoginParams loginParams,
+    public CompletableFuture<GetSsoTokenResult> getSsoToken(final LoginType loginType, final LoginParams loginParams,
                 final boolean loginOnInvalidToken) {
         GetSsoTokenParams getSsoTokenParams = createGetSsoTokenParams(loginType, loginOnInvalidToken);
         return lspProvider.getAmazonQServer()
@@ -70,7 +70,7 @@ public CompletableFuture<SsoToken> getSsoToken(final LoginType loginType, final
                 })
                 .thenCompose(server -> server.getSsoToken(getSsoTokenParams))
                 .thenApply(response -> {
-                    return response.ssoToken();
+                    return response;
                 })
                 .exceptionally(throwable -> {
                     throw new AmazonQPluginException("Failed to fetch SSO token", throwable);

plugin/src/software/aws/toolkits/eclipse/amazonq/lsp/auth/DefaultLoginService.java

@@ -15,8 +15,7 @@
 import software.aws.toolkits.eclipse.amazonq.lsp.auth.model.InvalidateSsoTokenParams;
 import software.aws.toolkits.eclipse.amazonq.lsp.auth.model.LoginParams;
 import software.aws.toolkits.eclipse.amazonq.lsp.auth.model.LoginType;
-import software.aws.toolkits.eclipse.amazonq.lsp.encryption.DefaultLspEncryptionManager;
-import software.aws.toolkits.eclipse.amazonq.lsp.encryption.LspEncryptionManager;
+import software.aws.toolkits.eclipse.amazonq.lsp.model.UpdateCredentialsPayload;
 import software.aws.toolkits.eclipse.amazonq.providers.LspProvider;
 import software.aws.toolkits.eclipse.amazonq.util.AuthUtil;
 import software.aws.toolkits.eclipse.amazonq.plugin.Activator;
@@ -115,7 +114,7 @@ public CompletableFuture<Void> logout() {
     public CompletableFuture<Void> expire() {
         Activator.getLogger().info("Attempting to expire credentials...");
 
-        return authCredentialsService.updateTokenCredentials(null, false)
+        return authCredentialsService.updateTokenCredentials(new UpdateCredentialsPayload(null, false))
                 .thenRun(() -> {
                     authStateManager.toExpired();
                     Activator.getLogger().info("Successfully expired credentials");
@@ -157,16 +156,16 @@ CompletableFuture<Void> processLogin(final LoginType loginType, final LoginParam
 
         return authTokenService.getSsoToken(loginType, loginParams, loginOnInvalidToken)
                 .thenApply(ssoToken -> {
-                    ssoTokenId.set(ssoToken.id());
+                    ssoTokenId.set(ssoToken.ssoToken().id());
                     return ssoToken;
                 })
                 .thenAccept(ssoToken -> {
-                    authCredentialsService.updateTokenCredentials(ssoToken.accessToken(), true);
+                    authCredentialsService.updateTokenCredentials(ssoToken.updateCredentialsParams());
                 })
                 .thenRun(() -> {
-                  authStateManager.toLoggedIn(loginType, loginParams, ssoTokenId.get());
-                  Activator.getLogger().info("Successfully logged in");
-                  CustomizationUtil.triggerChangeConfigurationNotification();
+                    authStateManager.toLoggedIn(loginType, loginParams, ssoTokenId.get());
+                    Activator.getLogger().info("Successfully logged in");
+                    CustomizationUtil.triggerChangeConfigurationNotification();
               })
               .exceptionally(throwable -> {
                   throw new AmazonQPluginException("Failed to process log in", throwable);
@@ -176,7 +175,6 @@ CompletableFuture<Void> processLogin(final LoginType loginType, final LoginParam
     public static class Builder {
         private LspProvider lspProvider;
         private PluginStore pluginStore;
-        private LspEncryptionManager encryptionManager;
         private AuthStateManager authStateManager;
         private AuthCredentialsService authCredentialsService;
         private AuthTokenService authTokenService;
@@ -190,10 +188,6 @@ public final Builder withPluginStore(final PluginStore pluginStore) {
             this.pluginStore = pluginStore;
             return this;
         }
-        public final Builder withEncryptionManager(final LspEncryptionManager encryptionManager) {
-            this.encryptionManager = encryptionManager;
-            return this;
-        }
         public final Builder withAuthStateManager(final AuthStateManager authStateManager) {
             this.authStateManager = authStateManager;
             return this;
@@ -217,16 +211,12 @@ public final DefaultLoginService build() {
             if (pluginStore == null) {
                 pluginStore = Activator.getPluginStore();
             }
-            if (encryptionManager == null) {
-                encryptionManager = DefaultLspEncryptionManager.getInstance();
-            }
             if (authStateManager == null) {
                 authStateManager = new DefaultAuthStateManager(pluginStore);
             }
             if (authCredentialsService == null) {
                 authCredentialsService = DefaultAuthCredentialsService.builder()
                         .withLspProvider(lspProvider)
-                        .withEncryptionManager(encryptionManager)
                         .build();
             }
             if (authTokenService == null) {

plugin/src/software/aws/toolkits/eclipse/amazonq/lsp/auth/model/GetSsoTokenResult.java

@@ -3,4 +3,6 @@
 
 package software.aws.toolkits.eclipse.amazonq.lsp.auth.model;
 
-public record GetSsoTokenResult(SsoToken ssoToken) { }
+import software.aws.toolkits.eclipse.amazonq.lsp.model.UpdateCredentialsPayload;
+
+public record GetSsoTokenResult(SsoToken ssoToken, UpdateCredentialsPayload updateCredentialsParams) { }

plugin/tst/software/aws/toolkits/eclipse/amazonq/lsp/auth/DefaultAuthCredentialsServiceTest.java

@@ -19,6 +19,7 @@
 
 import software.aws.toolkits.eclipse.amazonq.lsp.AmazonQLspServer;
 import software.aws.toolkits.eclipse.amazonq.lsp.encryption.LspEncryptionManager;
+import software.aws.toolkits.eclipse.amazonq.lsp.model.UpdateCredentialsPayload;
 import software.aws.toolkits.eclipse.amazonq.providers.LspProvider;
 
 public class DefaultAuthCredentialsServiceTest {
@@ -47,7 +48,7 @@ void updateTokenCredentialsUnencryptedSuccess() {
         when(mockedAmazonQServer.updateTokenCredentials(any()))
             .thenReturn(CompletableFuture.completedFuture(new ResponseMessage()));
 
-        authCredentialsService.updateTokenCredentials(accessToken, isEncrypted);
+        authCredentialsService.updateTokenCredentials(new UpdateCredentialsPayload(accessToken, isEncrypted));
 
         verify(mockedLspEncryptionManager, never()).decrypt(accessToken);
         verify(mockedAmazonQServer).updateTokenCredentials(any());
@@ -64,7 +65,7 @@ void updateTokenCredentialsEncryptedSuccess() {
         when(mockedAmazonQServer.updateTokenCredentials(any()))
             .thenReturn(CompletableFuture.completedFuture(new ResponseMessage()));
 
-        authCredentialsService.updateTokenCredentials("encryptedToken", isEncrypted);
+        authCredentialsService.updateTokenCredentials(new UpdateCredentialsPayload("encryptedToken", isEncrypted));
 
         verify(mockedLspEncryptionManager).decrypt(encryptedToken);
         verify(mockedAmazonQServer).updateTokenCredentials(any());
@@ -82,7 +83,6 @@ void deleteTokenCredentialsSuccess() {
     private void resetAuthTokenService() {
         authCredentialsService = DefaultAuthCredentialsService.builder()
                 .withLspProvider(mockLspProvider)
-                .withEncryptionManager(mockedLspEncryptionManager)
                 .build();
         authCredentialsService = spy(authCredentialsService);
       }