Skip to content

Commit

Permalink
Pass through auth token to Q server without decrypting (#363)
Browse files Browse the repository at this point in the history
  • Loading branch information
@breedloj
breedloj authored Feb 13, 2025
1 parent b9d5bdc commit a87f46b
Show file tree
Hide file tree
Showing 9 changed files with 136 additions and 202 deletions.
4 changes: 3 additions & 1 deletion plugin/src/software/aws/toolkits/eclipse/amazonq/lsp/auth/AuthCredentialsService.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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();
}
4 changes: 2 additions & 2 deletions plugin/src/software/aws/toolkits/eclipse/amazonq/lsp/auth/AuthTokenService.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
33 changes: 2 additions & 31 deletions plugin/src/software/aws/toolkits/eclipse/amazonq/lsp/auth/DefaultAuthCredentialsService.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
});
Expand All @@ -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) {
Expand Down
6 changes: 3 additions & 3 deletions plugin/src/software/aws/toolkits/eclipse/amazonq/lsp/auth/DefaultAuthTokenService.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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()
Expand Down Expand Up @@ -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);
Expand Down
24 changes: 7 additions & 17 deletions plugin/src/software/aws/toolkits/eclipse/amazonq/lsp/auth/DefaultLoginService.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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");
Expand Down Expand Up @@ -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);
Expand All @@ -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;
Expand All @@ -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;
Expand All @@ -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) {
Expand Down
4 changes: 3 additions & 1 deletion plugin/src/software/aws/toolkits/eclipse/amazonq/lsp/auth/model/GetSsoTokenResult.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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) { }
15 changes: 3 additions & 12 deletions ...tst/software/aws/toolkits/eclipse/amazonq/lsp/auth/DefaultAuthCredentialsServiceTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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();
Expand All @@ -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);
}
Expand All @@ -82,7 +74,6 @@ void deleteTokenCredentialsSuccess() {
private void resetAuthTokenService() {
authCredentialsService = DefaultAuthCredentialsService.builder()
.withLspProvider(mockLspProvider)
.withEncryptionManager(mockedLspEncryptionManager)
.build();
authCredentialsService = spy(authCredentialsService);
}
Expand Down
30 changes: 15 additions & 15 deletions plugin/tst/software/aws/toolkits/eclipse/amazonq/lsp/auth/DefaultAuthTokenServiceTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Expand All @@ -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);
}
Expand All @@ -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);
}
Expand All @@ -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() {
Expand Down
Loading

0 comments on commit a87f46b

Please sign in to comment.