feat: IMDSv2 client (#330)

Author: aajtodd

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

@@ -0,0 +1,34 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0.
+ */
+
+description = "Support for AWS configuration"
+extra["moduleName"] = "aws.sdk.kotlin.runtime.config"
+
+val smithyKotlinVersion: String by project
+
+kotlin {
+    sourceSets {
+        commonMain {
+            dependencies {
+                api(project(":aws-runtime:aws-core"))
+                implementation("aws.smithy.kotlin:logging:$smithyKotlinVersion")
+                implementation("aws.smithy.kotlin:http:$smithyKotlinVersion")
+                implementation("aws.smithy.kotlin:utils:$smithyKotlinVersion")
+                implementation(project(":aws-runtime:http-client-engine-crt"))
+                implementation(project(":aws-runtime:protocols:http"))
+            }
+        }
+        commonTest {
+            dependencies {
+                implementation(project(":aws-runtime:testing"))
+                implementation("aws.smithy.kotlin:http-test:$smithyKotlinVersion")
+                val kotlinxSerializationVersion: String by project
+                val mockkVersion: String by project
+                implementation("org.jetbrains.kotlinx:kotlinx-serialization-json:$kotlinxSerializationVersion")
+                implementation("io.mockk:mockk:$mockkVersion")
+            }
+        }
+    }
+}

aws-runtime/aws-config/common/src/aws/sdk/kotlin/runtime/config/CachedValue.kt

@@ -0,0 +1,64 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0.
+ */
+
+package aws.sdk.kotlin.runtime.config
+
+import aws.smithy.kotlin.runtime.time.Clock
+import aws.smithy.kotlin.runtime.time.Instant
+import kotlinx.coroutines.sync.Mutex
+import kotlinx.coroutines.sync.withLock
+import kotlin.time.Duration
+import kotlin.time.ExperimentalTime
+
+/**
+ * A value with an expiration
+ */
+internal data class ExpiringValue<T> (val value: T, val expiresAt: Instant)
+
+/**
+ * Expiry aware value
+ *
+ * @param value The value that expires
+ * @param bufferTime The amount of time before the actual expiration time when the value is considered expired. By default
+ * the buffer time is zero meaning the value expires at the expiration time. A non-zero buffer time means the value will
+ * expire BEFORE the actual expiration.
+ * @param clock The clock to use for system time
+ */
+@OptIn(ExperimentalTime::class)
+internal class CachedValue<T> (
+    private var value: ExpiringValue<T>? = null,
+    private val bufferTime: Duration = Duration.seconds(0),
+    private val clock: Clock = Clock.System
+) {
+    constructor(value: T, expiresAt: Instant, bufferTime: Duration = Duration.seconds(0), clock: Clock = Clock.System) : this(ExpiringValue(value, expiresAt), bufferTime, clock)
+    private val mu = Mutex()
+
+    /**
+     * Check if the value is expired or not as compared to the time [now]
+     */
+    suspend fun isExpired(): Boolean = mu.withLock { isExpiredUnlocked() }
+
+    private fun isExpiredUnlocked(): Boolean {
+        val curr = value ?: return true
+        return clock.now() >= (curr.expiresAt - bufferTime)
+    }
+
+    /**
+     * Get the value if it has not expired yet. Returns null if the value has expired
+     */
+    suspend fun get(): T? = mu.withLock {
+        if (!isExpiredUnlocked()) return value!!.value else null
+    }
+
+    /**
+     * Attempt to get the value or refresh it with [initializer] if it is expired
+     */
+    suspend fun getOrLoad(initializer: suspend () -> ExpiringValue<T>): T = mu.withLock {
+        if (!isExpiredUnlocked()) return@withLock value!!.value
+
+        val refreshed = initializer().also { value = it }
+        return refreshed.value
+    }
+}

aws-runtime/aws-config/common/src/aws/sdk/kotlin/runtime/config/imds/ImdsClient.kt

@@ -0,0 +1,233 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0.
+ */
+
+package aws.sdk.kotlin.runtime.config.imds
+
+import aws.sdk.kotlin.runtime.AwsServiceException
+import aws.sdk.kotlin.runtime.ConfigurationException
+import aws.sdk.kotlin.runtime.client.AwsClientOption
+import aws.sdk.kotlin.runtime.endpoint.Endpoint
+import aws.sdk.kotlin.runtime.http.ApiMetadata
+import aws.sdk.kotlin.runtime.http.AwsUserAgentMetadata
+import aws.sdk.kotlin.runtime.http.engine.crt.CrtHttpEngine
+import aws.sdk.kotlin.runtime.http.middleware.ServiceEndpointResolver
+import aws.sdk.kotlin.runtime.http.middleware.UserAgent
+import aws.smithy.kotlin.runtime.client.ExecutionContext
+import aws.smithy.kotlin.runtime.http.*
+import aws.smithy.kotlin.runtime.http.engine.HttpClientEngine
+import aws.smithy.kotlin.runtime.http.engine.HttpClientEngineConfig
+import aws.smithy.kotlin.runtime.http.operation.*
+import aws.smithy.kotlin.runtime.http.response.HttpResponse
+import aws.smithy.kotlin.runtime.io.Closeable
+import aws.smithy.kotlin.runtime.io.middleware.Phase
+import aws.smithy.kotlin.runtime.logging.Logger
+import aws.smithy.kotlin.runtime.time.Clock
+import aws.smithy.kotlin.runtime.util.Platform
+import aws.smithy.kotlin.runtime.util.PlatformProvider
+import kotlin.time.Duration
+import kotlin.time.ExperimentalTime
+
+/**
+ * Maximum time allowed by default (6 hours)
+ */
+internal const val DEFAULT_TOKEN_TTL_SECONDS: Int = 21_600
+internal const val DEFAULT_MAX_RETRIES: UInt = 3u
+
+private const val SERVICE = "imds"
+
+/**
+ * IMDSv2 Client
+ *
+ * This client supports fetching tokens, retrying failures, and token caching according to the specified TTL.
+ *
+ * NOTE: This client ONLY supports IMDSv2. It will not fallback to IMDSv1.
+ * See [transitioning to IMDSv2](https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/configuring-instance-metadata-service.html#instance-metadata-transition-to-version-2)
+ * for more information.
+ */
+@OptIn(ExperimentalTime::class)
+public class ImdsClient private constructor(builder: Builder) : Closeable {
+    public constructor() : this(Builder())
+
+    private val logger = Logger.getLogger<ImdsClient>()
+
+    private val maxRetries: UInt = builder.maxRetries
+    private val endpointConfiguration: EndpointConfiguration = builder.endpointConfiguration
+    private val tokenTtl: Duration = builder.tokenTTL
+    private val clock: Clock = builder.clock
+    private val platformProvider: PlatformProvider = builder.platformProvider
+
+    init {
+        // validate the override at construction time
+        if (endpointConfiguration is EndpointConfiguration.Custom) {
+            val url = endpointConfiguration.endpoint.toUrl()
+            try {
+                Url.parse(url.toString())
+            } catch (ex: Exception) {
+                throw ConfigurationException("Invalid endpoint configuration: `$url` is not a valid URI", ex)
+            }
+        }
+    }
+
+    // TODO connect/socket timeouts
+    private val httpClient = sdkHttpClient(builder.engine ?: CrtHttpEngine(HttpClientEngineConfig()))
+
+    // cached middleware instances
+    private val middleware: List<Feature> = listOf(
+        ServiceEndpointResolver.create {
+            serviceId = SERVICE
+            resolver = ImdsEndpointResolver(platformProvider, endpointConfiguration)
+        },
+        UserAgent.create {
+            metadata = AwsUserAgentMetadata.fromEnvironment(ApiMetadata(SERVICE, "unknown"))
+        },
+        TokenMiddleware.create {
+            httpClient = this@ImdsClient.httpClient
+            ttl = tokenTtl
+            clock = this@ImdsClient.clock
+        }
+    )
+
+    public companion object {
+        public operator fun invoke(block: Builder.() -> Unit): ImdsClient = ImdsClient(Builder().apply(block))
+    }
+
+    /**
+     * Retrieve information from instance metadata service (IMDS).
+     *
+     * This method will combine [path] with the configured endpoint and return the response as a string.
+     *
+     * For more information about IMDSv2 methods and functionality, see
+     * [Instance metadata and user data](https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/ec2-instance-metadata.html)
+     *
+     * Example:
+     *
+     * ```kotlin
+     * val client = EC2Metadata()
+     * val amiId = client.get("/latest/meta-data/ami-id")
+     * ```
+     */
+    public suspend fun get(path: String): String {
+        val op = SdkHttpOperation.build<Unit, String> {
+            serializer = UnitSerializer
+            deserializer = object : HttpDeserialize<String> {
+                override suspend fun deserialize(context: ExecutionContext, response: HttpResponse): String {
+                    if (response.status.isSuccess()) {
+                        val payload = response.body.readAll() ?: throw EC2MetadataError(response.status.value, "no metadata payload")
+                        return payload.decodeToString()
+                    } else {
+                        throw EC2MetadataError(response.status.value, "error retrieving instance metadata")
+                    }
+                }
+            }
+            context {
+                operationName = path
+                service = SERVICE
+                // artifact of re-using ServiceEndpointResolver middleware
+                set(AwsClientOption.Region, "not-used")
+            }
+        }
+        middleware.forEach { it.install(op) }
+        op.execution.mutate.intercept(Phase.Order.Before) { req, next ->
+            req.subject.url.path = path
+            next.call(req)
+        }
+
+        // TODO - retries
+        return op.roundTrip(httpClient, Unit)
+    }
+
+    override fun close() {
+        httpClient.close()
+    }
+
+    public class Builder {
+        /**
+         * The maximum number of retries for fetching tokens and metadata
+         */
+        public var maxRetries: UInt = DEFAULT_MAX_RETRIES
+
+        /**
+         * The endpoint configuration to use when making requests
+         */
+        public var endpointConfiguration: EndpointConfiguration = EndpointConfiguration.Default
+
+        /**
+         * Override the time-to-live for the session token
+         */
+        public var tokenTTL: Duration = Duration.seconds(DEFAULT_TOKEN_TTL_SECONDS)
+
+        /**
+         * The HTTP engine to use to make requests with. This is here to facilitate testing and can otherwise be ignored
+         */
+        internal var engine: HttpClientEngine? = null
+
+        /**
+         * The source of time for token refreshes. This is here to facilitate testing and can otherwise be ignored
+         */
+        internal var clock: Clock = Clock.System
+
+        /**
+         * The platform provider. This is here to facilitate testing and can otherwise be ignored
+         */
+        internal var platformProvider: PlatformProvider = Platform
+    }
+}
+
+public sealed class EndpointConfiguration {
+    /**
+     * Detected from the execution environment
+     */
+    public object Default : EndpointConfiguration()
+
+    /**
+     * Override the endpoint to make requests to
+     */
+    public data class Custom(val endpoint: Endpoint) : EndpointConfiguration()
+
+    /**
+     * Override the [EndpointMode] to use
+     */
+    public data class ModeOverride(val mode: EndpointMode) : EndpointConfiguration()
+}
+
+public enum class EndpointMode(internal val defaultEndpoint: Endpoint) {
+    /**
+     * IPv4 mode. This is the default unless otherwise specified
+     * e.g. `http://169.254.169.254'
+     */
+    IPv4(Endpoint("169.254.169.254", "http")),
+
+    /**
+     * IPv6 mode
+     * e.g. `http://[fd00:ec2::254]`
+     */
+    IPv6(Endpoint("[fd00:ec2::254]", "http"));
+
+    public companion object {
+        public fun fromValue(value: String): EndpointMode = when (value.lowercase()) {
+            "ipv4" -> IPv4
+            "ipv6" -> IPv6
+            else -> throw IllegalArgumentException("invalid EndpointMode: `$value`")
+        }
+    }
+}
+
+/**
+ * Exception thrown when an error occurs retrieving metadata from IMDS
+ *
+ * @param statusCode The raw HTTP status code of the response
+ * @param message The error message
+ */
+public class EC2MetadataError(public val statusCode: Int, message: String) : AwsServiceException(message)
+
+private fun Endpoint.toUrl(): Url {
+    val endpoint = this
+    val protocol = Protocol.parse(endpoint.protocol)
+    return Url(
+        scheme = protocol,
+        host = endpoint.hostname,
+        port = endpoint.port ?: protocol.defaultPort,
+    )
+}