Re-resolve SSO access token on each credential refresh instead of caching it at construction time

.changes/next-release/bugfix-AWSSDKforJavav2-0f86113.json

@@ -0,0 +1,6 @@
+{
+    "type": "bugfix",
+    "category": "AWS SDK for Java v2",
+    "contributor": "",
+    "description": "Re-resolve SSO access token on each credential refresh in the SSOCredentialsProvider instead of caching it at construction time ensuring that refreshed tokens (for example from running `aws sso login`) are always used."
+}

services/sso/src/main/java/software/amazon/awssdk/services/sso/auth/ExpiredTokenException.java

@@ -31,6 +31,10 @@
 @SdkPublicApi
 public final class ExpiredTokenException extends SdkClientException {
 
+    public static final String DEFAULT_MESSAGE =
+        "The SSO session associated with this profile has expired or is otherwise invalid."
+        + " To refresh this SSO session run aws sso login with the corresponding profile.";
+
     private static final List<SdkField<?>> SDK_FIELDS = Collections.unmodifiableList(Arrays.asList());
 
     private ExpiredTokenException(Builder b) {
@@ -88,6 +92,9 @@ public BuilderImpl writableStackTrace(Boolean writableStackTrace) {
 
         @Override
         public ExpiredTokenException build() {
+            if (this.message == null) {
+                this.message = DEFAULT_MESSAGE;
+            }
             return new ExpiredTokenException(this);
         }
 

services/sso/src/main/java/software/amazon/awssdk/services/sso/auth/SsoProfileCredentialsProviderFactory.java

@@ -32,6 +32,7 @@
 import software.amazon.awssdk.auth.token.credentials.SdkToken;
 import software.amazon.awssdk.auth.token.credentials.SdkTokenProvider;
 import software.amazon.awssdk.auth.token.internal.LazyTokenProvider;
+import software.amazon.awssdk.core.exception.SdkServiceException;
 import software.amazon.awssdk.profiles.Profile;
 import software.amazon.awssdk.profiles.ProfileFile;
 import software.amazon.awssdk.profiles.ProfileProperty;
@@ -63,7 +64,7 @@ public class SsoProfileCredentialsProviderFactory implements ProfileCredentialsP
      */
     @Override
     public AwsCredentialsProvider create(ProfileProviderCredentialsContext credentialsContext) {
-        return new SsoProfileCredentialsProvider(credentialsContext, sdkTokenProvider(credentialsContext));
+        return new SsoProfileCredentialsProvider(credentialsContext, sdkTokenProvider(credentialsContext), null);
     }
 
     /**
@@ -73,7 +74,18 @@ public AwsCredentialsProvider create(ProfileProviderCredentialsContext credentia
     @SdkTestInternalApi
     public AwsCredentialsProvider create(ProfileProviderCredentialsContext credentialsContext,
                                          SdkTokenProvider tokenProvider) {
-        return new SsoProfileCredentialsProvider(credentialsContext, tokenProvider);
+        return new SsoProfileCredentialsProvider(credentialsContext, tokenProvider, null);
+    }
+
+    /**
+     * Alternative method to create the {@link SsoProfileCredentialsProvider} with a customized {@link SdkTokenProvider}
+     * and {@link SsoClient}. This method is only used for testing.
+     */
+    @SdkTestInternalApi
+    public AwsCredentialsProvider create(ProfileProviderCredentialsContext credentialsContext,
+                                         SdkTokenProvider tokenProvider,
+                                         SsoClient ssoClient) {
+        return new SsoProfileCredentialsProvider(credentialsContext, tokenProvider, ssoClient);
     }
 
     /**
@@ -87,30 +99,34 @@ private static final class SsoProfileCredentialsProvider implements AwsCredentia
         private final SsoCredentialsProvider credentialsProvider;
 
         private SsoProfileCredentialsProvider(ProfileProviderCredentialsContext credentialsContext,
-                                              SdkTokenProvider tokenProvider) {
+                                              SdkTokenProvider tokenProvider,
+                                              SsoClient ssoClient) {
             Profile profile = credentialsContext.profile();
             String ssoAccountId = profile.properties().get(ProfileProperty.SSO_ACCOUNT_ID);
             String ssoRoleName = profile.properties().get(ProfileProperty.SSO_ROLE_NAME);
             String ssoRegion = regionFromProfileOrSession(profile, credentialsContext.profileFile());
 
-            this.ssoClient = SsoClient.builder()
-                                      .credentialsProvider(AnonymousCredentialsProvider.create())
-                                      .region(Region.of(ssoRegion))
-                                      .build();
+            this.ssoClient = ssoClient != null
+                             ? ssoClient
+                             : SsoClient.builder()
+                                        .credentialsProvider(AnonymousCredentialsProvider.create())
+                                        .region(Region.of(ssoRegion))
+                                        .build();
 
             GetRoleCredentialsRequest request = GetRoleCredentialsRequest.builder()
                                                                          .accountId(ssoAccountId)
                                                                          .roleName(ssoRoleName)
                                                                          .build();
-            SdkToken sdkToken = tokenProvider.resolveToken();
-            Validate.paramNotNull(sdkToken, "Token provided by the TokenProvider is null");
-            Supplier<GetRoleCredentialsRequest> supplier = () -> request.toBuilder()
-                                                                        .accessToken(sdkToken.token())
-                                                                        .build();
+            Supplier<GetRoleCredentialsRequest> supplier = () -> {
+                SdkToken token = resolveTokenOrThrow(tokenProvider);
+                return request.toBuilder()
+                              .accessToken(token.token())
+                              .build();
+            };
 
 
             this.credentialsProvider = SsoCredentialsProvider.builder()
-                                                             .ssoClient(ssoClient)
+                                                             .ssoClient(this.ssoClient)
                                                              .refreshRequest(supplier)
                                                              .sourceChain(credentialsContext.sourceChain())
                                                              .build();
@@ -127,6 +143,23 @@ public void close() {
             IoUtils.closeQuietly(ssoClient, null);
         }
 
+        private static SdkToken resolveTokenOrThrow(SdkTokenProvider tokenProvider) {
+            SdkToken token;
+            try {
+                token = tokenProvider.resolveToken();
+            } catch (ExpiredTokenException | SdkServiceException e) {
+                throw e;
+            } catch (RuntimeException e) {
+                throw ExpiredTokenException.builder()
+                                           .cause(e)
+                                           .build();
+            }
+            if (token == null || token.token() == null) {
+                throw ExpiredTokenException.builder().build();
+            }
+            return token;
+        }
+
         private static String regionFromProfileOrSession(Profile profile, ProfileFile profileFile) {
             Optional<String> ssoSession = profile.property(ProfileSection.SSO_SESSION.getPropertyKeyName());
             String profileRegion = profile.properties().get(ProfileProperty.SSO_REGION);

services/sso/src/main/java/software/amazon/awssdk/services/sso/internal/SsoAccessTokenProvider.java

@@ -29,7 +29,6 @@
 import software.amazon.awssdk.services.sso.auth.ExpiredTokenException;
 import software.amazon.awssdk.services.sso.auth.SsoCredentialsProvider;
 import software.amazon.awssdk.utils.IoUtils;
-import software.amazon.awssdk.utils.Validate;
 
 /**
  * Resolve the access token from the cached token file. If the token has expired then throw out an exception to ask the users to
@@ -48,7 +47,15 @@ public SsoAccessTokenProvider(Path cachedTokenFilePath) {
 
     @Override
     public SdkToken resolveToken() {
-        return tokenFromFile();
+        try {
+            return tokenFromFile();
+        } catch (ExpiredTokenException e) {
+            throw e;
+        } catch (Exception e) {
+            throw ExpiredTokenException.builder()
+                                       .cause(e)
+                                       .build();
+        }
     }
 
     private SdkToken tokenFromFile() {
@@ -61,25 +68,22 @@ private SdkToken tokenFromFile() {
 
     private SdkToken getTokenFromJson(String json) {
         JsonNode jsonNode = PARSER.parse(json);
-        String expiration = jsonNode.field("expiresAt").map(JsonNode::text).orElse(null);
+        String expirationStr = jsonNode.field("expiresAt").map(JsonNode::text).orElse(null);
+
+        if (expirationStr == null) {
+            throw ExpiredTokenException.builder().build();
+        }
 
-        Validate.notNull(expiration,
-                         "The SSO session's expiration time could not be determined. Please refresh your SSO session.");
+        Instant expiration = Instant.parse(expirationStr);
 
-        if (tokenIsInvalid(expiration)) {
-            throw ExpiredTokenException.builder().message("The SSO session associated with this profile has expired or is"
-                                                          + " otherwise invalid. To refresh this SSO session run aws sso"
-                                                          + " login with the corresponding profile.").build();
+        if (Instant.now().isAfter(expiration)) {
+            throw ExpiredTokenException.builder().build();
         }
 
         return SsoAccessToken.builder()
                              .accessToken(jsonNode.asObject().get("accessToken").text())
-                             .expiresAt(Instant.parse(expiration)).build();
-
+                             .expiresAt(expiration).build();
     }
 
-    private boolean tokenIsInvalid(String expirationTime) {
-        return Instant.now().isAfter(Instant.parse(expirationTime));
-    }
 
 }