Revert "feat: add support for account ID in IMDS credentials (#1573)"

This reverts commit a707e1b.

Author: lauzadis

aws-runtime/aws-config/build.gradle.kts

@@ -8,7 +8,6 @@ import aws.sdk.kotlin.gradle.codegen.smithyKotlinProjectionSrcDir
 
 plugins {
     alias(libs.plugins.aws.kotlin.repo.tools.smithybuild)
-    alias(libs.plugins.kotlinx.serialization)
 }
 
 description = "Support for AWS configuration"
@@ -53,7 +52,6 @@ kotlin {
                 implementation(libs.kotlinx.coroutines.test)
                 implementation(libs.smithy.kotlin.http.test)
                 implementation(libs.kotlinx.serialization.json)
-                implementation(libs.kotest.framework.datatest)
             }
         }
         jvmTest {

aws-runtime/aws-config/common/src/aws/sdk/kotlin/runtime/auth/credentials/DefaultChainCredentialsProvider.kt

@@ -41,7 +41,7 @@ import aws.smithy.kotlin.runtime.util.PlatformProvider
  * @param region the region to make credentials requests to.
  * @return the newly-constructed credentials provider
  */
-public class DefaultChainCredentialsProvider(
+public class DefaultChainCredentialsProvider constructor(
     public val profileName: String? = null,
     public val platformProvider: PlatformProvider = PlatformProvider.System,
     httpClient: HttpClientEngine? = null,
@@ -51,11 +51,6 @@ public class DefaultChainCredentialsProvider(
     private val manageEngine = httpClient == null
     private val engine = httpClient ?: DefaultHttpEngine()
 
-    private val imdsClient = ImdsClient {
-        platformProvider = this@DefaultChainCredentialsProvider.platformProvider
-        engine = this@DefaultChainCredentialsProvider.engine
-    }
-
     private val chain = CredentialsProviderChain(
         SystemPropertyCredentialsProvider(platformProvider::getProperty),
         EnvironmentCredentialsProvider(platformProvider::getenv),
@@ -64,7 +59,12 @@ public class DefaultChainCredentialsProvider(
         ProfileCredentialsProvider(profileName = profileName, platformProvider = platformProvider, httpClient = engine, region = region),
         EcsCredentialsProvider(platformProvider, engine),
         ImdsCredentialsProvider(
-            client = imdsClient,
+            client = lazy {
+                ImdsClient {
+                    platformProvider = this@DefaultChainCredentialsProvider.platformProvider
+                    engine = this@DefaultChainCredentialsProvider.engine
+                }
+            },
             platformProvider = platformProvider,
         ),
     )
@@ -75,7 +75,6 @@ public class DefaultChainCredentialsProvider(
 
     override fun close() {
         provider.close()
-        imdsClient.close()
         if (manageEngine) {
             engine.closeIfCloseable()
         }

aws-runtime/aws-config/common/src/aws/sdk/kotlin/runtime/auth/credentials/ImdsCredentialsProvider.kt

@@ -5,232 +5,150 @@
 
 package aws.sdk.kotlin.runtime.auth.credentials
 
-import aws.sdk.kotlin.runtime.auth.credentials.internal.credentials
-import aws.sdk.kotlin.runtime.config.imds.*
+import aws.sdk.kotlin.runtime.config.AwsSdkSetting
+import aws.sdk.kotlin.runtime.config.imds.EC2MetadataError
+import aws.sdk.kotlin.runtime.config.imds.ImdsClient
+import aws.sdk.kotlin.runtime.config.imds.InstanceMetadataProvider
 import aws.sdk.kotlin.runtime.http.interceptors.businessmetrics.AwsBusinessMetric
 import aws.sdk.kotlin.runtime.http.interceptors.businessmetrics.withBusinessMetric
 import aws.smithy.kotlin.runtime.auth.awscredentials.*
 import aws.smithy.kotlin.runtime.collections.Attributes
+import aws.smithy.kotlin.runtime.config.resolve
 import aws.smithy.kotlin.runtime.http.HttpStatusCode
 import aws.smithy.kotlin.runtime.io.IOException
 import aws.smithy.kotlin.runtime.serde.json.JsonDeserializer
 import aws.smithy.kotlin.runtime.telemetry.logging.info
+import aws.smithy.kotlin.runtime.telemetry.logging.warn
 import aws.smithy.kotlin.runtime.time.Clock
+import aws.smithy.kotlin.runtime.time.Instant
 import aws.smithy.kotlin.runtime.util.PlatformEnvironProvider
 import aws.smithy.kotlin.runtime.util.PlatformProvider
-import aws.smithy.kotlin.runtime.util.SingleFlightGroup
-import aws.smithy.kotlin.runtime.util.asyncLazy
+import kotlinx.coroutines.sync.Mutex
+import kotlinx.coroutines.sync.withLock
 import kotlin.coroutines.coroutineContext
+import kotlin.time.Duration.Companion.seconds
 
+private const val CREDENTIALS_BASE_PATH: String = "/latest/meta-data/iam/security-credentials/"
 private const val CODE_ASSUME_ROLE_UNAUTHORIZED_ACCESS: String = "AssumeRoleUnauthorizedAccess"
 private const val PROVIDER_NAME = "IMDSv2"
 
 /**
  * [CredentialsProvider] that uses EC2 instance metadata service (IMDS) to provide credentials information.
- * This provider requires that the EC2 instance has an
- * [instance profile](https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/iam-roles-for-amazon-ec2.html#ec2-instance-profile)
+ * This provider requires that the EC2 instance has an [instance profile](https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/iam-roles-for-amazon-ec2.html#ec2-instance-profile)
  * configured.
- * @param instanceProfileName overrides the instance profile name. When set, this provider skips querying IMDS for the
- * name of the active profile.
- * @param client a preconfigured IMDS client with which to retrieve instance metadata. If an instance is passed, the
- * caller is responsible for closing it. If no instance is passed, a default instance is created and will be closed when
- * this credentials provider is closed.
- * @param platformProvider a platform provider used for env vars and system properties
+ *
+ * See [EC2 IAM Roles](https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/iam-roles-for-amazon-ec2.html) for more
+ * information.
+ *
+ * @param profileOverride override the instance profile name. When retrieving credentials, a call must first be made to
+ * `<IMDS_BASE_URL>/latest/meta-data/iam/security-credentials/`. This returns the instance profile used. If
+ * [profileOverride] is set, the initial call to retrieve the profile is skipped and the provided value is used instead.
+ * @param client the IMDS client to use to resolve credentials information with. This provider takes ownership over
+ * the lifetime of the given [ImdsClient] and will close it when the provider is closed.
+ * @param platformProvider the [PlatformEnvironProvider] instance
  */
 public class ImdsCredentialsProvider(
-    instanceProfileName: String? = null,
-    client: InstanceMetadataProvider? = null,
-    platformProvider: PlatformProvider = PlatformProvider.System,
+    public val profileOverride: String? = null,
+    public val client: Lazy<InstanceMetadataProvider> = lazy { ImdsClient() },
+    public val platformProvider: PlatformEnvironProvider = PlatformProvider.System,
+    private val clock: Clock = Clock.System,
 ) : CloseableCredentialsProvider {
-
-    @Deprecated("This constructor supports parameters which are no longer used in the implementation. It will be removed in version 1.5.")
-    public constructor(
-        profileOverride: String? = null,
-        client: Lazy<InstanceMetadataProvider> = lazy { ImdsClient() },
-        platformProvider: PlatformEnvironProvider = PlatformProvider.System,
-        @Suppress("UNUSED_PARAMETER") clock: Clock = Clock.System,
-    ) : this(
-        profileOverride,
-        client.value,
-        platformProvider = platformProvider as? PlatformProvider ?: PlatformProvider.System,
-    )
-
-    private val actualPlatformProvider = platformProvider
-
-    @Deprecated("This property is retained for backwards compatibility but no longer needs to be public and will be removed in version 1.5.")
-    public val platformProvider: PlatformEnvironProvider = actualPlatformProvider
-
-    private val manageClient: Boolean = client == null
-
-    private val actualClient = client ?: ImdsClient {
-        this.platformProvider = actualPlatformProvider
-    }
-
-    @Deprecated("This property is retained for backwards compatibility but no longer needs to be public and will be removed in version 1.5.")
-    public val client: Lazy<InstanceMetadataProvider>
-        get() = lazyOf(actualClient)
-
-    private val instanceProfileName = asyncLazy {
-        instanceProfileName ?: resolveEc2InstanceProfileName(platformProvider)
-    }
-
-    @Deprecated("This property is retained for backwards compatibility but no longer needs to be public and will be removed in version 1.5.")
-    public val profileOverride: String? = instanceProfileName
-
-    private val providerDisabled = asyncLazy {
-        resolveDisableEc2Metadata(platformProvider) ?: false
-    }
-
-    /**
-     * Tracks the known-good version of IMDS APIs available in the local environment. This starts as `null` and will be
-     * updated after the first successful API call.
-     */
-    private var apiVersion: ApiVersion? = null
-
-    private val urlBase: String
-        get() = (apiVersion ?: ApiVersion.EXTENDED).urlBase
-
     private var previousCredentials: Credentials? = null
 
-    /**
-     * Tracks the instance profile name resolved from IMDS. This starts as `null` and will be updated after a
-     * successful API call. Note that if [instanceProfileName] is set, profile name resolution  will be skipped.
-     */
-    private var resolvedProfileName: String? = null
-
-    /**
-     * A deduplicator for resolving credentials and tracking mutable state about IMDS
-     */
-    private val sfg = SingleFlightGroup<Credentials>()
+    // the time to refresh the Credentials. If set, it will take precedence over the Credentials' expiration time
+    private var nextRefresh: Instant? = null
 
-    override suspend fun resolve(attributes: Attributes): Credentials = sfg.singleFlight(::resolveSingleFlight)
+    // protects previousCredentials and nextRefresh
+    private val mu = Mutex()
 
-    private suspend fun resolveSingleFlight(): Credentials {
-        if (providerDisabled.get()) {
+    override suspend fun resolve(attributes: Attributes): Credentials {
+        if (AwsSdkSetting.AwsEc2MetadataDisabled.resolve(platformProvider) == true) {
             throw CredentialsNotLoadedException("AWS EC2 metadata is explicitly disabled; credentials not loaded")
         }
 
-        val profileName = instanceProfileName.get() ?: resolvedProfileName ?: try {
-            actualClient.get(urlBase).also {
-                if (apiVersion == null) {
-                    // Tried EXTENDED and it worked; remember that for the future
-                    apiVersion = ApiVersion.EXTENDED
+        // if we have previously served IMDS credentials and it's not time for a refresh, just return the previous credentials
+        mu.withLock {
+            previousCredentials?.run {
+                nextRefresh?.takeIf { clock.now() < it }?.run {
+                    return previousCredentials!!
                 }
             }
-        } catch (ex: EC2MetadataError) {
-            when {
-                apiVersion == null && ex.status == HttpStatusCode.NotFound -> {
-                    // Tried EXTENDED and that didn't work; fallback to LEGACY
-                    apiVersion = ApiVersion.LEGACY
-                    return resolveSingleFlight()
-                }
-
-                ex.status == HttpStatusCode.NotFound -> {
-                    coroutineContext.info<ImdsCredentialsProvider> {
-                        "Received 404 when loading profile name. This instance may not have an associated profile."
-                    }
-                    throw ex
-                }
-
-                else -> return usePreviousCredentials()
-                    ?: throw ImdsProfileException(ex).wrapAsCredentialsProviderException()
-            }
-        } catch (ex: IOException) {
-            return usePreviousCredentials() ?: throw ImdsProfileException(ex).wrapAsCredentialsProviderException()
-        } catch (ex: Exception) {
-            throw ImdsProfileException(ex).wrapAsCredentialsProviderException()
         }
 
-        val credsPayload = try {
-            actualClient.get("$urlBase$profileName")
-        } catch (ex: EC2MetadataError) {
-            when {
-                apiVersion == null && ex.status == HttpStatusCode.NotFound -> {
-                    // Tried EXTENDED and that didn't work; fallback to LEGACY
-                    apiVersion = ApiVersion.LEGACY
-                    return resolveSingleFlight()
-                }
-
-                instanceProfileName.get() == null && ex.status == HttpStatusCode.NotFound -> {
-                    // A previously-resolved profile is now invalid; forget the resolved name and re-resolve
-                    resolvedProfileName = null
-                    return resolveSingleFlight()
-                }
-
-                else -> return usePreviousCredentials()
-                    ?: throw ImdsCredentialsException(profileName, ex).wrapAsCredentialsProviderException()
-            }
-        } catch (ex: IOException) {
-            return usePreviousCredentials()
-                ?: throw ImdsCredentialsException(profileName, ex).wrapAsCredentialsProviderException()
+        val profileName = try {
+            profileOverride ?: loadProfile()
         } catch (ex: Exception) {
-            throw ImdsCredentialsException(profileName, ex).wrapAsCredentialsProviderException()
+            return useCachedCredentials(ex) ?: throw CredentialsProviderException("failed to load instance profile", ex)
         }
 
-        if (instanceProfileName.get() == null) {
-            // No profile name was provided at construction time; cache the resolved name
-            resolvedProfileName = profileName
+        val payload = try {
+            client.value.get("$CREDENTIALS_BASE_PATH$profileName")
+        } catch (ex: Exception) {
+            return useCachedCredentials(ex) ?: throw CredentialsProviderException("failed to load credentials", ex)
         }
 
-        val deserializer = JsonDeserializer(credsPayload.encodeToByteArray())
+        val deserializer = JsonDeserializer(payload.encodeToByteArray())
 
         return when (val resp = deserializeJsonCredentials(deserializer)) {
             is JsonCredentialsResponse.SessionCredentials -> {
-                val creds = credentials(
+                nextRefresh = if (resp.expiration != null && resp.expiration < clock.now()) {
+                    coroutineContext.warn<ImdsCredentialsProvider> {
+                        "Attempting credential expiration extension due to a credential service availability issue. " +
+                            "A refresh of these credentials will be attempted again in " +
+                            "${ DEFAULT_CREDENTIALS_REFRESH_SECONDS / 60 } minutes."
+                    }
+                    clock.now() + DEFAULT_CREDENTIALS_REFRESH_SECONDS.seconds
+                } else {
+                    null
+                }
+
+                val creds = Credentials(
                     resp.accessKeyId,
                     resp.secretAccessKey,
                     resp.sessionToken,
                     resp.expiration,
                     PROVIDER_NAME,
-                    resp.accountId,
                 ).withBusinessMetric(AwsBusinessMetric.Credentials.CREDENTIALS_IMDS)
 
-                creds.also { previousCredentials = it }
+                creds.also {
+                    mu.withLock { previousCredentials = it }
+                }
             }
-            is JsonCredentialsResponse.Error -> when (resp.code) {
-                CODE_ASSUME_ROLE_UNAUTHORIZED_ACCESS -> throw ProviderConfigurationException("Incorrect IMDS/IAM configuration: [${resp.code}] ${resp.message}. Hint: Does this role have a trust relationship with EC2?")
-                else -> throw CredentialsProviderException("Error retrieving credentials from IMDS: code=${resp.code}; ${resp.message}")
+            is JsonCredentialsResponse.Error -> {
+                when (resp.code) {
+                    CODE_ASSUME_ROLE_UNAUTHORIZED_ACCESS -> throw ProviderConfigurationException("Incorrect IMDS/IAM configuration: [${resp.code}] ${resp.message}. Hint: Does this role have a trust relationship with EC2?")
+                    else -> throw CredentialsProviderException("Error retrieving credentials from IMDS: code=${resp.code}; ${resp.message}")
+                }
             }
         }
     }
 
     override fun close() {
-        if (manageClient) {
-            actualClient.close()
+        if (client.isInitialized()) {
+            client.value.close()
         }
     }
 
-    private suspend fun usePreviousCredentials(): Credentials? =
-        previousCredentials?.apply {
+    private suspend fun loadProfile() = try {
+        client.value.get(CREDENTIALS_BASE_PATH)
+    } catch (ex: EC2MetadataError) {
+        if (ex.statusCode == HttpStatusCode.NotFound.value) {
             coroutineContext.info<ImdsCredentialsProvider> {
-                "Attempting to reuse previously-fetched credentials (expiration = $expiration)"
+                "Received 404 from IMDS when loading profile information. Hint: This instance may not have an " +
+                    "IAM role associated."
             }
         }
-
-    override fun toString(): String = this.simpleClassName
-
-    /**
-     * Identifies different versions of IMDS APIs for fetching credentials
-     */
-    private enum class ApiVersion(val urlBase: String) {
-        /**
-         * The original, now-deprecated API
-         */
-        LEGACY("/latest/meta-data/iam/security-credentials/"),
-
-        /**
-         * The new API which provides `AccountId` and potentially other fields in the future
-         */
-        EXTENDED("/latest/meta-data/iam/security-credentials-extended/"),
+        throw ex
     }
-}
-
-internal class ImdsCredentialsException(
-    profileName: String,
-    cause: Throwable,
-) : RuntimeException("Failed to load credentials for EC2 instance profile \"$profileName\"", cause)
 
-internal class ImdsProfileException(cause: Throwable) : RuntimeException("Failed to load instance profile name", cause)
+    private suspend fun useCachedCredentials(ex: Exception): Credentials? = when {
+        ex is IOException || ex is EC2MetadataError && ex.statusCode == HttpStatusCode.InternalServerError.value -> {
+            mu.withLock {
+                previousCredentials?.apply { nextRefresh = clock.now() + DEFAULT_CREDENTIALS_REFRESH_SECONDS.seconds }
+            }
+        }
+        else -> null
+    }
 
-private fun Throwable.wrapAsCredentialsProviderException() =
-    CredentialsProviderException(message.orEmpty(), this)
+    override fun toString(): String = this.simpleClassName
+}

aws-runtime/aws-config/common/src/aws/sdk/kotlin/runtime/auth/credentials/ProfileCredentialsProvider.kt

@@ -106,10 +106,12 @@ public class ProfileCredentialsProvider @InternalSdkApi constructor(
     private val namedProviders = mapOf(
         "Environment" to EnvironmentCredentialsProvider(platformProvider::getenv),
         "Ec2InstanceMetadata" to ImdsCredentialsProvider(
-            instanceProfileName = profileName,
-            client = ImdsClient {
-                platformProvider = this@ProfileCredentialsProvider.platformProvider
-                engine = httpClient
+            profileOverride = profileName,
+            client = lazy {
+                ImdsClient {
+                    platformProvider = this@ProfileCredentialsProvider.platformProvider
+                    engine = httpClient
+                }
             },
             platformProvider = platformProvider,
         ),

aws-runtime/aws-config/common/src/aws/sdk/kotlin/runtime/auth/credentials/internal/CredentialsExt.kt

@@ -29,6 +29,3 @@ internal fun credentials(
     }
     return Credentials(accessKeyId, secretAccessKey, sessionToken, expiration, attributes = attributes)
 }
-
-internal val Credentials.accountId: String?
-    get() = attributes.getOrNull(AwsClientOption.AccountId)