ProfileTokenProviderLoader can use Supplier<Profilefile> internally

dave-fn · dave-fn · commit c594e620aa1c · 2022-12-18T15:55:35.000-07:00
diff --git a/core/auth/src/main/java/software/amazon/awssdk/auth/token/credentials/ProfileTokenProvider.java b/core/auth/src/main/java/software/amazon/awssdk/auth/token/credentials/ProfileTokenProvider.java
@@ -15,7 +15,6 @@
 
 package software.amazon.awssdk.auth.token.credentials;
 
-import java.util.Objects;
 import java.util.Optional;
 import java.util.function.Supplier;
 import software.amazon.awssdk.annotations.SdkPublicApi;
@@ -38,35 +37,45 @@
  */
 @SdkPublicApi
 public final class ProfileTokenProvider implements SdkTokenProvider, SdkAutoCloseable {
-    private SdkTokenProvider tokenProvider;
+    private final SdkTokenProvider tokenProvider;
     private final RuntimeException loadException;
-    private final Supplier<ProfileFile> profileFile;
-    private volatile ProfileFile currentProfileFile;
+
     private final String profileName;
 
     /**
      * @see #builder()
      */
     private ProfileTokenProvider(BuilderImpl builder) {
+        SdkTokenProvider sdkTokenProvider = null;
         RuntimeException thrownException = null;
-        Supplier<ProfileFile> selectedProfileFileSupplier = null;
+        Supplier<ProfileFile> selectedProfileFile = null;
         String selectedProfileName = null;
 
         try {
             selectedProfileName = Optional.ofNullable(builder.profileName)
                                           .orElseGet(ProfileFileSystemSetting.AWS_PROFILE::getStringValueOrThrow);
 
-            selectedProfileFileSupplier = Optional.ofNullable(builder.profileFile)
-                                                  .orElse(builder.defaultProfileFileLoader);
+            // Load the profiles file
+            selectedProfileFile = Optional.ofNullable(builder.profileFile)
+                                          .orElse(builder.defaultProfileFileLoader);
+
+            // Load the profile and token provider
+            sdkTokenProvider = createTokenProvider(selectedProfileFile, selectedProfileName);
 
         } catch (RuntimeException e) {
             // If we couldn't load the provider for some reason, save an exception describing why.
             thrownException = e;
         }
 
-        this.loadException = thrownException;
-        this.profileFile = selectedProfileFileSupplier;
-        this.profileName = selectedProfileName;
+        if (thrownException != null) {
+            this.loadException = thrownException;
+            this.tokenProvider = null;
+            this.profileName = null;
+        } else {
+            this.loadException = null;
+            this.tokenProvider = sdkTokenProvider;
+            this.profileName = selectedProfileName;
+        }
     }
 
     /**
@@ -99,21 +108,13 @@ public SdkToken resolveToken() {
         if (loadException != null) {
             throw loadException;
         }
-
-        ProfileFile cachedOrRefreshedProfileFile = refreshProfileFile();
-        if (isNewProfileFile(cachedOrRefreshedProfileFile)) {
-            currentProfileFile = cachedOrRefreshedProfileFile;
-            handleProfileFileReload(cachedOrRefreshedProfileFile);
-        }
-
         return tokenProvider.resolveToken();
     }
 
     @Override
     public String toString() {
         return ToString.builder("ProfileTokenProvider")
                        .add("profileName", profileName)
-                       .add("profileFile", currentProfileFile)
                        .build();
     }
 
@@ -123,27 +124,14 @@ public void close() {
         IoUtils.closeIfCloseable(tokenProvider, null);
     }
 
-    private SdkTokenProvider createTokenProvider(ProfileFile profileFile, String profileName) {
-        // Load the profile and token provider
-        return profileFile.profile(profileName)
-                          .flatMap(p -> new ProfileTokenProviderLoader(profileFile, p).tokenProvider())
-                          .orElseThrow(() -> {
-                              String errorMessage = String.format("Profile file contained no information for " +
-                                                                  "profile '%s': %s", profileName, profileFile);
-                              return SdkClientException.builder().message(errorMessage).build();
-                          });
-    }
-
-    private void handleProfileFileReload(ProfileFile profileFile) {
-        tokenProvider = createTokenProvider(profileFile, profileName);
-    }
-
-    private ProfileFile refreshProfileFile() {
-        return profileFile.get();
-    }
-
-    private boolean isNewProfileFile(ProfileFile profileFile) {
-        return !Objects.equals(currentProfileFile, profileFile);
+    private SdkTokenProvider createTokenProvider(Supplier<ProfileFile> profileFile, String profileName) {
+        return new ProfileTokenProviderLoader(profileFile, profileName)
+            .tokenProvider()
+            .orElseThrow(() -> {
+                String errorMessage = String.format("Profile file contained no information for " +
+                                                    "profile '%s'", profileName);
+                return SdkClientException.builder().message(errorMessage).build();
+            });
     }
 
     /**
diff --git a/core/auth/src/main/java/software/amazon/awssdk/auth/token/internal/ProfileTokenProviderLoader.java b/core/auth/src/main/java/software/amazon/awssdk/auth/token/internal/ProfileTokenProviderLoader.java
@@ -17,15 +17,19 @@
 
 import java.lang.reflect.InvocationTargetException;
 import java.util.Arrays;
+import java.util.Objects;
 import java.util.Optional;
+import java.util.function.Supplier;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.auth.token.credentials.ChildProfileTokenProviderFactory;
 import software.amazon.awssdk.auth.token.credentials.SdkTokenProvider;
+import software.amazon.awssdk.core.exception.SdkClientException;
 import software.amazon.awssdk.core.internal.util.ClassLoaderHelper;
 import software.amazon.awssdk.profiles.Profile;
 import software.amazon.awssdk.profiles.ProfileFile;
 import software.amazon.awssdk.profiles.ProfileProperty;
 import software.amazon.awssdk.profiles.internal.ProfileSection;
+import software.amazon.awssdk.utils.Lazy;
 import software.amazon.awssdk.utils.Validate;
 
 /**
@@ -36,12 +40,33 @@ public final class ProfileTokenProviderLoader {
     private static final String SSO_OIDC_TOKEN_PROVIDER_FACTORY =
         "software.amazon.awssdk.services.ssooidc.SsoOidcProfileTokenProviderFactory";
 
+    private final boolean createProviderFromSupplier;
+
     private final Profile profile;
     private final ProfileFile profileFile;
+    private final Supplier<ProfileFile> profileFileSupplier;
+    private final String profileName;
+    private volatile ProfileFile currentProfileFile;
+    private volatile SdkTokenProvider currentTokenProvider;
+
+    private final Lazy<ChildProfileTokenProviderFactory> factory;
 
-    public ProfileTokenProviderLoader(ProfileFile profileFile, Profile profile) {
-        this.profile = Validate.paramNotNull(profile, "profile");
+    public ProfileTokenProviderLoader(ProfileFile profileFile, String profileName) {
+        this.createProviderFromSupplier = false;
         this.profileFile = Validate.paramNotNull(profileFile, "profileFile");
+        this.profileFileSupplier = null;
+        this.profileName = Validate.paramNotNull(profileName, "profileName");
+        this.profile = resolveProfile(profileFile, profileName);
+        this.factory = new Lazy<>(this::ssoTokenProviderFactory);
+    }
+
+    public ProfileTokenProviderLoader(Supplier<ProfileFile> profileFile, String profileName) {
+        createProviderFromSupplier = true;
+        this.profileFile = null;
+        this.profileFileSupplier = Validate.paramNotNull(profileFile, "profileFile");
+        this.profileName = Validate.paramNotNull(profileName, "profileName");
+        this.profile = null;
+        this.factory = new Lazy<>(this::ssoTokenProviderFactory);
     }
 
     /**
@@ -55,19 +80,53 @@ public Optional<SdkTokenProvider> tokenProvider() {
      * Create the SSO credentials provider based on the related profile properties.
      */
     private SdkTokenProvider ssoProfileCredentialsProvider() {
+        if (createProviderFromSupplier) {
+            return () -> ssoProfileCredentialsProvider(profileFileSupplier, profileName).resolveToken();
+        }
 
-        String profileSsoSectionName = profile.property(ProfileSection.SSO_SESSION.getPropertyKeyName())
-                                              .orElseThrow(() -> new IllegalArgumentException(
-                                                  "Profile " + profile.name() + " does not have sso_session property"));
+        return ssoProfileCredentialsProvider(profileFile, profile);
+    }
+
+    private SdkTokenProvider ssoProfileCredentialsProvider(ProfileFile profileFile, Profile profile) {
+        String profileSsoSectionName = profileSsoSectionName(profile);
+        Profile ssoProfile = ssoProfile(profileFile, profileSsoSectionName);
 
-        Profile ssoProfile = profileFile.getSection(ProfileSection.SSO_SESSION.getSectionTitle(), profileSsoSectionName)
-                                        .orElseThrow(() -> new IllegalArgumentException(
-                                            "Sso-session section not found with sso-session title " + profileSsoSectionName));
+        validateRequiredProperties(ssoProfile, ProfileProperty.SSO_REGION, ProfileProperty.SSO_START_URL);
+
+        return factory.getValue().create(profileFile, profile);
+    }
+
+    private synchronized SdkTokenProvider ssoProfileCredentialsProvider(Supplier<ProfileFile> profileFile, String profileName) {
+        ProfileFile profileFileInstance = profileFile.get();
+        Profile profileInstance = resolveProfile(profileFileInstance, profileName);
+        if (!Objects.equals(profileFileInstance, currentProfileFile)) {
+            currentProfileFile = profileFileInstance;
+            currentTokenProvider = ssoProfileCredentialsProvider(profileFileInstance, profileInstance);
+        }
+
+        return currentTokenProvider;
+    }
+
+    private Profile resolveProfile(ProfileFile profileFile, String profileName) {
+        return profileFile.profile(profileName)
+                          .orElseThrow(() -> {
+                              String errorMessage = String.format("Profile file contained no information for profile '%s': %s",
+                                                                  profileName, profileFile);
+                              return SdkClientException.builder().message(errorMessage).build();
+                          });
+    }
+
+    private String profileSsoSectionName(Profile profile) {
+        return Optional.ofNullable(profile)
+                       .flatMap(p -> p.property(ProfileSection.SSO_SESSION.getPropertyKeyName()))
+                       .orElseThrow(() -> new IllegalArgumentException(
+                           "Profile " + profileName + " does not have sso_session property"));
+    }
 
-        validateRequiredProperties(ssoProfile,
-                                   ProfileProperty.SSO_REGION,
-                                   ProfileProperty.SSO_START_URL);
-        return ssoTokenProviderFactory().create(profileFile, profile);
+    private Profile ssoProfile(ProfileFile profileFile, String profileSsoSectionName) {
+        return profileFile.getSection(ProfileSection.SSO_SESSION.getSectionTitle(), profileSsoSectionName)
+                          .orElseThrow(() -> new IllegalArgumentException(
+                              "Sso-session section not found with sso-session title " + profileSsoSectionName));
     }
 
     /**
@@ -76,7 +135,8 @@ private SdkTokenProvider ssoProfileCredentialsProvider() {
     private void validateRequiredProperties(Profile ssoProfile, String... requiredProperties) {
         Arrays.stream(requiredProperties)
               .forEach(p -> Validate.isTrue(ssoProfile.properties().containsKey(p),
-                                            "Property '%s' was not configured for profile '%s'.", p, this.profile.name()));
+                                            "Property '%s' was not configured for profile '%s'.",
+                                            p, profileName));
     }
 
     /**
@@ -88,10 +148,10 @@ private ChildProfileTokenProviderFactory ssoTokenProviderFactory() {
                                                                                getClass());
             return (ChildProfileTokenProviderFactory) ssoOidcTokenProviderFactory.getConstructor().newInstance();
         } catch (ClassNotFoundException e) {
-            throw new IllegalStateException("To use SSO OIDC related properties in the '" + profile.name() + "' profile, "
+            throw new IllegalStateException("To use SSO OIDC related properties in the '" + profileName + "' profile, "
                                             + "the 'ssooidc' service module must be on the class path.", e);
         } catch (NoSuchMethodException | InvocationTargetException | InstantiationException | IllegalAccessException e) {
-            throw new IllegalStateException("Failed to create the '" + profile.name() + "' token provider factory.", e);
+            throw new IllegalStateException("Failed to create the '%s" + profileName + "' token provider factory.", e);
         }
     }
 }
diff --git a/services/ssooidc/src/test/java/software/amazon/awssdk/services/ssooidc/internal/common/SsoOidcTokenRefreshTestBase.java b/services/ssooidc/src/test/java/software/amazon/awssdk/services/ssooidc/internal/common/SsoOidcTokenRefreshTestBase.java
@@ -102,8 +102,7 @@ protected void initializeProfileProperties() {
                                           .content(new StringInputStream(profileContent))
                                           .type(ProfileFile.Type.CONFIGURATION)
                                           .build();
-        Optional<Profile> profile = profiles.profile("sso-refresh");
-        profileTokenProviderLoader = new ProfileTokenProviderLoader(profiles, profile.get());
+        profileTokenProviderLoader = new ProfileTokenProviderLoader(profiles, "sso-refresh");
     }
 
     @Test
diff --git a/test/auth-tests/src/it/java/software/amazon/awssdk/auth/ssooidc/ProfileTokenProviderLoaderTest.java b/test/auth-tests/src/it/java/software/amazon/awssdk/auth/ssooidc/ProfileTokenProviderLoaderTest.java
@@ -17,10 +17,11 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
+import static org.assertj.core.api.Assertions.assertThatNoException;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
-import java.util.HashMap;
 import java.util.Optional;
+import java.util.function.Supplier;
 import java.util.stream.Stream;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
@@ -29,7 +30,6 @@
 import software.amazon.awssdk.auth.token.credentials.SdkTokenProvider;
 import software.amazon.awssdk.auth.token.internal.ProfileTokenProviderLoader;
 import software.amazon.awssdk.core.exception.SdkClientException;
-import software.amazon.awssdk.profiles.Profile;
 import software.amazon.awssdk.profiles.ProfileFile;
 import software.amazon.awssdk.utils.StringInputStream;
 
@@ -43,18 +43,47 @@ public void noProfile_throwsException() {
 
     @Test
     public void noProfileFile_throwsException() {
-        assertThatThrownBy(() -> new ProfileTokenProviderLoader(null, Profile.builder().name("sso").properties(new HashMap<>()).build()))
+        assertThatThrownBy(() -> new ProfileTokenProviderLoader((ProfileFile) null, "sso"))
             .hasMessageContaining("profileFile must not be null");
     }
 
+    @Test
+    public void profileTokenProviderLoader_noProfileFileSupplier_throwsException() {
+        assertThatThrownBy(() -> new ProfileTokenProviderLoader((ProfileFile) null, "sso'"))
+            .hasMessageContaining("profileFile must not be null");
+    }
+
+    @Test
+    public void profileTokenProviderLoader_noProfileName_throwsException() {
+        assertThatThrownBy(() -> new ProfileTokenProviderLoader(ProfileFile::defaultProfileFile, null))
+            .hasMessageContaining("profileName must not be null");
+    }
+
     @ParameterizedTest
     @MethodSource("ssoErrorValues")
     public void incorrectSsoProperties_throwsException(String profileContent, String msg) {
         ProfileFile profileFile = configFile(profileContent);
 
-        assertThat(profileFile.profile("sso")).hasValueSatisfying(profile -> {
-            ProfileTokenProviderLoader providerLoader = new ProfileTokenProviderLoader(profileFile, profile);
-            assertThatExceptionOfType(IllegalArgumentException.class).isThrownBy(providerLoader::tokenProvider).withMessageContaining(msg);
+        ProfileTokenProviderLoader providerLoader = new ProfileTokenProviderLoader(profileFile, "sso");
+        assertThatExceptionOfType(IllegalArgumentException.class).isThrownBy(providerLoader::tokenProvider)
+                                                                 .withMessageContaining(msg);
+    }
+
+    @ParameterizedTest
+    @MethodSource("ssoErrorValues")
+    public void incorrectSsoProperties_supplier_delaysThrowingExceptionUntilResolvingToken(String profileContent, String msg) {
+        ProfileFile profileFile = configFile(profileContent);
+        Supplier<ProfileFile> supplier = () -> profileFile;
+
+        ProfileTokenProviderLoader providerLoader = new ProfileTokenProviderLoader(supplier, "sso");
+
+        assertThatNoException().isThrownBy(providerLoader::tokenProvider);
+
+        assertThat(providerLoader.tokenProvider()).satisfies(tokenProviderOptional -> {
+           assertThat(tokenProviderOptional).isPresent().get().satisfies(tokenProvider-> {
+               assertThatExceptionOfType(IllegalArgumentException.class).isThrownBy(tokenProvider::resolveToken)
+                                                                        .withMessageContaining(msg);
+           });
         });
     }
 
@@ -66,8 +95,7 @@ public void correctSsoProperties_createsTokenProvider() {
                                 "sso_region=us-east-1\n" +
                                 "sso_start_url=https://d-abc123.awsapps.com/start\n";
 
-        Optional<Profile> ssoProfile = configFile(profileContent).profile("sso");
-        ProfileTokenProviderLoader providerLoader = new ProfileTokenProviderLoader(configFile(profileContent), ssoProfile.get());
+        ProfileTokenProviderLoader providerLoader = new ProfileTokenProviderLoader(configFile(profileContent), "sso");
         Optional<SdkTokenProvider> tokenProvider = providerLoader.tokenProvider();
         assertThat(tokenProvider).isPresent();
         assertThatThrownBy(() -> tokenProvider.get().resolveToken())

Original file line number	Diff line number	Diff line change
`@@ -102,8 +102,7 @@ protected void initializeProfileProperties() {`
`102`	`102`	`.content(new StringInputStream(profileContent))`
`103`	`103`	`.type(ProfileFile.Type.CONFIGURATION)`
`104`	`104`	`.build();`
`105`		`- Optional<Profile> profile = profiles.profile("sso-refresh");`
`106`		`- profileTokenProviderLoader = new ProfileTokenProviderLoader(profiles, profile.get());`
	`105`	`+ profileTokenProviderLoader = new ProfileTokenProviderLoader(profiles, "sso-refresh");`
`107`	`106`	`}`
`108`	`107`
`109`	`108`	`@Test`