Adding ec2InstanceProfileName configuration to specify IMDS instance profile

Author: S-Saranya1

core/auth/src/main/java/software/amazon/awssdk/auth/credentials/InstanceProfileCredentialsProvider.java

@@ -18,6 +18,7 @@
 import static java.time.temporal.ChronoUnit.MINUTES;
 import static software.amazon.awssdk.utils.ComparableUtils.maximum;
 import static software.amazon.awssdk.utils.FunctionalUtils.invokeSafely;
+import static software.amazon.awssdk.utils.StringUtils.isBlank;
 import static software.amazon.awssdk.utils.cache.CachedSupplier.StaleValueBehavior.ALLOW;
 
 import java.net.URI;
@@ -103,6 +104,8 @@ private enum ApiVersion {
     private final Supplier<ProfileFile> profileFile;
 
     private final String profileName;
+    
+    private final String ec2InstanceProfileName;
 
     private final Duration staleTime;
 
@@ -118,6 +121,13 @@ private InstanceProfileCredentialsProvider(BuilderImpl builder) {
                                    .orElseGet(() -> ProfileFileSupplier.fixedProfileFile(ProfileFile.defaultProfileFile()));
         this.profileName = Optional.ofNullable(builder.profileName)
                                    .orElseGet(ProfileFileSystemSetting.AWS_PROFILE::getStringValueOrThrow);
+        this.ec2InstanceProfileName = builder.ec2InstanceProfileName;
+        
+        if (isBlank(ec2InstanceProfileName) && ec2InstanceProfileName != null) {
+            throw SdkClientException.builder()
+                                    .message("ec2InstanceProfileName cannot be blank")
+                                    .build();
+        }
 
         this.httpCredentialsLoader = HttpCredentialsLoader.create(PROVIDER_NAME);
         this.configProvider =
@@ -173,6 +183,11 @@ private RefreshResult<AwsCredentials> refreshCredentials() {
 
         try {
             LoadedCredentials credentials = httpCredentialsLoader.loadCredentials(createEndpointProvider());
+
+            if (apiVersion == ApiVersion.UNKNOWN) {
+                apiVersion = ApiVersion.EXTENDED;
+            }
+            
             Instant expiration = credentials.getExpiration().orElse(null);
             log.debug(() -> "Loaded credentials from IMDS with expiration time of " + expiration);
 
@@ -186,19 +201,25 @@ private RefreshResult<AwsCredentials> refreshCredentials() {
         } catch (Ec2MetadataClientException e) {
             if (e.statusCode() == 404) {
                 log.debug(() -> "Resolved profile is no longer available. Resetting it and trying again.");
-                resolvedProfile = null;
 
                 if (apiVersion == ApiVersion.UNKNOWN) {
                     apiVersion = ApiVersion.LEGACY;
                     return refreshCredentials();
+                } else if (ec2InstanceProfileName == null && configProvider.ec2InstanceProfileName() == null) {
+                    // Resolved profile name is invalid, reset it and try again
+                    resolvedProfile = null;
+                    
+                    profileRetryCount++;
+                    if (profileRetryCount <= MAX_PROFILE_RETRIES) {
+                        log.debug(() -> "Profile name not found, retrying fetching the profile name again.");
+                        return refreshCredentials();
+                    }
+                } else {
+                    throw SdkClientException.builder()
+                                           .message("Invalid profile name")
+                                           .cause(e)
+                                           .build();
                 }
-
-                profileRetryCount++;
-                if (profileRetryCount <= MAX_PROFILE_RETRIES) {
-                    log.debug(() -> "Profile name not found, retrying fetching the profile name again.");
-                    return refreshCredentials();
-                }
-                throw SdkClientException.create(FAILED_TO_LOAD_CREDENTIALS_ERROR, e);
             }
             throw SdkClientException.create(FAILED_TO_LOAD_CREDENTIALS_ERROR, e);
         } catch (RuntimeException e) {
@@ -328,6 +349,15 @@ private boolean isInsecureFallbackDisabled() {
     }
 
     private String[] getSecurityCredentials(String imdsHostname, String metadataToken) {
+        if (ec2InstanceProfileName != null) {
+            return new String[]{ec2InstanceProfileName};
+        }
+
+        String configuredProfileName = this.configProvider.ec2InstanceProfileName();
+        if (configuredProfileName != null) {
+            return new String[]{configuredProfileName};
+        }
+
         if (resolvedProfile != null) {
             return new String[]{resolvedProfile};
         }
@@ -383,6 +413,19 @@ public Builder toBuilder() {
      */
     public interface Builder extends HttpCredentialsProvider.Builder<InstanceProfileCredentialsProvider, Builder>,
                                      CopyableBuilder<Builder, InstanceProfileCredentialsProvider> {
+        /**
+         * Configure the EC2 instance profile name to use for retrieving credentials.
+         * 
+         * <p>When this is set, the provider will skip fetching the list of available instance profiles
+         * and use this name directly. This can improve performance by reducing the number of calls to IMDS.
+         * 
+         * <p>By default, this is not set and the provider will discover the instance profile name from IMDS.
+         * 
+         * @param ec2InstanceProfileName The EC2 instance profile name to use
+         * @return This builder for method chaining
+         */
+        Builder ec2InstanceProfileName(String ec2InstanceProfileName);
+        
         /**
          * Configure the profile file used for loading IMDS-related configuration, like the endpoint mode (IPv4 vs IPv6).
          *
@@ -434,6 +477,7 @@ static final class BuilderImpl implements Builder {
         private String asyncThreadName;
         private Supplier<ProfileFile> profileFile;
         private String profileName;
+        private String ec2InstanceProfileName;
         private Duration staleTime;
 
         private BuilderImpl() {
@@ -447,6 +491,7 @@ private BuilderImpl(InstanceProfileCredentialsProvider provider) {
             this.asyncThreadName = provider.asyncThreadName;
             this.profileFile = provider.profileFile;
             this.profileName = provider.profileName;
+            this.ec2InstanceProfileName = provider.ec2InstanceProfileName;
             this.staleTime = provider.staleTime;
         }
 
@@ -515,6 +560,17 @@ public Builder profileName(String profileName) {
         public void setProfileName(String profileName) {
             profileName(profileName);
         }
+        
+        @Override
+        public Builder ec2InstanceProfileName(String ec2InstanceProfileName) {
+            this.ec2InstanceProfileName = ec2InstanceProfileName;
+            return this;
+        }
+        
+        public void setEc2InstanceProfileName(String ec2InstanceProfileName) {
+            ec2InstanceProfileName(ec2InstanceProfileName);
+        }
+
 
         @Override
         public Builder staleTime(Duration duration) {

core/auth/src/main/java/software/amazon/awssdk/auth/credentials/internal/Ec2MetadataConfigProvider.java

@@ -46,12 +46,14 @@ public final class Ec2MetadataConfigProvider {
 
     private final Lazy<Boolean> metadataV1Disabled;
     private final Lazy<Long> serviceTimeout;
+    private final Lazy<String> ec2InstanceProfileName;
 
     private Ec2MetadataConfigProvider(Builder builder) {
         this.profileFile = builder.profileFile;
         this.profileName = builder.profileName;
         this.metadataV1Disabled = new Lazy<>(this::resolveMetadataV1Disabled);
         this.serviceTimeout = new Lazy<>(this::resolveServiceTimeout);
+        this.ec2InstanceProfileName = new Lazy<>(this::resolveEc2InstanceProfileName);
     }
 
     public enum EndpointMode {
@@ -127,6 +129,14 @@ public boolean isMetadataV1Disabled() {
     public long serviceTimeout() {
         return serviceTimeout.getValue();
     }
+    
+    /**
+     * Resolves the EC2 Instance Profile Name to use.
+     * @return the EC2 Instance Profile Name or null if not specified.
+     */
+    public String ec2InstanceProfileName() {
+        return ec2InstanceProfileName.getValue();
+    }
 
     // Internal resolution logic for Metadata V1 disabled
     private boolean resolveMetadataV1Disabled() {
@@ -146,6 +156,14 @@ private long resolveServiceTimeout() {
                             .orElseGet(() -> parseTimeoutValue(SdkSystemSetting.AWS_METADATA_SERVICE_TIMEOUT.defaultValue()));
     }
 
+    private String resolveEc2InstanceProfileName() {
+        return OptionalUtils.firstPresent(
+                                fromSystemSettingsEc2InstanceProfileName(),
+                                () -> fromProfileFileEc2InstanceProfileName(profileFile, profileName)
+                            )
+                            .orElse(null);
+    }
+
     // System settings resolution for Metadata V1 disabled
     private static Optional<Boolean> fromSystemSettingsMetadataV1Disabled() {
         return SdkSystemSetting.AWS_EC2_METADATA_V1_DISABLED.getBooleanValue();
@@ -171,6 +189,22 @@ private static Optional<Long> fromProfileFileServiceTimeout(Supplier<ProfileFile
                           .flatMap(p -> p.property(ProfileProperty.METADATA_SERVICE_TIMEOUT))
                           .map(Ec2MetadataConfigProvider::parseTimeoutValue);
     }
+    
+    // System settings resolution for EC2 Instance Profile Name
+    private static Optional<String> fromSystemSettingsEc2InstanceProfileName() {
+        return SdkSystemSetting.AWS_EC2_INSTANCE_PROFILE_NAME.getNonDefaultStringValue();
+    }
+    
+    // Profile file resolution for EC2 Instance Profile Name
+    private static Optional<String> fromProfileFileEc2InstanceProfileName(Supplier<ProfileFile> profileFile, String profileName) {
+        try {
+            return profileFile.get()
+                              .profile(profileName)
+                              .flatMap(p -> p.property(ProfileProperty.EC2_INSTANCE_PROFILE_NAME));
+        } catch (Exception e) {
+            return Optional.empty();
+        }
+    }
 
     // Parses a timeout value from a string to milliseconds
     private static long parseTimeoutValue(String timeoutValue) {