fix: enable DefaultRegionProviderChainTest to pass in environments where an IMDS client is already present (#369)

Author: ianbotsf

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

@@ -8,6 +8,7 @@ package aws.sdk.kotlin.runtime.auth.credentials
 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.config.resolve
 import aws.smithy.kotlin.runtime.http.HttpStatusCode
 import aws.smithy.kotlin.runtime.io.Closeable
@@ -36,7 +37,7 @@ private const val CODE_ASSUME_ROLE_UNAUTHORIZED_ACCESS: String = "AssumeRoleUnau
  */
 public class ImdsCredentialsProvider(
     private val profileOverride: String? = null,
-    private val client: Lazy<ImdsClient> = lazy { ImdsClient() },
+    private val client: Lazy<InstanceMetadataProvider> = lazy { ImdsClient() },
     private val platformProvider: PlatformEnvironProvider = Platform
 ) : CredentialsProvider, Closeable {
     private val logger = Logger.getLogger<ImdsCredentialsProvider>()

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

@@ -38,6 +38,16 @@ internal const val DEFAULT_MAX_RETRIES: UInt = 3u
 
 private const val SERVICE = "imds"
 
+/**
+ * Represents a generic client that can fetch instance metadata.
+ */
+public interface InstanceMetadataProvider : Closeable {
+    /**
+     * Gets the specified instance metadata value by the given path.
+     */
+    public suspend fun get(path: String): String
+}
+
 /**
  * IMDSv2 Client
  *
@@ -48,7 +58,7 @@ private const val SERVICE = "imds"
  * for more information.
  */
 @OptIn(ExperimentalTime::class)
-public class ImdsClient private constructor(builder: Builder) : Closeable {
+public class ImdsClient private constructor(builder: Builder) : InstanceMetadataProvider {
     public constructor() : this(Builder())
 
     private val logger = Logger.getLogger<ImdsClient>()
@@ -115,7 +125,7 @@ public class ImdsClient private constructor(builder: Builder) : Closeable {
      * val amiId = client.get("/latest/meta-data/ami-id")
      * ```
      */
-    public suspend fun get(path: String): String {
+    public override suspend fun get(path: String): String {
         val op = SdkHttpOperation.build<Unit, String> {
             serializer = UnitSerializer
             deserializer = object : HttpDeserialize<String> {

aws-runtime/aws-config/common/src/aws/sdk/kotlin/runtime/region/DefaultRegionProviderChain.kt

@@ -5,6 +5,8 @@
 
 package aws.sdk.kotlin.runtime.region
 
+import aws.sdk.kotlin.runtime.config.imds.ImdsClient
+import aws.sdk.kotlin.runtime.config.imds.InstanceMetadataProvider
 import aws.smithy.kotlin.runtime.io.Closeable
 import aws.smithy.kotlin.runtime.util.Platform
 import aws.smithy.kotlin.runtime.util.PlatformProvider
@@ -17,5 +19,6 @@ import aws.smithy.kotlin.runtime.util.PlatformProvider
  *  4. If running on EC2, check the EC2 metadata service for region
  */
 internal expect class DefaultRegionProviderChain constructor(
-    platformProvider: PlatformProvider = Platform
+    platformProvider: PlatformProvider = Platform,
+    imdsClient: Lazy<InstanceMetadataProvider> = lazy { ImdsClient() },
 ) : RegionProvider, Closeable

aws-runtime/aws-config/common/src/aws/sdk/kotlin/runtime/region/ImdsRegionProvider.kt

@@ -7,6 +7,7 @@ package aws.sdk.kotlin.runtime.region
 
 import aws.sdk.kotlin.runtime.config.AwsSdkSetting
 import aws.sdk.kotlin.runtime.config.imds.ImdsClient
+import aws.sdk.kotlin.runtime.config.imds.InstanceMetadataProvider
 import aws.sdk.kotlin.runtime.config.resolve
 import aws.smithy.kotlin.runtime.io.Closeable
 import aws.smithy.kotlin.runtime.util.Platform
@@ -22,7 +23,7 @@ private const val REGION_PATH: String = "/latest/meta-data/placement/region"
  * @param platformProvider the [PlatformEnvironProvider] instance
  */
 internal class ImdsRegionProvider(
-    private val client: Lazy<ImdsClient> = lazy { ImdsClient() },
+    private val client: Lazy<InstanceMetadataProvider> = lazy { ImdsClient() },
     private val platformProvider: PlatformEnvironProvider = Platform,
 ) : RegionProvider, Closeable {
     private val resolvedRegion = asyncLazy(::loadRegion)

aws-runtime/aws-config/common/test/aws/sdk/kotlin/runtime/region/DefaultRegionProviderChainTest.kt

@@ -7,6 +7,7 @@ package aws.sdk.kotlin.runtime.region
 
 import aws.sdk.kotlin.runtime.testing.TestPlatformProvider
 import aws.sdk.kotlin.runtime.testing.runSuspendTest
+import aws.sdk.kotlin.runtime.util.TestInstanceMetadataProvider
 import kotlinx.serialization.json.*
 import kotlin.test.Test
 import kotlin.test.assertEquals
@@ -15,6 +16,7 @@ class DefaultRegionProviderChainTest {
     private data class RegionProviderChainTest(
         val name: String,
         val platformProvider: TestPlatformProvider,
+        val instanceMetadataProvider: TestInstanceMetadataProvider,
         val region: String?,
         val targets: List<String> = emptyList()
     )
@@ -26,12 +28,16 @@ class DefaultRegionProviderChainTest {
             .map {
                 val name = it["name"]!!.jsonPrimitive.content
                 val platform = TestPlatformProvider.fromJsonNode(it["platform"]!!.jsonObject)
+                val instanceMetadata = TestInstanceMetadataProvider.fromJsonNode(it["imds"]!!.jsonObject)
                 val region = it["region"]!!.jsonPrimitive.contentOrNull
-                RegionProviderChainTest(name, platform, region)
+                RegionProviderChainTest(name, platform, instanceMetadata, region)
             }
 
         tests.forEach { test ->
-            val provider = DefaultRegionProviderChain(test.platformProvider)
+            val provider = DefaultRegionProviderChain(
+                platformProvider = test.platformProvider,
+                imdsClient = lazy { test.instanceMetadataProvider },
+            )
             val actual = provider.getRegion()
             assertEquals(test.region, actual, test.name)
         }
@@ -62,6 +68,14 @@ fun TestPlatformProvider.Companion.fromJsonNode(obj: JsonObject): TestPlatformPr
     return TestPlatformProvider(env, props, fs)
 }
 
+/**
+ * Construct a [TestInstanceMetadataProvider] from a JSON object containing metadata as key-value pairs.
+ */
+fun TestInstanceMetadataProvider.Companion.fromJsonNode(obj: JsonObject): TestInstanceMetadataProvider {
+    val metadata = obj.jsonObject.mapValues { it.value.jsonPrimitive.content }
+    return TestInstanceMetadataProvider(metadata)
+}
+
 // language=JSON
 private const val regionProviderChainTestSuite = """
 [
@@ -72,6 +86,7 @@ private const val regionProviderChainTestSuite = """
             "props": {},
             "fs": {}
         },
+        "imds": {},
         "region": null
     },
     {
@@ -83,6 +98,7 @@ private const val regionProviderChainTestSuite = """
             "props": {},
             "fs": {}
         },
+        "imds": {},
         "region": "us-east-2"
     },
     {
@@ -96,6 +112,7 @@ private const val regionProviderChainTestSuite = """
             },
             "fs": {}
         },
+        "imds": {},
         "region": "us-west-1"
     },
     {
@@ -109,6 +126,7 @@ private const val regionProviderChainTestSuite = """
               "config": "[default]\nregion = us-east-2"
             }
         },
+        "imds": {},
         "region": "us-east-2"
     },
     {
@@ -123,7 +141,64 @@ private const val regionProviderChainTestSuite = """
               "config": "[default]\nregion = us-east-2\n[profile test-profile]\nregion = us-west-1"
             }
         },
+        "imds": {},
         "region": "us-west-1"
+    },
+    {
+        "name": "imds configured",
+        "platform": {
+            "env": {},
+            "props": {},
+            "fs": {}
+        },
+        "imds": {
+            "/latest/meta-data/placement/region": "us-east-1"
+        },
+        "region": "us-east-1"
+    },
+    {
+        "name": "jvm system properties are favored over imds",
+        "platform": {
+            "env": {
+                "AWS_REGION": "us-east-2"
+            },
+            "props": {},
+            "fs": {}
+        },
+        "imds": {
+            "/latest/meta-data/placement/region": "us-east-1"
+        },
+        "region": "us-east-2"
+    },
+    {
+        "name": "environment variables are favored over imds",
+        "platform": {
+            "env": {},
+            "props": {
+                "aws.region": "us-west-1"
+            },
+            "fs": {}
+        },
+        "imds": {
+            "/latest/meta-data/placement/region": "us-east-1"
+        },
+        "region": "us-west-1"
+    },
+    {
+        "name": "profile is favored over imds",
+        "platform": {
+            "env": {
+              "AWS_CONFIG_FILE": "config"
+            },
+            "props": {},
+            "fs": {
+              "config": "[default]\nregion = us-west-2"
+            }
+        },
+        "imds": {
+            "/latest/meta-data/placement/region": "us-east-1"
+        },
+        "region": "us-west-2"
     }
 ]
 """

aws-runtime/aws-config/common/test/aws/sdk/kotlin/runtime/util/TestInstanceMetadataProvider.kt

@@ -0,0 +1,15 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0.
+ */
+
+package aws.sdk.kotlin.runtime.util
+
+import aws.sdk.kotlin.runtime.config.imds.InstanceMetadataProvider
+
+public class TestInstanceMetadataProvider(private val metadata: Map<String, String>) : InstanceMetadataProvider {
+    public companion object { }
+
+    override fun close(): Unit = Unit
+    override suspend fun get(path: String): String = metadata[path] ?: throw IllegalArgumentException("$path missing")
+}

aws-runtime/aws-config/jvm/src/aws/sdk/kotlin/runtime/region/DefaultRegionProviderChainJVM.kt

@@ -5,18 +5,20 @@
 
 package aws.sdk.kotlin.runtime.region
 
+import aws.sdk.kotlin.runtime.config.imds.InstanceMetadataProvider
 import aws.smithy.kotlin.runtime.io.Closeable
 import aws.smithy.kotlin.runtime.util.PlatformProvider
 
 internal actual class DefaultRegionProviderChain actual constructor(
-    platformProvider: PlatformProvider
+    platformProvider: PlatformProvider,
+    imdsClient: Lazy<InstanceMetadataProvider>,
 ) : RegionProvider,
     Closeable,
     RegionProviderChain(
         JvmSystemPropRegionProvider(platformProvider),
         EnvironmentRegionProvider(platformProvider),
         ProfileRegionProvider(platformProvider),
-        ImdsRegionProvider(platformProvider = platformProvider)
+        ImdsRegionProvider(client = imdsClient, platformProvider = platformProvider)
     ) {
 
     override fun close() {