aws · breedloj · Feb 13, 2025 · Feb 12, 2025
diff --git a/plugin/src/software/aws/toolkits/eclipse/amazonq/lsp/auth/AuthCredentialsService.java b/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();
 }
diff --git a/plugin/src/software/aws/toolkits/eclipse/amazonq/lsp/auth/AuthTokenService.java b/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);
 }
diff --git a/plugin/src/software/aws/toolkits/eclipse/amazonq/lsp/auth/DefaultAuthCredentialsService.java b/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) {

diff --git a/plugin/src/software/aws/toolkits/eclipse/amazonq/lsp/auth/DefaultAuthTokenService.java b/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);

diff --git a/plugin/src/software/aws/toolkits/eclipse/amazonq/lsp/auth/DefaultLoginService.java b/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) {

diff --git a/plugin/src/software/aws/toolkits/eclipse/amazonq/lsp/auth/model/GetSsoTokenResult.java b/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) { }
diff --git a/...tst/software/aws/toolkits/eclipse/amazonq/lsp/auth/DefaultAuthCredentialsServiceTest.java b/...tst/software/aws/toolkits/eclipse/amazonq/lsp/auth/DefaultAuthCredentialsServiceTest.java
@@ -5,7 +5,6 @@
 
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoMoreInteractions;
@@ -18,19 +17,17 @@
 import org.junit.jupiter.api.Test;
 
 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 {
     private static DefaultAuthCredentialsService authCredentialsService;
     private static LspProvider mockLspProvider;
-    private static LspEncryptionManager mockedLspEncryptionManager;
     private static AmazonQLspServer mockedAmazonQServer;
 
     @BeforeEach
     public final void setUp() {
         mockLspProvider = mock(LspProvider.class);
-        mockedLspEncryptionManager = mock(LspEncryptionManager.class);
         mockedAmazonQServer = mock(AmazonQLspServer.class);
 
         resetAuthTokenService();
@@ -47,26 +44,21 @@ 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());
         verifyNoMoreInteractions(mockedAmazonQServer);
     }
 
     @Test
     void updateTokenCredentialsEncryptedSuccess() {
-        String encryptedToken = "encryptedToken";
-        String accessToken = "accessToken";
         boolean isEncrypted = true;
 
-        when(mockedLspEncryptionManager.decrypt(encryptedToken)).thenReturn(accessToken);
         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());
         verifyNoMoreInteractions(mockedAmazonQServer);
     }
@@ -82,7 +74,6 @@ void deleteTokenCredentialsSuccess() {
     private void resetAuthTokenService() {
         authCredentialsService = DefaultAuthCredentialsService.builder()
                 .withLspProvider(mockLspProvider)
-                .withEncryptionManager(mockedLspEncryptionManager)
                 .build();
         authCredentialsService = spy(authCredentialsService);
       }

diff --git a/plugin/tst/software/aws/toolkits/eclipse/amazonq/lsp/auth/DefaultAuthTokenServiceTest.java b/plugin/tst/software/aws/toolkits/eclipse/amazonq/lsp/auth/DefaultAuthTokenServiceTest.java
@@ -77,10 +77,10 @@ void getSsoTokenBuilderIdNoLoginOnInvalidTokenSuccess() throws Exception {
         when(mockSsoTokenResult.ssoToken()).thenReturn(expectedToken);
         boolean loginOnInvalidToken = false;
 
-        SsoToken actualToken = invokeGetSsoToken(loginType, loginParams, loginOnInvalidToken);
+        GetSsoTokenResult result = invokeGetSsoToken(loginType, loginParams, loginOnInvalidToken);
 
-        assertEquals(expectedToken.id(), actualToken.id());
-        assertEquals(expectedToken.accessToken(), actualToken.accessToken());
+        assertEquals(expectedToken.id(), result.ssoToken().id());
+        assertEquals(expectedToken.accessToken(), result.ssoToken().accessToken());
         verify(mockAmazonQServer).getSsoToken(any(GetSsoTokenParams.class));
         verifyNoMoreInteractions(mockAmazonQServer);
     }
@@ -93,10 +93,10 @@ void getSsoTokenBuilderIdWithLoginOnInvalidTokenSuccess() throws Exception {
         when(mockSsoTokenResult.ssoToken()).thenReturn(expectedToken);
         boolean loginOnInvalidToken = true;
 
-        SsoToken actualToken = invokeGetSsoToken(loginType, loginParams, loginOnInvalidToken);
+        GetSsoTokenResult result = invokeGetSsoToken(loginType, loginParams, loginOnInvalidToken);
 
-        assertEquals(expectedToken.id(), actualToken.id());
-        assertEquals(expectedToken.accessToken(), actualToken.accessToken());
+        assertEquals(expectedToken.id(), result.ssoToken().id());
+        assertEquals(expectedToken.accessToken(), result.ssoToken().accessToken());
         verify(mockAmazonQServer).getSsoToken(any(GetSsoTokenParams.class));
         verifyNoMoreInteractions(mockAmazonQServer);
     }
@@ -109,10 +109,10 @@ void getSsoTokenIDCNoLoginOnInvalidTokenSuccess() throws Exception {
         when(mockSsoTokenResult.ssoToken()).thenReturn(expectedToken);
         boolean loginOnInvalidToken = false;
 
-        SsoToken actualToken = invokeGetSsoToken(loginType, loginParams, loginOnInvalidToken);
+        GetSsoTokenResult result = invokeGetSsoToken(loginType, loginParams, loginOnInvalidToken);
 
-        assertEquals(expectedToken.id(), actualToken.id());
-        assertEquals(expectedToken.accessToken(), actualToken.accessToken());
+        assertEquals(expectedToken.id(), result.ssoToken().id());
+        assertEquals(expectedToken.accessToken(), result.ssoToken().accessToken());
         verify(mockAmazonQServer).getSsoToken(any(GetSsoTokenParams.class));
         verifyNoMoreInteractions(mockAmazonQServer);
     }
@@ -129,26 +129,26 @@ void getSsoTokenIDCWithLoginOnInvalidTokenSuccess() throws Exception {
         when(mockSsoTokenResult.ssoToken()).thenReturn(expectedToken);
         boolean loginOnInvalidToken = true;
 
-        SsoToken actualToken = invokeGetSsoToken(loginType, loginParams, loginOnInvalidToken);
+        GetSsoTokenResult result = invokeGetSsoToken(loginType, loginParams, loginOnInvalidToken);
 
-        assertEquals(expectedToken.id(), actualToken.id());
-        assertEquals(expectedToken.accessToken(), actualToken.accessToken());
+        assertEquals(expectedToken.id(), result.ssoToken().id());
+        assertEquals(expectedToken.accessToken(), result.ssoToken().accessToken());
         verify(mockAmazonQServer).updateProfile(updateProfileParamsCaptor.capture());
         UpdateProfileParams actualParams = updateProfileParamsCaptor.getValue();
         verifyUpdateProfileParams(actualParams);
         verify(mockAmazonQServer).getSsoToken(any(GetSsoTokenParams.class));
         verifyNoMoreInteractions(mockAmazonQServer);
     }
 
-    private SsoToken invokeGetSsoToken(final LoginType loginType, final LoginParams loginParams, final boolean loginOnInvalidToken) throws Exception {
+    private GetSsoTokenResult invokeGetSsoToken(final LoginType loginType, final LoginParams loginParams, final boolean loginOnInvalidToken) throws Exception {
         Object getSsoTokenFuture = authTokenService.getSsoToken(loginType, loginParams, loginOnInvalidToken);
         assertTrue(getSsoTokenFuture instanceof CompletableFuture<?>, "Return value should be CompletableFuture");
 
         CompletableFuture<?> future = (CompletableFuture<?>) getSsoTokenFuture;
         Object result = future.get();
-        assertTrue(result instanceof SsoToken, "getSsoTokenFuture result should be SsoToken");
+        assertTrue(result instanceof GetSsoTokenResult, "getSsoTokenFuture result should be GetSsoTokenResult");
 
-        return (SsoToken) result;
+        return (GetSsoTokenResult) result;
     }
 
     private LoginParams createValidLoginParams() {