Skip to content

Commit e7e0e5a

Browse files
ianbotsfianbotsf
committed
feat: add support for account ID in IMDS credentials
1 parent ddb53bf commit e7e0e5a

File tree

20 files changed

+1015
-810
lines changed

20 files changed

+1015
-810
lines changed

‎.changes/525a1637-f0f1-4cbd-97dd-5e9c6bcd182e.json‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
{
2+
"id": "525a1637-f0f1-4cbd-97dd-5e9c6bcd182e",
3+
"type": "feature",
4+
"description": "Add support for fetching account ID from IMDS credentials on EC2"
5+
}

‎aws-runtime/aws-config/api/aws-config.api‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -77,12 +77,11 @@ public final class aws/sdk/kotlin/runtime/auth/credentials/EnvironmentCredential
7777

7878
public final class aws/sdk/kotlin/runtime/auth/credentials/ImdsCredentialsProvider : aws/smithy/kotlin/runtime/auth/awscredentials/CloseableCredentialsProvider {
7979
public fun <init> ()V
80+
public fun <init> (Ljava/lang/String;Laws/sdk/kotlin/runtime/config/imds/InstanceMetadataProvider;Laws/smithy/kotlin/runtime/util/PlatformProvider;)V
81+
public synthetic fun <init> (Ljava/lang/String;Laws/sdk/kotlin/runtime/config/imds/InstanceMetadataProvider;Laws/smithy/kotlin/runtime/util/PlatformProvider;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
8082
public fun <init> (Ljava/lang/String;Lkotlin/Lazy;Laws/smithy/kotlin/runtime/util/PlatformEnvironProvider;Laws/smithy/kotlin/runtime/time/Clock;)V
8183
public synthetic fun <init> (Ljava/lang/String;Lkotlin/Lazy;Laws/smithy/kotlin/runtime/util/PlatformEnvironProvider;Laws/smithy/kotlin/runtime/time/Clock;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
8284
public fun close ()V
83-
public final fun getClient ()Lkotlin/Lazy;
84-
public final fun getPlatformProvider ()Laws/smithy/kotlin/runtime/util/PlatformEnvironProvider;
85-
public final fun getProfileOverride ()Ljava/lang/String;
8685
public fun resolve (Laws/smithy/kotlin/runtime/collections/Attributes;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
8786
public fun toString ()Ljava/lang/String;
8887
}
@@ -277,6 +276,7 @@ public final class aws/sdk/kotlin/runtime/config/AwsSdkSetting {
277276
public final fun getAwsContainerCredentialsFullUri ()Laws/smithy/kotlin/runtime/config/EnvironmentSetting;
278277
public final fun getAwsContainerCredentialsRelativeUri ()Laws/smithy/kotlin/runtime/config/EnvironmentSetting;
279278
public final fun getAwsDisableRequestCompression ()Laws/smithy/kotlin/runtime/config/EnvironmentSetting;
279+
public final fun getAwsEc2InstanceProfileName ()Laws/smithy/kotlin/runtime/config/EnvironmentSetting;
280280
public final fun getAwsEc2MetadataDisabled ()Laws/smithy/kotlin/runtime/config/EnvironmentSetting;
281281
public final fun getAwsEc2MetadataServiceEndpoint ()Laws/smithy/kotlin/runtime/config/EnvironmentSetting;
282282
public final fun getAwsEc2MetadataServiceEndpointMode ()Laws/smithy/kotlin/runtime/config/EnvironmentSetting;
@@ -348,8 +348,8 @@ public final class aws/sdk/kotlin/runtime/config/endpoints/ResolversKt {
348348
}
349349

350350
public final class aws/sdk/kotlin/runtime/config/imds/EC2MetadataError : aws/sdk/kotlin/runtime/AwsServiceException {
351-
public fun <init> (ILjava/lang/String;)V
352-
public final fun getStatusCode ()I
351+
public fun <init> (Laws/smithy/kotlin/runtime/http/HttpStatusCode;Ljava/lang/String;)V
352+
public final fun getStatusCode ()Laws/smithy/kotlin/runtime/http/HttpStatusCode;
353353
}
354354

355355
public abstract class aws/sdk/kotlin/runtime/config/imds/EndpointConfiguration {

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import org.jetbrains.dokka.gradle.DokkaTaskPartial
99

1010
plugins {
1111
alias(libs.plugins.aws.kotlin.repo.tools.smithybuild)
12+
alias(libs.plugins.kotlinx.serialization)
1213
}
1314

1415
description = "Support for AWS configuration"
@@ -53,6 +54,7 @@ kotlin {
5354
implementation(libs.kotlinx.coroutines.test)
5455
implementation(libs.smithy.kotlin.http.test)
5556
implementation(libs.kotlinx.serialization.json)
57+
implementation(libs.kotest.framework.datatest)
5658
}
5759
}
5860
jvmTest {

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

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@ import aws.smithy.kotlin.runtime.util.PlatformProvider
4141
* @param region the region to make credentials requests to.
4242
* @return the newly-constructed credentials provider
4343
*/
44-
public class DefaultChainCredentialsProvider constructor(
44+
public class DefaultChainCredentialsProvider(
4545
public val profileName: String? = null,
4646
public val platformProvider: PlatformProvider = PlatformProvider.System,
4747
httpClient: HttpClientEngine? = null,
@@ -59,11 +59,9 @@ public class DefaultChainCredentialsProvider constructor(
5959
ProfileCredentialsProvider(profileName = profileName, platformProvider = platformProvider, httpClient = engine, region = region),
6060
EcsCredentialsProvider(platformProvider, engine),
6161
ImdsCredentialsProvider(
62-
client = lazy {
63-
ImdsClient {
64-
platformProvider = this@DefaultChainCredentialsProvider.platformProvider
65-
engine = this@DefaultChainCredentialsProvider.engine
66-
}
62+
client = ImdsClient {
63+
platformProvider = this@DefaultChainCredentialsProvider.platformProvider
64+
engine = this@DefaultChainCredentialsProvider.engine
6765
},
6866
platformProvider = platformProvider,
6967
),

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

Lines changed: 151 additions & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55

66
package aws.sdk.kotlin.runtime.auth.credentials
77

8+
import aws.sdk.kotlin.runtime.auth.credentials.internal.credentials
89
import aws.sdk.kotlin.runtime.config.AwsSdkSetting
910
import aws.sdk.kotlin.runtime.config.imds.EC2MetadataError
1011
import aws.sdk.kotlin.runtime.config.imds.ImdsClient
@@ -18,137 +19,208 @@ import aws.smithy.kotlin.runtime.http.HttpStatusCode
1819
import aws.smithy.kotlin.runtime.io.IOException
1920
import aws.smithy.kotlin.runtime.serde.json.JsonDeserializer
2021
import aws.smithy.kotlin.runtime.telemetry.logging.info
21-
import aws.smithy.kotlin.runtime.telemetry.logging.warn
2222
import aws.smithy.kotlin.runtime.time.Clock
23-
import aws.smithy.kotlin.runtime.time.Instant
2423
import aws.smithy.kotlin.runtime.util.PlatformEnvironProvider
2524
import aws.smithy.kotlin.runtime.util.PlatformProvider
26-
import kotlinx.coroutines.sync.Mutex
27-
import kotlinx.coroutines.sync.withLock
25+
import aws.smithy.kotlin.runtime.util.SingleFlightGroup
2826
import kotlin.coroutines.coroutineContext
29-
import kotlin.time.Duration.Companion.seconds
3027

31-
private const val CREDENTIALS_BASE_PATH: String = "/latest/meta-data/iam/security-credentials/"
3228
private const val CODE_ASSUME_ROLE_UNAUTHORIZED_ACCESS: String = "AssumeRoleUnauthorizedAccess"
3329
private const val PROVIDER_NAME = "IMDSv2"
3430

3531
/**
3632
* [CredentialsProvider] that uses EC2 instance metadata service (IMDS) to provide credentials information.
37-
* 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)
33+
* This provider requires that the EC2 instance has an
34+
* [instance profile](https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/iam-roles-for-amazon-ec2.html#ec2-instance-profile)
3835
* configured.
39-
*
40-
* See [EC2 IAM Roles](https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/iam-roles-for-amazon-ec2.html) for more
41-
* information.
42-
*
43-
* @param profileOverride override the instance profile name. When retrieving credentials, a call must first be made to
44-
* `<IMDS_BASE_URL>/latest/meta-data/iam/security-credentials/`. This returns the instance profile used. If
45-
* [profileOverride] is set, the initial call to retrieve the profile is skipped and the provided value is used instead.
46-
* @param client the IMDS client to use to resolve credentials information with. This provider takes ownership over
47-
* the lifetime of the given [ImdsClient] and will close it when the provider is closed.
48-
* @param platformProvider the [PlatformEnvironProvider] instance
36+
* @param instanceProfileName overrides the instance profile name. When set, this provider skips querying IMDS for the
37+
* name of the active profile.
38+
* @param client a preconfigured IMDS client with which to retrieve instance metadata. If an instance is passed, the
39+
* caller is responsible for closing it. If no instance is passed, a default instance is created and will be closed when
40+
* this credentials provider is closed.
41+
* @param platformProvider a platform provider used for env vars and system properties
4942
*/
5043
public class ImdsCredentialsProvider(
51-
public val profileOverride: String? = null,
52-
public val client: Lazy<InstanceMetadataProvider> = lazy { ImdsClient() },
53-
public val platformProvider: PlatformEnvironProvider = PlatformProvider.System,
54-
private val clock: Clock = Clock.System,
44+
instanceProfileName: String? = null,
45+
client: InstanceMetadataProvider? = null,
46+
private val platformProvider: PlatformProvider = PlatformProvider.System,
5547
) : CloseableCredentialsProvider {
48+
49+
@Deprecated("This constructor supports parameters which are no longer used in the implementation. It will be removed in version 1.5.")
50+
public constructor(
51+
profileOverride: String? = null,
52+
client: Lazy<InstanceMetadataProvider> = lazy { ImdsClient() },
53+
platformProvider: PlatformEnvironProvider = PlatformProvider.System,
54+
@Suppress("UNUSED_PARAMETER") clock: Clock = Clock.System,
55+
) : this(profileOverride, client.value, platformProvider = platformProvider as? PlatformProvider ?: PlatformProvider.System)
56+
57+
private val manageClient: Boolean = client == null
58+
59+
private val client: InstanceMetadataProvider = client ?: ImdsClient {
60+
this.platformProvider = this@ImdsCredentialsProvider.platformProvider
61+
}
62+
63+
// FIXME This only resolves from env vars and sys props but we need to resolve from profiles too
64+
private val instanceProfileName = instanceProfileName
65+
?: AwsSdkSetting.AwsEc2InstanceProfileName.resolve(platformProvider)
66+
67+
// FIXME This only resolves from env vars and sys props but we need to resolve from profiles too
68+
private val providerDisabled = AwsSdkSetting.AwsEc2MetadataDisabled.resolve(platformProvider) == true
69+
70+
/**
71+
* Tracks the known-good version of IMDS APIs available in the local environment. This starts as `null` and will be
72+
* updated after the first successful API call.
73+
*/
74+
private var apiVersion: ApiVersion? = null
75+
76+
private val urlBase: String
77+
get() = (apiVersion ?: ApiVersion.EXTENDED).urlBase
78+
5679
private var previousCredentials: Credentials? = null
5780

58-
// the time to refresh the Credentials. If set, it will take precedence over the Credentials' expiration time
59-
private var nextRefresh: Instant? = null
81+
/**
82+
* Tracks the instance profile name resolved from IMDS. This starts as `null` and will be updated after a
83+
* successful API call. Note that if [instanceProfileName] is set, profile name resolution will be skipped.
84+
*/
85+
private var resolvedProfileName: String? = null
86+
87+
/**
88+
* A deduplicator for resolving credentials and tracking mutable state about IMDS
89+
*/
90+
private val sfg = SingleFlightGroup<Credentials>()
6091

61-
// protects previousCredentials and nextRefresh
62-
private val mu = Mutex()
92+
override suspend fun resolve(attributes: Attributes): Credentials = sfg.singleFlight { resolveUnderLock() }
6393

64-
override suspend fun resolve(attributes: Attributes): Credentials {
65-
if (AwsSdkSetting.AwsEc2MetadataDisabled.resolve(platformProvider) == true) {
94+
private suspend fun resolveUnderLock(): Credentials {
95+
println("**** Resolving creds (instanceProfileName=$instanceProfileName; apiVersion=$apiVersion; urlBase=$urlBase)")
96+
97+
if (providerDisabled) {
98+
println("**** Explicitly disabled")
6699
throw CredentialsNotLoadedException("AWS EC2 metadata is explicitly disabled; credentials not loaded")
67100
}
68101

69-
// if we have previously served IMDS credentials and it's not time for a refresh, just return the previous credentials
70-
mu.withLock {
71-
previousCredentials?.run {
72-
nextRefresh?.takeIf { clock.now() < it }?.run {
73-
return previousCredentials!!
102+
val profileName = instanceProfileName ?: resolvedProfileName ?: try {
103+
println("**** Resolving profile")
104+
client.get(urlBase).also {
105+
if (apiVersion == null) {
106+
// Tried EXTENDED and it worked; remember that for the future
107+
apiVersion = ApiVersion.EXTENDED
74108
}
75109
}
76-
}
110+
} catch (ex: EC2MetadataError) {
111+
when {
112+
apiVersion == null && ex.statusCode == HttpStatusCode.NotFound -> {
113+
// Tried EXTENDED and that didn't work; fallback to LEGACY
114+
apiVersion = ApiVersion.LEGACY
115+
return resolveUnderLock()
116+
}
117+
118+
ex.statusCode == HttpStatusCode.NotFound -> {
119+
coroutineContext.info<ImdsCredentialsProvider> {
120+
"Received 404 when loading profile name. This instance may not have an associated profile."
121+
}
122+
throw ex
123+
}
77124

78-
val profileName = try {
79-
profileOverride ?: loadProfile()
125+
else -> return usePreviousCredentials()
126+
?: throw ImdsProfileException(ex).wrapAsCredentialsProviderException()
127+
}
128+
} catch (ex: IOException) {
129+
return usePreviousCredentials() ?: throw ImdsProfileException(ex).wrapAsCredentialsProviderException()
80130
} catch (ex: Exception) {
81-
return useCachedCredentials(ex) ?: throw CredentialsProviderException("failed to load instance profile", ex)
131+
throw ImdsProfileException(ex).wrapAsCredentialsProviderException()
82132
}
83133

84-
val payload = try {
85-
client.value.get("$CREDENTIALS_BASE_PATH$profileName")
134+
val credsPayload = try {
135+
client.get("$urlBase$profileName")
136+
} catch (ex: EC2MetadataError) {
137+
when {
138+
apiVersion == null && ex.statusCode == HttpStatusCode.NotFound -> {
139+
// Tried EXTENDED and that didn't work; fallback to LEGACY
140+
apiVersion = ApiVersion.LEGACY
141+
return resolveUnderLock()
142+
}
143+
144+
instanceProfileName == null && ex.statusCode == HttpStatusCode.NotFound -> {
145+
// A previously-resolved profile is now invalid; forget the resolved name and re-resolve
146+
resolvedProfileName = null
147+
return resolveUnderLock()
148+
}
149+
150+
else -> return usePreviousCredentials()
151+
?: throw ImdsCredentialsException(profileName, ex).wrapAsCredentialsProviderException()
152+
}
153+
} catch (ex: IOException) {
154+
return usePreviousCredentials()
155+
?: throw ImdsCredentialsException(profileName, ex).wrapAsCredentialsProviderException()
86156
} catch (ex: Exception) {
87-
return useCachedCredentials(ex) ?: throw CredentialsProviderException("failed to load credentials", ex)
157+
throw ImdsCredentialsException(profileName, ex).wrapAsCredentialsProviderException()
88158
}
89159

90-
val deserializer = JsonDeserializer(payload.encodeToByteArray())
160+
if (instanceProfileName == null) {
161+
// No profile name was provided at construction time; cache the resolved name
162+
resolvedProfileName = profileName
163+
}
164+
165+
val deserializer = JsonDeserializer(credsPayload.encodeToByteArray())
91166

92167
return when (val resp = deserializeJsonCredentials(deserializer)) {
93168
is JsonCredentialsResponse.SessionCredentials -> {
94-
nextRefresh = if (resp.expiration != null && resp.expiration < clock.now()) {
95-
coroutineContext.warn<ImdsCredentialsProvider> {
96-
"Attempting credential expiration extension due to a credential service availability issue. " +
97-
"A refresh of these credentials will be attempted again in " +
98-
"${ DEFAULT_CREDENTIALS_REFRESH_SECONDS / 60 } minutes."
99-
}
100-
clock.now() + DEFAULT_CREDENTIALS_REFRESH_SECONDS.seconds
101-
} else {
102-
null
103-
}
104-
105-
val creds = Credentials(
169+
val creds = credentials(
106170
resp.accessKeyId,
107171
resp.secretAccessKey,
108172
resp.sessionToken,
109173
resp.expiration,
110174
PROVIDER_NAME,
175+
resp.accountId,
111176
).withBusinessMetric(AwsBusinessMetric.Credentials.CREDENTIALS_IMDS)
112177

113-
creds.also {
114-
mu.withLock { previousCredentials = it }
115-
}
178+
creds.also { previousCredentials = it }
116179
}
117-
is JsonCredentialsResponse.Error -> {
118-
when (resp.code) {
119-
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?")
120-
else -> throw CredentialsProviderException("Error retrieving credentials from IMDS: code=${resp.code}; ${resp.message}")
121-
}
180+
is JsonCredentialsResponse.Error -> when (resp.code) {
181+
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?")
182+
else -> throw CredentialsProviderException("Error retrieving credentials from IMDS: code=${resp.code}; ${resp.message}")
122183
}
123184
}
124185
}
125186

126187
override fun close() {
127-
if (client.isInitialized()) {
128-
client.value.close()
188+
if (manageClient) {
189+
client.close()
129190
}
130191
}
131192

132-
private suspend fun loadProfile() = try {
133-
client.value.get(CREDENTIALS_BASE_PATH)
134-
} catch (ex: EC2MetadataError) {
135-
if (ex.statusCode == HttpStatusCode.NotFound.value) {
193+
private suspend fun usePreviousCredentials(): Credentials? =
194+
previousCredentials?.apply {
136195
coroutineContext.info<ImdsCredentialsProvider> {
137-
"Received 404 from IMDS when loading profile information. Hint: This instance may not have an " +
138-
"IAM role associated."
196+
"Attempting to reuse previously-fetched credentials (expiration = $expiration)"
139197
}
140198
}
141-
throw ex
142-
}
143-
144-
private suspend fun useCachedCredentials(ex: Exception): Credentials? = when {
145-
ex is IOException || ex is EC2MetadataError && ex.statusCode == HttpStatusCode.InternalServerError.value -> {
146-
mu.withLock {
147-
previousCredentials?.apply { nextRefresh = clock.now() + DEFAULT_CREDENTIALS_REFRESH_SECONDS.seconds }
148-
}
149-
}
150-
else -> null
151-
}
152199

153200
override fun toString(): String = this.simpleClassName
201+
202+
/**
203+
* Identifies different versions of IMDS APIs for fetching credentials
204+
*/
205+
private enum class ApiVersion(val urlBase: String) {
206+
/**
207+
* The original, now-deprecated API
208+
*/
209+
LEGACY("/latest/meta-data/iam/security-credentials/"),
210+
211+
/**
212+
* The new API which provides `AccountId` and potentially other fields in the future
213+
*/
214+
EXTENDED("/latest/meta-data/iam/security-credentials-extended/"),
215+
}
154216
}
217+
218+
internal class ImdsCredentialsException(
219+
profileName: String,
220+
cause: Throwable,
221+
) : RuntimeException("Failed to load credentials for EC2 instance profile \"$profileName\"", cause)
222+
223+
internal class ImdsProfileException(cause: Throwable) : RuntimeException("Failed to load instance profile name", cause)
224+
225+
private fun Throwable.wrapAsCredentialsProviderException() =
226+
CredentialsProviderException(message.orEmpty(), this)

0 commit comments

Comments
 (0)