Implement credentials refresh for LoginCredentialsProvider

Author: alextwoods

services/signin/src/main/java/software/amazon/awssdk/services/signin/auth/LoginCredentialsProvider.java

@@ -22,14 +22,24 @@
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.time.Duration;
+import java.time.Instant;
 import java.util.Optional;
 import software.amazon.awssdk.annotations.SdkPublicApi;
 import software.amazon.awssdk.annotations.ThreadSafe;
 import software.amazon.awssdk.auth.credentials.AwsCredentials;
 import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
+import software.amazon.awssdk.auth.credentials.AwsSessionCredentials;
+import software.amazon.awssdk.core.exception.SdkClientException;
 import software.amazon.awssdk.core.useragent.BusinessMetricFeatureId;
 import software.amazon.awssdk.services.signin.SigninClient;
+import software.amazon.awssdk.services.signin.internal.AccessTokenManager;
+import software.amazon.awssdk.services.signin.internal.LoginAccessToken;
+import software.amazon.awssdk.services.signin.internal.LoginCacheDirectorySystemSetting;
+import software.amazon.awssdk.services.signin.internal.OnDiskTokenManager;
 import software.amazon.awssdk.services.signin.model.CreateOAuth2TokenRequest;
+import software.amazon.awssdk.services.signin.model.CreateOAuth2TokenResponse;
+import software.amazon.awssdk.services.signin.model.SigninException;
+import software.amazon.awssdk.utils.Logger;
 import software.amazon.awssdk.utils.SdkAutoCloseable;
 import software.amazon.awssdk.utils.StringUtils;
 import software.amazon.awssdk.utils.builder.CopyableBuilder;
@@ -56,6 +66,8 @@
 public class LoginCredentialsProvider implements
                                       AwsCredentialsProvider, SdkAutoCloseable,
                                       ToCopyableBuilder<LoginCredentialsProvider.Builder, LoginCredentialsProvider> {
+    private static final Logger log = Logger.loggerFor(LoginCredentialsProvider.class);
+
     private static final String PROVIDER_NAME = BusinessMetricFeatureId.CREDENTIALS_LOGIN.value();
 
     private static final Duration DEFAULT_STALE_TIME = Duration.ofMinutes(1);
@@ -74,10 +86,15 @@ public class LoginCredentialsProvider implements
     private final Path tokenCacheLocation;
 
     private final CachedSupplier<AwsCredentials> credentialCache;
+    private final AccessTokenManager onDiskTokenManager;
 
     private final Boolean asyncCredentialUpdateEnabled;
 
-    public LoginCredentialsProvider(BuilderImpl builder) {
+    /**
+     *
+     * @see #builder()
+     */
+    private LoginCredentialsProvider(BuilderImpl builder) {
         this.signinClient = notNull(builder.signinClient, "SigninClient must not be null.");
         this.loginSession = paramNotBlank(builder.loginSession, "LoginSession");
 
@@ -94,6 +111,8 @@ public LoginCredentialsProvider(BuilderImpl builder) {
                                                         .map(p -> Paths.get(p))
                                                         .orElse(DEFAULT_TOKEN_LOCATION));
 
+        this.onDiskTokenManager = OnDiskTokenManager.create(this.tokenCacheLocation, this.loginSession);
+
         this.asyncCredentialUpdateEnabled = builder.asyncCredentialUpdateEnabled;
         CachedSupplier.Builder<AwsCredentials> cacheBuilder =
             CachedSupplier.builder(this::updateSigninCredentials)
@@ -110,8 +129,75 @@ public LoginCredentialsProvider(BuilderImpl builder) {
      * close to expiring.
      */
     private RefreshResult<AwsCredentials> updateSigninCredentials() {
-        // TODO: Future PRs will implement this
-        throw new UnsupportedOperationException();
+        // always re-load token from the disk in case it has been updated elsewhere
+        LoginAccessToken tokenFromDisc = onDiskTokenManager.loadToken().orElseThrow(
+            () -> SdkClientException.create("Token cache file for login_session `" + loginSession + "` not found. "
+                                            + "You must re-authenticate."));
+
+        Instant currentExpirationTime = tokenFromDisc.getAccessToken().expirationTime().orElseThrow(
+            () -> SdkClientException.create("Invalid token expiration time. You must re-authenticate.")
+        );
+
+        if (isWithinRefreshWindow(currentExpirationTime, staleTime)
+            && isWithinRefreshWindow(currentExpirationTime, prefetchTime)) {
+            log.debug(() -> "Using access token from disk, current expiration time is : " + currentExpirationTime);
+            AwsCredentials credentials = tokenFromDisc.getAccessToken()
+                .toBuilder()
+                .providerName(this.providerName)
+                .build();
+
+            return RefreshResult.builder(credentials)
+                                .staleTime(currentExpirationTime.minus(staleTime))
+                                .prefetchTime(currentExpirationTime.minus(prefetchTime))
+                                .build();
+        }
+
+        return refreshFromSigninService(tokenFromDisc);
+    }
+
+    private RefreshResult<AwsCredentials> refreshFromSigninService(LoginAccessToken tokenFromDisc) {
+        log.debug(() -> "Credentials are near expiration/expired, refreshing from Signin service.");
+
+        try {
+            CreateOAuth2TokenRequest refreshRequest =
+                CreateOAuth2TokenRequest
+                    .builder()
+                    .tokenInput(t -> t
+                        .clientId(tokenFromDisc.getClientId())
+                        .refreshToken(tokenFromDisc.getRefreshToken())
+                        .grantType("refresh_token"))
+                    .dpopProof("TODO") // TODO: This will be implemented in a separate PR
+                    .build();
+
+            CreateOAuth2TokenResponse createTokenResponse = signinClient.createOAuth2Token(refreshRequest);
+
+            Instant newExpiration = Instant.now().plusSeconds(createTokenResponse.tokenOutput().expiresIn());
+            AwsSessionCredentials updatedCredentials = AwsSessionCredentials
+                .builder()
+                .accessKeyId(createTokenResponse.tokenOutput().accessToken().accessKeyId())
+                .secretAccessKey(createTokenResponse.tokenOutput().accessToken().secretAccessKey())
+                .sessionToken(createTokenResponse.tokenOutput().accessToken().sessionToken())
+                .accountId(tokenFromDisc.getAccessToken().accountId().orElseThrow(
+                    () -> SdkClientException.create("Invalid access token, missing account ID. You must re-authenticate.")
+                ))
+                .expirationTime(newExpiration)
+                .providerName(this.providerName)
+                .build();
+
+            onDiskTokenManager.storeToken(tokenFromDisc.toBuilder()
+                                                       .accessToken(updatedCredentials)
+                                                       .refreshToken(createTokenResponse.tokenOutput().refreshToken())
+                                                       .build());
+
+            return RefreshResult.builder((AwsCredentials) updatedCredentials)
+                                .staleTime(newExpiration.minus(staleTime))
+                                .prefetchTime(newExpiration.minus(prefetchTime))
+                                .build();
+        } catch (SigninException serviceException) {
+            throw SdkClientException.create(
+                "Unable to refresh AWS Signin Access Token: You must re-authenticate.",
+                serviceException);
+        }
     }
 
     /**
@@ -152,6 +238,11 @@ public Builder toBuilder() {
         return new BuilderImpl(this);
     }
 
+    private static boolean isWithinRefreshWindow(Instant expiration, Duration staleTime) {
+        Instant now = Instant.now();
+        return expiration.isAfter(now.plus(staleTime));
+    }
+
     /**
      * A builder for creating a custom {@link LoginCredentialsProvider}.
      */

services/signin/src/main/java/software/amazon/awssdk/services/signin/internal/AccessTokenManager.java

@@ -0,0 +1,27 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.services.signin.internal;
+
+import java.util.Optional;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.utils.SdkAutoCloseable;
+
+@SdkInternalApi
+public interface AccessTokenManager extends SdkAutoCloseable {
+    Optional<LoginAccessToken> loadToken();
+
+    void storeToken(LoginAccessToken token);
+}