Rename service-specific token generators, add more optional parameter…

…s, change [Credentials] param to [CredentialsProvider]
awslabs · Jan 9, 2025 · 3b9b2b2 · 3b9b2b2
1 parent 68fa563
commit 3b9b2b2
Show file tree

Hide file tree

Showing 5 changed files with 89 additions and 45 deletions.
diff --git a/aws-runtime/aws-config/common/src/aws/sdk/kotlin/runtime/auth/AuthTokenGenerator.kt b/aws-runtime/aws-config/common/src/aws/sdk/kotlin/runtime/auth/AuthTokenGenerator.kt
@@ -6,43 +6,63 @@ package aws.sdk.kotlin.runtime.auth
 
 import aws.sdk.kotlin.runtime.auth.credentials.DefaultChainCredentialsProvider
 import aws.smithy.kotlin.runtime.auth.awscredentials.Credentials
+import aws.smithy.kotlin.runtime.auth.awscredentials.CredentialsProvider
 import aws.smithy.kotlin.runtime.auth.awssigning.AwsSignatureType
+import aws.smithy.kotlin.runtime.auth.awssigning.AwsSigner
 import aws.smithy.kotlin.runtime.auth.awssigning.AwsSigningConfig
 import aws.smithy.kotlin.runtime.auth.awssigning.AwsSigningConfig.Companion.invoke
 import aws.smithy.kotlin.runtime.auth.awssigning.DefaultAwsSigner
 import aws.smithy.kotlin.runtime.http.HttpMethod
 import aws.smithy.kotlin.runtime.http.request.HttpRequest
 import aws.smithy.kotlin.runtime.net.url.Url
 import aws.smithy.kotlin.runtime.time.Clock
-import kotlinx.coroutines.runBlocking
+import aws.smithy.kotlin.runtime.util.ExpiringValue
 import kotlin.time.Duration
+import kotlin.time.Duration.Companion.minutes
+import kotlin.time.Duration.Companion.seconds
+
+// The default expiration value to use for [Credentials] when none is provided.
+private val DEFAULT_CREDENTIALS_EXPIRATION = 10.minutes
 
 /**
  * Generates an authentication token, which is a SigV4-signed URL with the HTTP scheme removed.
  * @param service The name of the service the token is being generated for
- * @param credentials The credentials to use when generating the auth token, defaults to resolving credentials from the [DefaultChainCredentialsProvider]
+ * @param credentialsProvider The [CredentialsProvider] which will provide credentials to use when generating the auth token, defaults to [DefaultChainCredentialsProvider]
+ * @param credentialsRefreshBuffer The amount of time before the resolved [Credentials] expire in which they are considered expired, defaults to 10 seconds.
+ * @param signer The [AwsSigner] implementation to use when creating the authentication token, defaults to [DefaultAwsSigner]
+ * @param clock The [Clock] implementation to use
  */
 public class AuthTokenGenerator(
     public val service: String,
-    public val credentials: Credentials? = runBlocking { DefaultChainCredentialsProvider().resolve() },
+    public val credentialsProvider: CredentialsProvider = DefaultChainCredentialsProvider(),
+    public val credentialsRefreshBuffer: Duration = 10.seconds,
+    public val signer: AwsSigner = DefaultAwsSigner,
+    public val clock: Clock = Clock.System
 ) {
-    private fun String.trimScheme() = removePrefix("http://").removePrefix("https://")
+    private lateinit var credentials: ExpiringValue<Credentials>
+
+    private fun Url.trimScheme(): String = toString().removePrefix(scheme.protocolName).removePrefix("://")
 
     public suspend fun generateAuthToken(endpoint: Url, region: String, expiration: Duration): String {
+        if (!::credentials.isInitialized || (credentials.expiresAt - clock.now()).absoluteValue <= credentialsRefreshBuffer) {
+            val resolved = credentialsProvider.resolve()
+            credentials = ExpiringValue(resolved, resolved.expiration ?: (clock.now() + DEFAULT_CREDENTIALS_EXPIRATION))
+        }
+
         val req = HttpRequest(HttpMethod.GET, endpoint)
 
-        val creds = credentials
+        val creds = credentials.value
         val serv = service
 
         val config = AwsSigningConfig {
             credentials = creds
             this.region = region
             service = serv
-            signingDate = Clock.System.now()
+            signingDate = clock.now()
             expiresAfter = expiration
             signatureType = AwsSignatureType.HTTP_REQUEST_VIA_QUERY_PARAMS
         }
 
-        return DefaultAwsSigner.sign(req, config).output.url.toString().trimScheme()
+        return signer.sign(req, config).output.url.trimScheme()
     }
 }
diff --git a/...otlin/services/dsql/AuthTokenGenerator.kt → ...n/services/dsql/DsqlAuthTokenGenerator.kt b/...otlin/services/dsql/AuthTokenGenerator.kt → ...n/services/dsql/DsqlAuthTokenGenerator.kt
@@ -7,19 +7,28 @@ package aws.sdk.kotlin.services.dsql
 import aws.sdk.kotlin.runtime.auth.AuthTokenGenerator
 import aws.sdk.kotlin.runtime.auth.credentials.DefaultChainCredentialsProvider
 import aws.smithy.kotlin.runtime.auth.awscredentials.Credentials
+import aws.smithy.kotlin.runtime.auth.awscredentials.CredentialsProvider
+import aws.smithy.kotlin.runtime.auth.awssigning.AwsSigner
+import aws.smithy.kotlin.runtime.auth.awssigning.DefaultAwsSigner
 import aws.smithy.kotlin.runtime.net.url.Url
-import kotlinx.coroutines.runBlocking
+import aws.smithy.kotlin.runtime.time.Clock
 import kotlin.time.Duration
 import kotlin.time.Duration.Companion.seconds
 
 /**
  * Generates an IAM authentication token for use with DSQL databases
- * @param credentials The credentials to use when generating the auth token, defaults to resolving credentials from the [DefaultChainCredentialsProvider]
+ * @param credentialsProvider The [CredentialsProvider] which will provide credentials to use when generating the auth token, defaults to [DefaultChainCredentialsProvider]
+ * @param credentialsRefreshBuffer The amount of time before the resolved [Credentials] expire in which they are considered expired, defaults to 10 seconds.
+ * @param signer The [AwsSigner] implementation to use when creating the authentication token, defaults to [DefaultAwsSigner]
+ * @param clock The [Clock] implementation to use
  */
-public class AuthTokenGenerator(
-    public val credentials: Credentials? = runBlocking { DefaultChainCredentialsProvider().resolve() },
+public class DsqlAuthTokenGenerator(
+    public val credentialsProvider: CredentialsProvider = DefaultChainCredentialsProvider(),
+    public val credentialsRefreshBuffer: Duration = 10.seconds,
+    public val signer: AwsSigner = DefaultAwsSigner,
+    public val clock: Clock = Clock.System
 ) {
-    private val generator = AuthTokenGenerator("dsql", credentials)
+    private val generator = AuthTokenGenerator("dsql", credentialsProvider, credentialsRefreshBuffer, signer, clock)
 
     /**
      * Generates an auth token for the DbConnect action.

diff --git a/...n/services/dsql/AuthTokenGeneratorTest.kt → ...rvices/dsql/DsqlAuthTokenGeneratorTest.kt b/...n/services/dsql/AuthTokenGeneratorTest.kt → ...rvices/dsql/DsqlAuthTokenGeneratorTest.kt
@@ -4,22 +4,28 @@
  */
 package aws.sdk.kotlin.services.dsql
 
+import aws.sdk.kotlin.runtime.auth.credentials.StaticCredentialsProvider
 import aws.smithy.kotlin.runtime.auth.awscredentials.Credentials
 import aws.smithy.kotlin.runtime.net.Host
 import aws.smithy.kotlin.runtime.net.url.Url
+import aws.smithy.kotlin.runtime.time.Instant
+import aws.smithy.kotlin.runtime.time.ManualClock
 import kotlinx.coroutines.test.runTest
 import kotlin.test.Test
 import kotlin.test.assertContains
 import kotlin.test.assertFalse
 import kotlin.test.assertTrue
 import kotlin.time.Duration.Companion.seconds
 
-class AuthTokenGeneratorTest {
+class DsqlAuthTokenGeneratorTest {
     @Test
     fun testGenerateDbConnectAuthToken() = runTest {
+        val clock = ManualClock(Instant.fromEpochSeconds(1724716800))
+
         val credentials = Credentials("akid", "secret")
+        val credentialsProvider = StaticCredentialsProvider(credentials)
 
-        val token = AuthTokenGenerator(credentials)
+        val token = DsqlAuthTokenGenerator(credentialsProvider, clock = clock)
             .generateDbConnectAuthToken(
                 endpoint = Url { host = Host.parse("peccy.dsql.us-east-1.on.aws") },
                 region = "us-east-1",
@@ -28,11 +34,7 @@ class AuthTokenGeneratorTest {
 
         // Token should have a parameter Action=DbConnect
         assertContains(token, "peccy.dsql.us-east-1.on.aws?Action=DbConnect")
-
-        // Match the X-Amz-Credential parameter for any signing date
-        val credentialRegex = Regex("X-Amz-Credential=akid%2F(\\d{8})%2Fus-east-1%2Fdsql%2Faws4_request")
-        assertTrue(token.contains(credentialRegex))
-
+        assertContains(token, "X-Amz-Credential=akid%2F20240827%2Fus-east-1%2Fdsql%2Faws4_request")
         assertContains(token, "X-Amz-Expires=450")
 
         // Token should not contain a scheme
@@ -43,9 +45,12 @@ class AuthTokenGeneratorTest {
 
     @Test
     fun testGenerateDbConnectAuthAdminToken() = runTest {
+        val clock = ManualClock(Instant.fromEpochSeconds(1724716800))
+
         val credentials = Credentials("akid", "secret")
+        val credentialsProvider = StaticCredentialsProvider(credentials)
 
-        val token = AuthTokenGenerator(credentials)
+        val token = DsqlAuthTokenGenerator(credentialsProvider, clock = clock)
             .generateDbConnectAdminAuthToken(
                 endpoint = Url { host = Host.parse("peccy.dsql.us-east-1.on.aws") },
                 region = "us-east-1",
@@ -54,11 +59,7 @@ class AuthTokenGeneratorTest {
 
         // Token should have a parameter Action=DbConnectAdmin
         assertContains(token, "peccy.dsql.us-east-1.on.aws?Action=DbConnectAdmin")
-
-        // Match the X-Amz-Credential parameter for any signing date
-        val credentialRegex = Regex("X-Amz-Credential=akid%2F(\\d{8})%2Fus-east-1%2Fdsql%2Faws4_request")
-        assertTrue(token.contains(credentialRegex))
-
+        assertContains(token, "X-Amz-Credential=akid%2F20240827%2Fus-east-1%2Fdsql%2Faws4_request")
         assertContains(token, "X-Amz-Expires=450")
 
         // Token should not contain a scheme

diff --git a/...kotlin/services/rds/AuthTokenGenerator.kt → ...lin/services/rds/RdsAuthTokenGenerator.kt b/...kotlin/services/rds/AuthTokenGenerator.kt → ...lin/services/rds/RdsAuthTokenGenerator.kt
@@ -7,20 +7,30 @@ package aws.sdk.kotlin.services.rds
 import aws.sdk.kotlin.runtime.auth.AuthTokenGenerator
 import aws.sdk.kotlin.runtime.auth.credentials.DefaultChainCredentialsProvider
 import aws.smithy.kotlin.runtime.auth.awscredentials.Credentials
+import aws.smithy.kotlin.runtime.auth.awscredentials.CredentialsProvider
+import aws.smithy.kotlin.runtime.auth.awssigning.AwsSigner
+import aws.smithy.kotlin.runtime.auth.awssigning.DefaultAwsSigner
 import aws.smithy.kotlin.runtime.net.url.Url
+import aws.smithy.kotlin.runtime.time.Clock
 import kotlinx.coroutines.runBlocking
 import kotlin.apply
 import kotlin.time.Duration
 import kotlin.time.Duration.Companion.seconds
 
 /**
  * Generates an IAM authentication token for use with RDS databases
- * @param credentials The credentials to use when generating the auth token, defaults to resolving credentials from the [DefaultChainCredentialsProvider]
+ * @param credentialsProvider The [CredentialsProvider] which will provide credentials to use when generating the auth token, defaults to [DefaultChainCredentialsProvider]
+ * @param credentialsRefreshBuffer The amount of time before the resolved [Credentials] expire in which they are considered expired, defaults to 10 seconds.
+ * @param signer The [AwsSigner] implementation to use when creating the authentication token, defaults to [DefaultAwsSigner]
+ * @param clock The [Clock] implementation to use
  */
-public class AuthTokenGenerator(
-    public val credentials: Credentials? = runBlocking { DefaultChainCredentialsProvider().resolve() },
+public class RdsAuthTokenGenerator(
+    public val credentialsProvider: CredentialsProvider = DefaultChainCredentialsProvider(),
+    public val credentialsRefreshBuffer: Duration = 10.seconds,
+    public val signer: AwsSigner = DefaultAwsSigner,
+    public val clock: Clock = Clock.System
 ) {
-    private val generator = AuthTokenGenerator("rds-db", credentials)
+    private val generator = AuthTokenGenerator("rds-db", credentialsProvider, credentialsRefreshBuffer, signer, clock)
 
     /**
      * Generates an auth token for the `connect` action.

diff --git a/...rds/common/test/AuthTokenGeneratorTest.kt → .../common/test/RdsAuthTokenGeneratorTest.kt b/...rds/common/test/AuthTokenGeneratorTest.kt → .../common/test/RdsAuthTokenGeneratorTest.kt
@@ -4,39 +4,43 @@
  */
 package aws.sdk.kotlin.services.rds
 
+import aws.sdk.kotlin.runtime.auth.credentials.StaticCredentialsProvider
 import aws.smithy.kotlin.runtime.auth.awscredentials.Credentials
 import aws.smithy.kotlin.runtime.net.Host
 import aws.smithy.kotlin.runtime.net.url.Url
+import aws.smithy.kotlin.runtime.time.Instant
+import aws.smithy.kotlin.runtime.time.ManualClock
 import kotlinx.coroutines.test.runTest
 import kotlin.test.Test
 import kotlin.test.assertContains
 import kotlin.test.assertFalse
 import kotlin.test.assertTrue
 import kotlin.time.Duration.Companion.seconds
 
-class AuthTokenGeneratorTest {
+class RdsAuthTokenGeneratorTest {
     @Test
     fun testGenerateAuthToken() = runTest {
+        val clock = ManualClock(Instant.fromEpochSeconds(1724716800))
+        println(clock.now())
+
         val credentials = Credentials("akid", "secret")
+        val credentialsProvider = StaticCredentialsProvider(credentials)
+
+        val generator = RdsAuthTokenGenerator(credentialsProvider, clock = clock)
 
-        val token = AuthTokenGenerator(credentials)
-            .generateAuthToken(
-                endpoint = Url {
-                    host = Host.parse("prod-instance.us-east-1.rds.amazonaws.com")
-                    port = 3306
-                },
-                region = "us-east-1",
-                username = "peccy",
-                expiration = 450.seconds,
-            )
+        val token = generator.generateAuthToken(
+            endpoint = Url {
+                host = Host.parse("prod-instance.us-east-1.rds.amazonaws.com")
+                port = 3306
+            },
+            region = "us-east-1",
+            username = "peccy",
+            expiration = 450.seconds,
+        )
 
         // Token should have a parameter Action=connect, DBUser=peccy
         assertContains(token, "prod-instance.us-east-1.rds.amazonaws.com:3306?Action=connect&DBUser=peccy")
-
-        // Match the X-Amz-Credential parameter for any signing date
-        val credentialRegex = Regex("X-Amz-Credential=akid%2F(\\d{8})%2Fus-east-1%2Frds-db%2Faws4_request")
-        assertTrue(token.contains(credentialRegex))
-
+        assertContains(token, "X-Amz-Credential=akid%2F20240827%2Fus-east-1%2Frds-db%2Faws4_request")
         assertContains(token, "X-Amz-Expires=450")
 
         // Token should not contain a scheme