From b2624eab8891e77e4cb9666e5345777a1f061a55 Mon Sep 17 00:00:00 2001
From: Matas Lauzadis <lauzmata@amazon.com>
Date: Wed, 8 Jan 2025 11:22:09 -0500
Subject: [PATCH 01/24] Implement DSQL auth token generator

---
 .../services/dsql/AuthTokenGenerator.kt       | 81 +++++++++++++++++++
 .../services/dsql/AuthTokenGeneratorTest.kt   | 69 ++++++++++++++++
 2 files changed, 150 insertions(+)
 create mode 100644 services/dsql/common/src/aws/sdk/kotlin/services/dsql/AuthTokenGenerator.kt
 create mode 100644 services/dsql/common/test/aws/sdk/kotlin/services/dsql/AuthTokenGeneratorTest.kt

diff --git a/services/dsql/common/src/aws/sdk/kotlin/services/dsql/AuthTokenGenerator.kt b/services/dsql/common/src/aws/sdk/kotlin/services/dsql/AuthTokenGenerator.kt
new file mode 100644
index 00000000000..9e916f59804
--- /dev/null
+++ b/services/dsql/common/src/aws/sdk/kotlin/services/dsql/AuthTokenGenerator.kt
@@ -0,0 +1,81 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package aws.sdk.kotlin.services.dsql
+
+import aws.sdk.kotlin.runtime.auth.credentials.DefaultChainCredentialsProvider
+import aws.smithy.kotlin.runtime.auth.awscredentials.Credentials
+import aws.smithy.kotlin.runtime.auth.awssigning.AwsSignatureType
+import aws.smithy.kotlin.runtime.auth.awssigning.AwsSigningConfig
+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 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]
+ */
+public class AuthTokenGenerator(
+    public val credentials: Credentials? = runBlocking { DefaultChainCredentialsProvider().resolve() }
+) {
+    private fun String.trimScheme() = removePrefix("http://").removePrefix("https://")
+
+    /**
+     * Generates an auth token for the DbConnect action.
+     * @param endpoint the endpoint of the database
+     * @param region the region of the database
+     * @param expiration how long the auth token should be valid for. Defaults to 900.seconds
+     */
+    public suspend fun generateDbConnectAuthToken(endpoint: Url, region: String, expiration: Duration = 900.seconds): String {
+        val dbConnectEndpoint = endpoint.toBuilder().apply {
+            parameters.apply {
+                decodedParameters {
+                    add("Action", "DbConnect")
+                }
+            }
+        }.build()
+
+        return generateAuthToken(dbConnectEndpoint, region, expiration)
+    }
+
+    /**
+     * Generates an auth token for the DbConnectAdmin action.
+     * @param endpoint the endpoint of the database
+     * @param region the region of the database
+     * @param expiration how long the auth token should be valid for. Defaults to 900.seconds
+     */
+    public suspend fun generateDbConnectAdminAuthToken(endpoint: Url, region: String, expiration: Duration = 900.seconds): String {
+        val dbConnectAdminEndpoint = endpoint.toBuilder().apply {
+            parameters.apply {
+                decodedParameters {
+                    add("Action", "DbConnectAdmin")
+                }
+            }
+        }.build()
+
+        return generateAuthToken(dbConnectAdminEndpoint, region, expiration)
+    }
+
+    private suspend fun generateAuthToken(endpoint: Url, region: String, expiration: Duration): String {
+        val req = HttpRequest(HttpMethod.GET, endpoint)
+
+        val creds = credentials
+
+        val config = AwsSigningConfig {
+            credentials = creds ?: DefaultChainCredentialsProvider().resolve()
+            this.region = region
+            service = "dsql"
+            signingDate = Clock.System.now()
+            expiresAfter = expiration
+            signatureType = AwsSignatureType.HTTP_REQUEST_VIA_QUERY_PARAMS
+        }
+
+        return DefaultAwsSigner.sign(req, config).output.url.toString().trimScheme()
+    }
+}
\ No newline at end of file
diff --git a/services/dsql/common/test/aws/sdk/kotlin/services/dsql/AuthTokenGeneratorTest.kt b/services/dsql/common/test/aws/sdk/kotlin/services/dsql/AuthTokenGeneratorTest.kt
new file mode 100644
index 00000000000..f8028a6ae56
--- /dev/null
+++ b/services/dsql/common/test/aws/sdk/kotlin/services/dsql/AuthTokenGeneratorTest.kt
@@ -0,0 +1,69 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package aws.sdk.kotlin.services.dsql
+
+import aws.smithy.kotlin.runtime.auth.awscredentials.Credentials
+import aws.smithy.kotlin.runtime.net.Host
+import aws.smithy.kotlin.runtime.net.url.Url
+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 {
+    @Test
+    fun testGenerateDbConnectAuthToken() = runTest {
+        val credentials = Credentials("akid", "secret")
+
+        val token = AuthTokenGenerator(credentials)
+            .generateDbConnectAuthToken(
+                endpoint = Url { host = Host.parse("peccy.dsql.us-east-1.on.aws") },
+                region = "us-east-1",
+                expiration = 450.seconds
+            )
+
+        // 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-Expires=450")
+
+        // Token should not contain a scheme
+        listOf("http://", "https://").forEach {
+            assertFalse(token.contains(it))
+        }
+    }
+
+    @Test
+    fun testGenerateDbConnectAuthAdminToken() = runTest {
+        val credentials = Credentials("akid", "secret")
+
+        val token = AuthTokenGenerator(credentials)
+            .generateDbConnectAdminAuthToken(
+                endpoint = Url { host = Host.parse("peccy.dsql.us-east-1.on.aws") },
+                region = "us-east-1",
+                expiration = 450.seconds
+            )
+
+        // Token should have a parameter Action=DbConnect
+        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-Expires=450")
+
+        // Token should not contain a scheme
+        listOf("http://", "https://").forEach {
+            assertFalse(token.contains(it))
+        }
+    }
+}
\ No newline at end of file

From adb9d72c3a868bd940c20ef889d67b5753fee698 Mon Sep 17 00:00:00 2001
From: Matas Lauzadis <lauzmata@amazon.com>
Date: Wed, 8 Jan 2025 12:27:44 -0500
Subject: [PATCH 02/24] Abstract AuthTokenGenerator and implement RDS
 AuthTokenGenerator

---
 .../services/dsql/AuthTokenGenerator.kt       | 30 ++----------
 .../kotlin/services/rds/AuthTokenGenerator.kt | 43 +++++++++++++++++
 .../rds/common/test/AuthTokenGeneratorTest.kt | 47 +++++++++++++++++++
 3 files changed, 94 insertions(+), 26 deletions(-)
 create mode 100644 services/rds/common/src/aws/sdk/kotlin/services/rds/AuthTokenGenerator.kt
 create mode 100644 services/rds/common/test/AuthTokenGeneratorTest.kt

diff --git a/services/dsql/common/src/aws/sdk/kotlin/services/dsql/AuthTokenGenerator.kt b/services/dsql/common/src/aws/sdk/kotlin/services/dsql/AuthTokenGenerator.kt
index 9e916f59804..151ee1cbf36 100644
--- a/services/dsql/common/src/aws/sdk/kotlin/services/dsql/AuthTokenGenerator.kt
+++ b/services/dsql/common/src/aws/sdk/kotlin/services/dsql/AuthTokenGenerator.kt
@@ -6,16 +6,11 @@ package aws.sdk.kotlin.services.dsql
 
 import aws.sdk.kotlin.runtime.auth.credentials.DefaultChainCredentialsProvider
 import aws.smithy.kotlin.runtime.auth.awscredentials.Credentials
-import aws.smithy.kotlin.runtime.auth.awssigning.AwsSignatureType
-import aws.smithy.kotlin.runtime.auth.awssigning.AwsSigningConfig
-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 kotlin.time.Duration
 import kotlin.time.Duration.Companion.seconds
+import aws.sdk.kotlin.runtime.auth.AuthTokenGenerator
 
 /**
  * Generates an IAM authentication token for use with DSQL databases
@@ -24,7 +19,7 @@ import kotlin.time.Duration.Companion.seconds
 public class AuthTokenGenerator(
     public val credentials: Credentials? = runBlocking { DefaultChainCredentialsProvider().resolve() }
 ) {
-    private fun String.trimScheme() = removePrefix("http://").removePrefix("https://")
+    private val generator = AuthTokenGenerator("dsql", credentials)
 
     /**
      * Generates an auth token for the DbConnect action.
@@ -41,7 +36,7 @@ public class AuthTokenGenerator(
             }
         }.build()
 
-        return generateAuthToken(dbConnectEndpoint, region, expiration)
+        return generator.generateAuthToken(dbConnectEndpoint, region, expiration)
     }
 
     /**
@@ -59,23 +54,6 @@ public class AuthTokenGenerator(
             }
         }.build()
 
-        return generateAuthToken(dbConnectAdminEndpoint, region, expiration)
-    }
-
-    private suspend fun generateAuthToken(endpoint: Url, region: String, expiration: Duration): String {
-        val req = HttpRequest(HttpMethod.GET, endpoint)
-
-        val creds = credentials
-
-        val config = AwsSigningConfig {
-            credentials = creds ?: DefaultChainCredentialsProvider().resolve()
-            this.region = region
-            service = "dsql"
-            signingDate = Clock.System.now()
-            expiresAfter = expiration
-            signatureType = AwsSignatureType.HTTP_REQUEST_VIA_QUERY_PARAMS
-        }
-
-        return DefaultAwsSigner.sign(req, config).output.url.toString().trimScheme()
+        return generator.generateAuthToken(dbConnectAdminEndpoint, region, expiration)
     }
 }
\ No newline at end of file
diff --git a/services/rds/common/src/aws/sdk/kotlin/services/rds/AuthTokenGenerator.kt b/services/rds/common/src/aws/sdk/kotlin/services/rds/AuthTokenGenerator.kt
new file mode 100644
index 00000000000..3f567a6c957
--- /dev/null
+++ b/services/rds/common/src/aws/sdk/kotlin/services/rds/AuthTokenGenerator.kt
@@ -0,0 +1,43 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+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.net.url.Url
+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]
+ */
+public class AuthTokenGenerator(
+    public val credentials: Credentials? = runBlocking { DefaultChainCredentialsProvider().resolve() }
+) {
+    private val generator = AuthTokenGenerator("rds-db", credentials)
+
+    /**
+     * Generates an auth token for the DbConnect action.
+     * @param endpoint the endpoint of the database
+     * @param region the region of the database
+     * @param expiration how long the auth token should be valid for. Defaults to 900.seconds
+     */
+    public suspend fun generateAuthToken(endpoint: Url, region: String, username: String, expiration: Duration = 900.seconds): String {
+        val dbConnectEndpoint = endpoint.toBuilder().apply {
+            parameters.apply {
+                decodedParameters {
+                    add("Action", "connect")
+                    add("DBUser", username)
+                }
+            }
+        }.build()
+
+        return generator.generateAuthToken(dbConnectEndpoint, region, expiration)
+    }
+}
\ No newline at end of file
diff --git a/services/rds/common/test/AuthTokenGeneratorTest.kt b/services/rds/common/test/AuthTokenGeneratorTest.kt
new file mode 100644
index 00000000000..066cf3debb7
--- /dev/null
+++ b/services/rds/common/test/AuthTokenGeneratorTest.kt
@@ -0,0 +1,47 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package aws.sdk.kotlin.services.rds
+
+import aws.smithy.kotlin.runtime.auth.awscredentials.Credentials
+import aws.smithy.kotlin.runtime.net.Host
+import aws.smithy.kotlin.runtime.net.url.Url
+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 {
+    @Test
+    fun testGenerateDbConnectAuthToken() = runTest {
+        val credentials = Credentials("akid", "secret")
+
+        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
+            )
+
+        // Token should have a parameter Action=DbConnect
+        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-Expires=450")
+
+        // Token should not contain a scheme
+        listOf("http://", "https://").forEach {
+            assertFalse(token.contains(it))
+        }
+    }
+}
\ No newline at end of file

From 2631b3a9121a6610d892953f0536b6d439fb0dc5 Mon Sep 17 00:00:00 2001
From: Matas Lauzadis <lauzmata@amazon.com>
Date: Wed, 8 Jan 2025 12:28:08 -0500
Subject: [PATCH 03/24] ktlint

---
 .../src/aws/sdk/kotlin/services/dsql/AuthTokenGenerator.kt  | 6 +++---
 .../aws/sdk/kotlin/services/dsql/AuthTokenGeneratorTest.kt  | 6 +++---
 .../src/aws/sdk/kotlin/services/rds/AuthTokenGenerator.kt   | 4 ++--
 services/rds/common/test/AuthTokenGeneratorTest.kt          | 4 ++--
 4 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/services/dsql/common/src/aws/sdk/kotlin/services/dsql/AuthTokenGenerator.kt b/services/dsql/common/src/aws/sdk/kotlin/services/dsql/AuthTokenGenerator.kt
index 151ee1cbf36..e489619d74e 100644
--- a/services/dsql/common/src/aws/sdk/kotlin/services/dsql/AuthTokenGenerator.kt
+++ b/services/dsql/common/src/aws/sdk/kotlin/services/dsql/AuthTokenGenerator.kt
@@ -4,20 +4,20 @@
  */
 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.net.url.Url
 import kotlinx.coroutines.runBlocking
 import kotlin.time.Duration
 import kotlin.time.Duration.Companion.seconds
-import aws.sdk.kotlin.runtime.auth.AuthTokenGenerator
 
 /**
  * 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]
  */
 public class AuthTokenGenerator(
-    public val credentials: Credentials? = runBlocking { DefaultChainCredentialsProvider().resolve() }
+    public val credentials: Credentials? = runBlocking { DefaultChainCredentialsProvider().resolve() },
 ) {
     private val generator = AuthTokenGenerator("dsql", credentials)
 
@@ -56,4 +56,4 @@ public class AuthTokenGenerator(
 
         return generator.generateAuthToken(dbConnectAdminEndpoint, region, expiration)
     }
-}
\ No newline at end of file
+}
diff --git a/services/dsql/common/test/aws/sdk/kotlin/services/dsql/AuthTokenGeneratorTest.kt b/services/dsql/common/test/aws/sdk/kotlin/services/dsql/AuthTokenGeneratorTest.kt
index f8028a6ae56..3c44c7b5c0d 100644
--- a/services/dsql/common/test/aws/sdk/kotlin/services/dsql/AuthTokenGeneratorTest.kt
+++ b/services/dsql/common/test/aws/sdk/kotlin/services/dsql/AuthTokenGeneratorTest.kt
@@ -23,7 +23,7 @@ class AuthTokenGeneratorTest {
             .generateDbConnectAuthToken(
                 endpoint = Url { host = Host.parse("peccy.dsql.us-east-1.on.aws") },
                 region = "us-east-1",
-                expiration = 450.seconds
+                expiration = 450.seconds,
             )
 
         // Token should have a parameter Action=DbConnect
@@ -49,7 +49,7 @@ class AuthTokenGeneratorTest {
             .generateDbConnectAdminAuthToken(
                 endpoint = Url { host = Host.parse("peccy.dsql.us-east-1.on.aws") },
                 region = "us-east-1",
-                expiration = 450.seconds
+                expiration = 450.seconds,
             )
 
         // Token should have a parameter Action=DbConnect
@@ -66,4 +66,4 @@ class AuthTokenGeneratorTest {
             assertFalse(token.contains(it))
         }
     }
-}
\ No newline at end of file
+}
diff --git a/services/rds/common/src/aws/sdk/kotlin/services/rds/AuthTokenGenerator.kt b/services/rds/common/src/aws/sdk/kotlin/services/rds/AuthTokenGenerator.kt
index 3f567a6c957..3af79db0039 100644
--- a/services/rds/common/src/aws/sdk/kotlin/services/rds/AuthTokenGenerator.kt
+++ b/services/rds/common/src/aws/sdk/kotlin/services/rds/AuthTokenGenerator.kt
@@ -18,7 +18,7 @@ import kotlin.time.Duration.Companion.seconds
  * @param credentials The credentials to use when generating the auth token, defaults to resolving credentials from the [DefaultChainCredentialsProvider]
  */
 public class AuthTokenGenerator(
-    public val credentials: Credentials? = runBlocking { DefaultChainCredentialsProvider().resolve() }
+    public val credentials: Credentials? = runBlocking { DefaultChainCredentialsProvider().resolve() },
 ) {
     private val generator = AuthTokenGenerator("rds-db", credentials)
 
@@ -40,4 +40,4 @@ public class AuthTokenGenerator(
 
         return generator.generateAuthToken(dbConnectEndpoint, region, expiration)
     }
-}
\ No newline at end of file
+}
diff --git a/services/rds/common/test/AuthTokenGeneratorTest.kt b/services/rds/common/test/AuthTokenGeneratorTest.kt
index 066cf3debb7..d1f37694ff5 100644
--- a/services/rds/common/test/AuthTokenGeneratorTest.kt
+++ b/services/rds/common/test/AuthTokenGeneratorTest.kt
@@ -27,7 +27,7 @@ class AuthTokenGeneratorTest {
                 },
                 region = "us-east-1",
                 username = "peccy",
-                expiration = 450.seconds
+                expiration = 450.seconds,
             )
 
         // Token should have a parameter Action=DbConnect
@@ -44,4 +44,4 @@ class AuthTokenGeneratorTest {
             assertFalse(token.contains(it))
         }
     }
-}
\ No newline at end of file
+}

From e7db5b60fd2230dfc59141cfa24ba95ad30c3667 Mon Sep 17 00:00:00 2001
From: Matas Lauzadis <lauzmata@amazon.com>
Date: Wed, 8 Jan 2025 12:28:46 -0500
Subject: [PATCH 04/24] changelog

---
 .changes/a2520ae2-1cba-49b1-b720-10b70e52f9e0.json | 5 +++++
 1 file changed, 5 insertions(+)
 create mode 100644 .changes/a2520ae2-1cba-49b1-b720-10b70e52f9e0.json

diff --git a/.changes/a2520ae2-1cba-49b1-b720-10b70e52f9e0.json b/.changes/a2520ae2-1cba-49b1-b720-10b70e52f9e0.json
new file mode 100644
index 00000000000..9d81f0fdc20
--- /dev/null
+++ b/.changes/a2520ae2-1cba-49b1-b720-10b70e52f9e0.json
@@ -0,0 +1,5 @@
+{
+    "id": "a2520ae2-1cba-49b1-b720-10b70e52f9e0",
+    "type": "feature",
+    "description": "Add auth token generator for RDS and DSQL"
+}
\ No newline at end of file

From d46a242b6b5f630251a7a443789a6bc21a8ed0ba Mon Sep 17 00:00:00 2001
From: Matas Lauzadis <lauzmata@amazon.com>
Date: Wed, 8 Jan 2025 12:29:54 -0500
Subject: [PATCH 05/24] Update test comments

---
 services/rds/common/test/AuthTokenGeneratorTest.kt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/services/rds/common/test/AuthTokenGeneratorTest.kt b/services/rds/common/test/AuthTokenGeneratorTest.kt
index d1f37694ff5..207c826efbf 100644
--- a/services/rds/common/test/AuthTokenGeneratorTest.kt
+++ b/services/rds/common/test/AuthTokenGeneratorTest.kt
@@ -30,7 +30,7 @@ class AuthTokenGeneratorTest {
                 expiration = 450.seconds,
             )
 
-        // Token should have a parameter Action=DbConnect
+        // 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

From 6a12913940c0fb821ce63985b8af118800e66c11 Mon Sep 17 00:00:00 2001
From: Matas Lauzadis <lauzmata@amazon.com>
Date: Wed, 8 Jan 2025 12:31:13 -0500
Subject: [PATCH 06/24] Update test comments

---
 .../src/aws/sdk/kotlin/services/rds/AuthTokenGenerator.kt  | 7 ++++---
 services/rds/common/test/AuthTokenGeneratorTest.kt         | 2 +-
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/services/rds/common/src/aws/sdk/kotlin/services/rds/AuthTokenGenerator.kt b/services/rds/common/src/aws/sdk/kotlin/services/rds/AuthTokenGenerator.kt
index 3af79db0039..06786e5749b 100644
--- a/services/rds/common/src/aws/sdk/kotlin/services/rds/AuthTokenGenerator.kt
+++ b/services/rds/common/src/aws/sdk/kotlin/services/rds/AuthTokenGenerator.kt
@@ -23,13 +23,14 @@ public class AuthTokenGenerator(
     private val generator = AuthTokenGenerator("rds-db", credentials)
 
     /**
-     * Generates an auth token for the DbConnect action.
+     * Generates an auth token for the `connect` action.
      * @param endpoint the endpoint of the database
      * @param region the region of the database
+     * @param username the username to authenticate with
      * @param expiration how long the auth token should be valid for. Defaults to 900.seconds
      */
     public suspend fun generateAuthToken(endpoint: Url, region: String, username: String, expiration: Duration = 900.seconds): String {
-        val dbConnectEndpoint = endpoint.toBuilder().apply {
+        val endpoint = endpoint.toBuilder().apply {
             parameters.apply {
                 decodedParameters {
                     add("Action", "connect")
@@ -38,6 +39,6 @@ public class AuthTokenGenerator(
             }
         }.build()
 
-        return generator.generateAuthToken(dbConnectEndpoint, region, expiration)
+        return generator.generateAuthToken(endpoint, region, expiration)
     }
 }
diff --git a/services/rds/common/test/AuthTokenGeneratorTest.kt b/services/rds/common/test/AuthTokenGeneratorTest.kt
index 207c826efbf..727364ea612 100644
--- a/services/rds/common/test/AuthTokenGeneratorTest.kt
+++ b/services/rds/common/test/AuthTokenGeneratorTest.kt
@@ -16,7 +16,7 @@ import kotlin.time.Duration.Companion.seconds
 
 class AuthTokenGeneratorTest {
     @Test
-    fun testGenerateDbConnectAuthToken() = runTest {
+    fun testGenerateAuthToken() = runTest {
         val credentials = Credentials("akid", "secret")
 
         val token = AuthTokenGenerator(credentials)

From bbc53a555297bc0abea1ff377149c88311d85abf Mon Sep 17 00:00:00 2001
From: Matas Lauzadis <lauzmata@amazon.com>
Date: Wed, 8 Jan 2025 12:31:57 -0500
Subject: [PATCH 07/24] Update test comments

---
 .../test/aws/sdk/kotlin/services/dsql/AuthTokenGeneratorTest.kt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/services/dsql/common/test/aws/sdk/kotlin/services/dsql/AuthTokenGeneratorTest.kt b/services/dsql/common/test/aws/sdk/kotlin/services/dsql/AuthTokenGeneratorTest.kt
index 3c44c7b5c0d..8f21a277370 100644
--- a/services/dsql/common/test/aws/sdk/kotlin/services/dsql/AuthTokenGeneratorTest.kt
+++ b/services/dsql/common/test/aws/sdk/kotlin/services/dsql/AuthTokenGeneratorTest.kt
@@ -52,7 +52,7 @@ class AuthTokenGeneratorTest {
                 expiration = 450.seconds,
             )
 
-        // Token should have a parameter Action=DbConnect
+        // 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

From 2c665000836b35f61a82aa0d923b3b045e42f436 Mon Sep 17 00:00:00 2001
From: Matas Lauzadis <lauzmata@amazon.com>
Date: Wed, 8 Jan 2025 12:33:03 -0500
Subject: [PATCH 08/24] Commit abstracted AuthTokenGenerator

---
 .../kotlin/runtime/auth/AuthTokenGenerator.kt | 48 +++++++++++++++++++
 1 file changed, 48 insertions(+)
 create mode 100644 aws-runtime/aws-config/common/src/aws/sdk/kotlin/runtime/auth/AuthTokenGenerator.kt

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
new file mode 100644
index 00000000000..ca45fcd5538
--- /dev/null
+++ b/aws-runtime/aws-config/common/src/aws/sdk/kotlin/runtime/auth/AuthTokenGenerator.kt
@@ -0,0 +1,48 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+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.awssigning.AwsSignatureType
+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 kotlin.time.Duration
+
+/**
+ * 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]
+ */
+public class AuthTokenGenerator(
+    public val service: String,
+    public val credentials: Credentials? = runBlocking { DefaultChainCredentialsProvider().resolve() },
+) {
+    private fun String.trimScheme() = removePrefix("http://").removePrefix("https://")
+
+    public suspend fun generateAuthToken(endpoint: Url, region: String, expiration: Duration): String {
+        val req = HttpRequest(HttpMethod.GET, endpoint)
+
+        val creds = credentials
+        val serv = service
+
+        val config = AwsSigningConfig {
+            credentials = creds
+            this.region = region
+            service = serv
+            signingDate = Clock.System.now()
+            expiresAfter = expiration
+            signatureType = AwsSignatureType.HTTP_REQUEST_VIA_QUERY_PARAMS
+        }
+
+        return DefaultAwsSigner.sign(req, config).output.url.toString().trimScheme()
+    }
+}

From 68fa5639770c2c601c3f9e35baf9adde123ef11f Mon Sep 17 00:00:00 2001
From: Matas Lauzadis <lauzmata@amazon.com>
Date: Wed, 8 Jan 2025 12:39:11 -0500
Subject: [PATCH 09/24] Add API dump

---
 aws-runtime/aws-config/api/aws-config.api | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/aws-runtime/aws-config/api/aws-config.api b/aws-runtime/aws-config/api/aws-config.api
index 05589280306..7cc13dfd5c5 100644
--- a/aws-runtime/aws-config/api/aws-config.api
+++ b/aws-runtime/aws-config/api/aws-config.api
@@ -1,3 +1,11 @@
+public final class aws/sdk/kotlin/runtime/auth/AuthTokenGenerator {
+	public fun <init> (Ljava/lang/String;Laws/smithy/kotlin/runtime/auth/awscredentials/Credentials;)V
+	public synthetic fun <init> (Ljava/lang/String;Laws/smithy/kotlin/runtime/auth/awscredentials/Credentials;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
+	public final fun generateAuthToken-exY8QGI (Laws/smithy/kotlin/runtime/net/url/Url;Ljava/lang/String;JLkotlin/coroutines/Continuation;)Ljava/lang/Object;
+	public final fun getCredentials ()Laws/smithy/kotlin/runtime/auth/awscredentials/Credentials;
+	public final fun getService ()Ljava/lang/String;
+}
+
 public final class aws/sdk/kotlin/runtime/auth/credentials/AssumeRoleParameters {
 	public synthetic fun <init> (Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;JLjava/util/List;Ljava/lang/String;Ljava/util/Map;Ljava/util/List;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
 	public synthetic fun <init> (Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;JLjava/util/List;Ljava/lang/String;Ljava/util/Map;Ljava/util/List;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Lkotlin/jvm/internal/DefaultConstructorMarker;)V

From 3b9b2b211ce298b906ed0fe49d50b9b9eb6140b7 Mon Sep 17 00:00:00 2001
From: Matas Lauzadis <lauzmata@amazon.com>
Date: Thu, 9 Jan 2025 11:05:34 -0500
Subject: [PATCH 10/24] Rename service-specific token generators, add more
 optional parameters, change [Credentials] param to [CredentialsProvider]

---
 .../kotlin/runtime/auth/AuthTokenGenerator.kt | 34 ++++++++++++++----
 ...Generator.kt => DsqlAuthTokenGenerator.kt} | 19 +++++++---
 ...rTest.kt => DsqlAuthTokenGeneratorTest.kt} | 27 +++++++-------
 ...nGenerator.kt => RdsAuthTokenGenerator.kt} | 18 +++++++---
 ...orTest.kt => RdsAuthTokenGeneratorTest.kt} | 36 ++++++++++---------
 5 files changed, 89 insertions(+), 45 deletions(-)
 rename services/dsql/common/src/aws/sdk/kotlin/services/dsql/{AuthTokenGenerator.kt => DsqlAuthTokenGenerator.kt} (65%)
 rename services/dsql/common/test/aws/sdk/kotlin/services/dsql/{AuthTokenGeneratorTest.kt => DsqlAuthTokenGeneratorTest.kt} (69%)
 rename services/rds/common/src/aws/sdk/kotlin/services/rds/{AuthTokenGenerator.kt => RdsAuthTokenGenerator.kt} (57%)
 rename services/rds/common/test/{AuthTokenGeneratorTest.kt => RdsAuthTokenGeneratorTest.kt} (53%)

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
index ca45fcd5538..a54d9172789 100644
--- 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,7 +6,9 @@ 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
@@ -14,35 +16,53 @@ 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/services/dsql/common/src/aws/sdk/kotlin/services/dsql/AuthTokenGenerator.kt b/services/dsql/common/src/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGenerator.kt
similarity index 65%
rename from services/dsql/common/src/aws/sdk/kotlin/services/dsql/AuthTokenGenerator.kt
rename to services/dsql/common/src/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGenerator.kt
index e489619d74e..231f26c1395 100644
--- a/services/dsql/common/src/aws/sdk/kotlin/services/dsql/AuthTokenGenerator.kt
+++ b/services/dsql/common/src/aws/sdk/kotlin/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/services/dsql/common/test/aws/sdk/kotlin/services/dsql/AuthTokenGeneratorTest.kt b/services/dsql/common/test/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGeneratorTest.kt
similarity index 69%
rename from services/dsql/common/test/aws/sdk/kotlin/services/dsql/AuthTokenGeneratorTest.kt
rename to services/dsql/common/test/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGeneratorTest.kt
index 8f21a277370..50757464ef5 100644
--- a/services/dsql/common/test/aws/sdk/kotlin/services/dsql/AuthTokenGeneratorTest.kt
+++ b/services/dsql/common/test/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGeneratorTest.kt
@@ -4,9 +4,12 @@
  */
 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
@@ -14,12 +17,15 @@ 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/services/rds/common/src/aws/sdk/kotlin/services/rds/AuthTokenGenerator.kt b/services/rds/common/src/aws/sdk/kotlin/services/rds/RdsAuthTokenGenerator.kt
similarity index 57%
rename from services/rds/common/src/aws/sdk/kotlin/services/rds/AuthTokenGenerator.kt
rename to services/rds/common/src/aws/sdk/kotlin/services/rds/RdsAuthTokenGenerator.kt
index 06786e5749b..1911842a925 100644
--- a/services/rds/common/src/aws/sdk/kotlin/services/rds/AuthTokenGenerator.kt
+++ b/services/rds/common/src/aws/sdk/kotlin/services/rds/RdsAuthTokenGenerator.kt
@@ -7,7 +7,11 @@ 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
@@ -15,12 +19,18 @@ 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/services/rds/common/test/AuthTokenGeneratorTest.kt b/services/rds/common/test/RdsAuthTokenGeneratorTest.kt
similarity index 53%
rename from services/rds/common/test/AuthTokenGeneratorTest.kt
rename to services/rds/common/test/RdsAuthTokenGeneratorTest.kt
index 727364ea612..cb3e5087979 100644
--- a/services/rds/common/test/AuthTokenGeneratorTest.kt
+++ b/services/rds/common/test/RdsAuthTokenGeneratorTest.kt
@@ -4,9 +4,12 @@
  */
 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
@@ -14,29 +17,30 @@ 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

From 19466be6ce782a1675da3d494c4a74483058ce02 Mon Sep 17 00:00:00 2001
From: Matas Lauzadis <lauzmata@amazon.com>
Date: Thu, 9 Jan 2025 11:06:14 -0500
Subject: [PATCH 11/24] ktlint

---
 .../src/aws/sdk/kotlin/runtime/auth/AuthTokenGenerator.kt      | 2 +-
 .../src/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGenerator.kt | 2 +-
 .../aws/sdk/kotlin/services/dsql/DsqlAuthTokenGeneratorTest.kt | 1 -
 .../src/aws/sdk/kotlin/services/rds/RdsAuthTokenGenerator.kt   | 3 +--
 services/rds/common/test/RdsAuthTokenGeneratorTest.kt          | 1 -
 5 files changed, 3 insertions(+), 6 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
index a54d9172789..38955755a10 100644
--- 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
@@ -37,7 +37,7 @@ public class AuthTokenGenerator(
     public val credentialsProvider: CredentialsProvider = DefaultChainCredentialsProvider(),
     public val credentialsRefreshBuffer: Duration = 10.seconds,
     public val signer: AwsSigner = DefaultAwsSigner,
-    public val clock: Clock = Clock.System
+    public val clock: Clock = Clock.System,
 ) {
     private lateinit var credentials: ExpiringValue<Credentials>
 
diff --git a/services/dsql/common/src/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGenerator.kt b/services/dsql/common/src/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGenerator.kt
index 231f26c1395..afc4b7fd76b 100644
--- a/services/dsql/common/src/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGenerator.kt
+++ b/services/dsql/common/src/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGenerator.kt
@@ -26,7 +26,7 @@ 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
+    public val clock: Clock = Clock.System,
 ) {
     private val generator = AuthTokenGenerator("dsql", credentialsProvider, credentialsRefreshBuffer, signer, clock)
 
diff --git a/services/dsql/common/test/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGeneratorTest.kt b/services/dsql/common/test/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGeneratorTest.kt
index 50757464ef5..069bc8d6241 100644
--- a/services/dsql/common/test/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGeneratorTest.kt
+++ b/services/dsql/common/test/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGeneratorTest.kt
@@ -14,7 +14,6 @@ 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 DsqlAuthTokenGeneratorTest {
diff --git a/services/rds/common/src/aws/sdk/kotlin/services/rds/RdsAuthTokenGenerator.kt b/services/rds/common/src/aws/sdk/kotlin/services/rds/RdsAuthTokenGenerator.kt
index 1911842a925..88255047f0d 100644
--- a/services/rds/common/src/aws/sdk/kotlin/services/rds/RdsAuthTokenGenerator.kt
+++ b/services/rds/common/src/aws/sdk/kotlin/services/rds/RdsAuthTokenGenerator.kt
@@ -12,7 +12,6 @@ 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
@@ -28,7 +27,7 @@ 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
+    public val clock: Clock = Clock.System,
 ) {
     private val generator = AuthTokenGenerator("rds-db", credentialsProvider, credentialsRefreshBuffer, signer, clock)
 
diff --git a/services/rds/common/test/RdsAuthTokenGeneratorTest.kt b/services/rds/common/test/RdsAuthTokenGeneratorTest.kt
index cb3e5087979..917eb504c93 100644
--- a/services/rds/common/test/RdsAuthTokenGeneratorTest.kt
+++ b/services/rds/common/test/RdsAuthTokenGeneratorTest.kt
@@ -14,7 +14,6 @@ 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 RdsAuthTokenGeneratorTest {

From d3f877de0463bb6c1379f90afacc2446a89f1ba4 Mon Sep 17 00:00:00 2001
From: Matas Lauzadis <lauzmata@amazon.com>
Date: Thu, 9 Jan 2025 11:06:46 -0500
Subject: [PATCH 12/24] apiDump

---
 aws-runtime/aws-config/api/aws-config.api | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/aws-runtime/aws-config/api/aws-config.api b/aws-runtime/aws-config/api/aws-config.api
index 7cc13dfd5c5..d0eae382f10 100644
--- a/aws-runtime/aws-config/api/aws-config.api
+++ b/aws-runtime/aws-config/api/aws-config.api
@@ -1,9 +1,12 @@
 public final class aws/sdk/kotlin/runtime/auth/AuthTokenGenerator {
-	public fun <init> (Ljava/lang/String;Laws/smithy/kotlin/runtime/auth/awscredentials/Credentials;)V
-	public synthetic fun <init> (Ljava/lang/String;Laws/smithy/kotlin/runtime/auth/awscredentials/Credentials;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
+	public synthetic fun <init> (Ljava/lang/String;Laws/smithy/kotlin/runtime/auth/awscredentials/CredentialsProvider;JLaws/smithy/kotlin/runtime/auth/awssigning/AwsSigner;Laws/smithy/kotlin/runtime/time/Clock;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
+	public synthetic fun <init> (Ljava/lang/String;Laws/smithy/kotlin/runtime/auth/awscredentials/CredentialsProvider;JLaws/smithy/kotlin/runtime/auth/awssigning/AwsSigner;Laws/smithy/kotlin/runtime/time/Clock;Lkotlin/jvm/internal/DefaultConstructorMarker;)V
 	public final fun generateAuthToken-exY8QGI (Laws/smithy/kotlin/runtime/net/url/Url;Ljava/lang/String;JLkotlin/coroutines/Continuation;)Ljava/lang/Object;
-	public final fun getCredentials ()Laws/smithy/kotlin/runtime/auth/awscredentials/Credentials;
+	public final fun getClock ()Laws/smithy/kotlin/runtime/time/Clock;
+	public final fun getCredentialsProvider ()Laws/smithy/kotlin/runtime/auth/awscredentials/CredentialsProvider;
+	public final fun getCredentialsRefreshBuffer-UwyO8pc ()J
 	public final fun getService ()Ljava/lang/String;
+	public final fun getSigner ()Laws/smithy/kotlin/runtime/auth/awssigning/AwsSigner;
 }
 
 public final class aws/sdk/kotlin/runtime/auth/credentials/AssumeRoleParameters {

From 0e88c5a297881e9c8384b2658a4b76242c269987 Mon Sep 17 00:00:00 2001
From: Matas Lauzadis <lauzmata@amazon.com>
Date: Thu, 9 Jan 2025 11:39:34 -0500
Subject: [PATCH 13/24] Relocate AuthTokenGenerator to aws-signing-common

---
 .../kotlin/runtime/auth/AuthTokenGenerator.kt | 68 -------------------
 .../services/dsql/DsqlAuthTokenGenerator.kt   |  2 +-
 .../services/rds/RdsAuthTokenGenerator.kt     |  2 +-
 3 files changed, 2 insertions(+), 70 deletions(-)
 delete mode 100644 aws-runtime/aws-config/common/src/aws/sdk/kotlin/runtime/auth/AuthTokenGenerator.kt

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
deleted file mode 100644
index 38955755a10..00000000000
--- a/aws-runtime/aws-config/common/src/aws/sdk/kotlin/runtime/auth/AuthTokenGenerator.kt
+++ /dev/null
@@ -1,68 +0,0 @@
-/*
- * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
- * SPDX-License-Identifier: Apache-2.0
- */
-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 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 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 credentialsProvider: CredentialsProvider = DefaultChainCredentialsProvider(),
-    public val credentialsRefreshBuffer: Duration = 10.seconds,
-    public val signer: AwsSigner = DefaultAwsSigner,
-    public val clock: Clock = Clock.System,
-) {
-    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.value
-        val serv = service
-
-        val config = AwsSigningConfig {
-            credentials = creds
-            this.region = region
-            service = serv
-            signingDate = clock.now()
-            expiresAfter = expiration
-            signatureType = AwsSignatureType.HTTP_REQUEST_VIA_QUERY_PARAMS
-        }
-
-        return signer.sign(req, config).output.url.trimScheme()
-    }
-}
diff --git a/services/dsql/common/src/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGenerator.kt b/services/dsql/common/src/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGenerator.kt
index afc4b7fd76b..9b8e10faee1 100644
--- a/services/dsql/common/src/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGenerator.kt
+++ b/services/dsql/common/src/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGenerator.kt
@@ -4,10 +4,10 @@
  */
 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.AuthTokenGenerator
 import aws.smithy.kotlin.runtime.auth.awssigning.AwsSigner
 import aws.smithy.kotlin.runtime.auth.awssigning.DefaultAwsSigner
 import aws.smithy.kotlin.runtime.net.url.Url
diff --git a/services/rds/common/src/aws/sdk/kotlin/services/rds/RdsAuthTokenGenerator.kt b/services/rds/common/src/aws/sdk/kotlin/services/rds/RdsAuthTokenGenerator.kt
index 88255047f0d..aae531f4578 100644
--- a/services/rds/common/src/aws/sdk/kotlin/services/rds/RdsAuthTokenGenerator.kt
+++ b/services/rds/common/src/aws/sdk/kotlin/services/rds/RdsAuthTokenGenerator.kt
@@ -4,10 +4,10 @@
  */
 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.AuthTokenGenerator
 import aws.smithy.kotlin.runtime.auth.awssigning.AwsSigner
 import aws.smithy.kotlin.runtime.auth.awssigning.DefaultAwsSigner
 import aws.smithy.kotlin.runtime.net.url.Url

From cc6961f29c83d4111b4baf630c9c7c6ae6150257 Mon Sep 17 00:00:00 2001
From: Matas Lauzadis <lauzmata@amazon.com>
Date: Thu, 9 Jan 2025 11:40:29 -0500
Subject: [PATCH 14/24] apiDump

---
 aws-runtime/aws-config/api/aws-config.api | 11 -----------
 1 file changed, 11 deletions(-)

diff --git a/aws-runtime/aws-config/api/aws-config.api b/aws-runtime/aws-config/api/aws-config.api
index d0eae382f10..05589280306 100644
--- a/aws-runtime/aws-config/api/aws-config.api
+++ b/aws-runtime/aws-config/api/aws-config.api
@@ -1,14 +1,3 @@
-public final class aws/sdk/kotlin/runtime/auth/AuthTokenGenerator {
-	public synthetic fun <init> (Ljava/lang/String;Laws/smithy/kotlin/runtime/auth/awscredentials/CredentialsProvider;JLaws/smithy/kotlin/runtime/auth/awssigning/AwsSigner;Laws/smithy/kotlin/runtime/time/Clock;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
-	public synthetic fun <init> (Ljava/lang/String;Laws/smithy/kotlin/runtime/auth/awscredentials/CredentialsProvider;JLaws/smithy/kotlin/runtime/auth/awssigning/AwsSigner;Laws/smithy/kotlin/runtime/time/Clock;Lkotlin/jvm/internal/DefaultConstructorMarker;)V
-	public final fun generateAuthToken-exY8QGI (Laws/smithy/kotlin/runtime/net/url/Url;Ljava/lang/String;JLkotlin/coroutines/Continuation;)Ljava/lang/Object;
-	public final fun getClock ()Laws/smithy/kotlin/runtime/time/Clock;
-	public final fun getCredentialsProvider ()Laws/smithy/kotlin/runtime/auth/awscredentials/CredentialsProvider;
-	public final fun getCredentialsRefreshBuffer-UwyO8pc ()J
-	public final fun getService ()Ljava/lang/String;
-	public final fun getSigner ()Laws/smithy/kotlin/runtime/auth/awssigning/AwsSigner;
-}
-
 public final class aws/sdk/kotlin/runtime/auth/credentials/AssumeRoleParameters {
 	public synthetic fun <init> (Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;JLjava/util/List;Ljava/lang/String;Ljava/util/Map;Ljava/util/List;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
 	public synthetic fun <init> (Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;JLjava/util/List;Ljava/lang/String;Ljava/util/Map;Ljava/util/List;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Lkotlin/jvm/internal/DefaultConstructorMarker;)V

From 1c391b80fd2f8f9af037b251fdf87d25a9514cb7 Mon Sep 17 00:00:00 2001
From: Matas Lauzadis <lauzmata@amazon.com>
Date: Thu, 9 Jan 2025 17:18:43 -0500
Subject: [PATCH 15/24] Refactor based on upstream changes, simplify parameter
 addition, change `add` to `put`

---
 .../kotlin/services/dsql/DsqlAuthTokenGenerator.kt | 14 +++-----------
 .../kotlin/services/rds/RdsAuthTokenGenerator.kt   |  2 +-
 .../rds/common/test/RdsAuthTokenGeneratorTest.kt   |  1 -
 3 files changed, 4 insertions(+), 13 deletions(-)

diff --git a/services/dsql/common/src/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGenerator.kt b/services/dsql/common/src/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGenerator.kt
index 9b8e10faee1..2bfef0e72c9 100644
--- a/services/dsql/common/src/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGenerator.kt
+++ b/services/dsql/common/src/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGenerator.kt
@@ -28,7 +28,7 @@ public class DsqlAuthTokenGenerator(
     public val signer: AwsSigner = DefaultAwsSigner,
     public val clock: Clock = Clock.System,
 ) {
-    private val generator = AuthTokenGenerator("dsql", credentialsProvider, credentialsRefreshBuffer, signer, clock)
+    private val generator = AuthTokenGenerator("dsql", credentialsProvider, signer, clock)
 
     /**
      * Generates an auth token for the DbConnect action.
@@ -38,11 +38,7 @@ public class DsqlAuthTokenGenerator(
      */
     public suspend fun generateDbConnectAuthToken(endpoint: Url, region: String, expiration: Duration = 900.seconds): String {
         val dbConnectEndpoint = endpoint.toBuilder().apply {
-            parameters.apply {
-                decodedParameters {
-                    add("Action", "DbConnect")
-                }
-            }
+            parameters.decodedParameters.put("Action", "DbConnect")
         }.build()
 
         return generator.generateAuthToken(dbConnectEndpoint, region, expiration)
@@ -56,11 +52,7 @@ public class DsqlAuthTokenGenerator(
      */
     public suspend fun generateDbConnectAdminAuthToken(endpoint: Url, region: String, expiration: Duration = 900.seconds): String {
         val dbConnectAdminEndpoint = endpoint.toBuilder().apply {
-            parameters.apply {
-                decodedParameters {
-                    add("Action", "DbConnectAdmin")
-                }
-            }
+            parameters.decodedParameters.put("Action", "DbConnectAdmin")
         }.build()
 
         return generator.generateAuthToken(dbConnectAdminEndpoint, region, expiration)
diff --git a/services/rds/common/src/aws/sdk/kotlin/services/rds/RdsAuthTokenGenerator.kt b/services/rds/common/src/aws/sdk/kotlin/services/rds/RdsAuthTokenGenerator.kt
index aae531f4578..3d19af00d9b 100644
--- a/services/rds/common/src/aws/sdk/kotlin/services/rds/RdsAuthTokenGenerator.kt
+++ b/services/rds/common/src/aws/sdk/kotlin/services/rds/RdsAuthTokenGenerator.kt
@@ -29,7 +29,7 @@ public class RdsAuthTokenGenerator(
     public val signer: AwsSigner = DefaultAwsSigner,
     public val clock: Clock = Clock.System,
 ) {
-    private val generator = AuthTokenGenerator("rds-db", credentialsProvider, credentialsRefreshBuffer, signer, clock)
+    private val generator = AuthTokenGenerator("rds-db", credentialsProvider, signer, clock)
 
     /**
      * Generates an auth token for the `connect` action.
diff --git a/services/rds/common/test/RdsAuthTokenGeneratorTest.kt b/services/rds/common/test/RdsAuthTokenGeneratorTest.kt
index 917eb504c93..5a1f76ad37c 100644
--- a/services/rds/common/test/RdsAuthTokenGeneratorTest.kt
+++ b/services/rds/common/test/RdsAuthTokenGeneratorTest.kt
@@ -20,7 +20,6 @@ class RdsAuthTokenGeneratorTest {
     @Test
     fun testGenerateAuthToken() = runTest {
         val clock = ManualClock(Instant.fromEpochSeconds(1724716800))
-        println(clock.now())
 
         val credentials = Credentials("akid", "secret")
         val credentialsProvider = StaticCredentialsProvider(credentials)

From 30d625eec5fde0d76d72e3390df7d9a1f9db2029 Mon Sep 17 00:00:00 2001
From: Matas Lauzadis <lauzmata@amazon.com>
Date: Thu, 9 Jan 2025 17:19:29 -0500
Subject: [PATCH 16/24] Remove `credentialsRefreshBuffer` parameter

---
 .../src/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGenerator.kt  | 2 --
 .../src/aws/sdk/kotlin/services/rds/RdsAuthTokenGenerator.kt    | 2 --
 2 files changed, 4 deletions(-)

diff --git a/services/dsql/common/src/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGenerator.kt b/services/dsql/common/src/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGenerator.kt
index 2bfef0e72c9..e286c5f7514 100644
--- a/services/dsql/common/src/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGenerator.kt
+++ b/services/dsql/common/src/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGenerator.kt
@@ -18,13 +18,11 @@ import kotlin.time.Duration.Companion.seconds
 /**
  * Generates an IAM authentication token for use with DSQL databases
  * @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 DsqlAuthTokenGenerator(
     public val credentialsProvider: CredentialsProvider = DefaultChainCredentialsProvider(),
-    public val credentialsRefreshBuffer: Duration = 10.seconds,
     public val signer: AwsSigner = DefaultAwsSigner,
     public val clock: Clock = Clock.System,
 ) {
diff --git a/services/rds/common/src/aws/sdk/kotlin/services/rds/RdsAuthTokenGenerator.kt b/services/rds/common/src/aws/sdk/kotlin/services/rds/RdsAuthTokenGenerator.kt
index 3d19af00d9b..686ffaca67e 100644
--- a/services/rds/common/src/aws/sdk/kotlin/services/rds/RdsAuthTokenGenerator.kt
+++ b/services/rds/common/src/aws/sdk/kotlin/services/rds/RdsAuthTokenGenerator.kt
@@ -19,13 +19,11 @@ import kotlin.time.Duration.Companion.seconds
 /**
  * Generates an IAM authentication token for use with RDS databases
  * @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 RdsAuthTokenGenerator(
     public val credentialsProvider: CredentialsProvider = DefaultChainCredentialsProvider(),
-    public val credentialsRefreshBuffer: Duration = 10.seconds,
     public val signer: AwsSigner = DefaultAwsSigner,
     public val clock: Clock = Clock.System,
 ) {

From c1ddba96d9d85ab1b7f7421a3faaeefe228fa4d0 Mon Sep 17 00:00:00 2001
From: Matas Lauzadis <lauzmata@amazon.com>
Date: Thu, 9 Jan 2025 17:21:13 -0500
Subject: [PATCH 17/24] `put` to `decodedParameters`

---
 .../aws/sdk/kotlin/services/rds/RdsAuthTokenGenerator.kt  | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/services/rds/common/src/aws/sdk/kotlin/services/rds/RdsAuthTokenGenerator.kt b/services/rds/common/src/aws/sdk/kotlin/services/rds/RdsAuthTokenGenerator.kt
index 686ffaca67e..7816debc3cc 100644
--- a/services/rds/common/src/aws/sdk/kotlin/services/rds/RdsAuthTokenGenerator.kt
+++ b/services/rds/common/src/aws/sdk/kotlin/services/rds/RdsAuthTokenGenerator.kt
@@ -38,11 +38,9 @@ public class RdsAuthTokenGenerator(
      */
     public suspend fun generateAuthToken(endpoint: Url, region: String, username: String, expiration: Duration = 900.seconds): String {
         val endpoint = endpoint.toBuilder().apply {
-            parameters.apply {
-                decodedParameters {
-                    add("Action", "connect")
-                    add("DBUser", username)
-                }
+            parameters.decodedParameters.apply {
+                put("Action", "connect")
+                put("DBUser", username)
             }
         }.build()
 

From d3e1cfb45afbef4391a53f77ae21422e5bc63db5 Mon Sep 17 00:00:00 2001
From: Matas Lauzadis <lauzmata@amazon.com>
Date: Thu, 9 Jan 2025 18:48:00 -0500
Subject: [PATCH 18/24] ktlint

---
 .../src/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGenerator.kt   | 1 -
 .../src/aws/sdk/kotlin/services/rds/RdsAuthTokenGenerator.kt     | 1 -
 2 files changed, 2 deletions(-)

diff --git a/services/dsql/common/src/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGenerator.kt b/services/dsql/common/src/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGenerator.kt
index e286c5f7514..59359bf457b 100644
--- a/services/dsql/common/src/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGenerator.kt
+++ b/services/dsql/common/src/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGenerator.kt
@@ -5,7 +5,6 @@
 package aws.sdk.kotlin.services.dsql
 
 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.AuthTokenGenerator
 import aws.smithy.kotlin.runtime.auth.awssigning.AwsSigner
diff --git a/services/rds/common/src/aws/sdk/kotlin/services/rds/RdsAuthTokenGenerator.kt b/services/rds/common/src/aws/sdk/kotlin/services/rds/RdsAuthTokenGenerator.kt
index 7816debc3cc..a855a675f26 100644
--- a/services/rds/common/src/aws/sdk/kotlin/services/rds/RdsAuthTokenGenerator.kt
+++ b/services/rds/common/src/aws/sdk/kotlin/services/rds/RdsAuthTokenGenerator.kt
@@ -5,7 +5,6 @@
 package aws.sdk.kotlin.services.rds
 
 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.AuthTokenGenerator
 import aws.smithy.kotlin.runtime.auth.awssigning.AwsSigner

From d88a88dbe2e3337677749be54ebb4919f9efa86a Mon Sep 17 00:00:00 2001
From: Matas Lauzadis <lauzmata@amazon.com>
Date: Fri, 10 Jan 2025 11:56:49 -0500
Subject: [PATCH 19/24] Add tests comparing X-Amz-Date header to clock

---
 .../kotlin/services/dsql/DsqlAuthTokenGeneratorTest.kt   | 9 +++++++++
 services/rds/common/test/RdsAuthTokenGeneratorTest.kt    | 5 +++++
 2 files changed, 14 insertions(+)

diff --git a/services/dsql/common/test/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGeneratorTest.kt b/services/dsql/common/test/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGeneratorTest.kt
index 069bc8d6241..0fa84a87d9b 100644
--- a/services/dsql/common/test/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGeneratorTest.kt
+++ b/services/dsql/common/test/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGeneratorTest.kt
@@ -13,6 +13,7 @@ import aws.smithy.kotlin.runtime.time.ManualClock
 import kotlinx.coroutines.test.runTest
 import kotlin.test.Test
 import kotlin.test.assertContains
+import kotlin.test.assertEquals
 import kotlin.test.assertFalse
 import kotlin.time.Duration.Companion.seconds
 
@@ -40,6 +41,10 @@ class DsqlAuthTokenGeneratorTest {
         listOf("http://", "https://").forEach {
             assertFalse(token.contains(it))
         }
+
+        val urlToken = Url.parse("https://$token")
+        val xAmzDate = urlToken.parameters.decodedParameters.getValue("X-Amz-Date").single()
+        assertEquals(clock.now(), Instant.fromIso8601(xAmzDate))
     }
 
     @Test
@@ -65,5 +70,9 @@ class DsqlAuthTokenGeneratorTest {
         listOf("http://", "https://").forEach {
             assertFalse(token.contains(it))
         }
+
+        val urlToken = Url.parse("https://$token")
+        val xAmzDate = urlToken.parameters.decodedParameters.getValue("X-Amz-Date").single()
+        assertEquals(clock.now(), Instant.fromIso8601(xAmzDate))
     }
 }
diff --git a/services/rds/common/test/RdsAuthTokenGeneratorTest.kt b/services/rds/common/test/RdsAuthTokenGeneratorTest.kt
index 5a1f76ad37c..ceb34f77d23 100644
--- a/services/rds/common/test/RdsAuthTokenGeneratorTest.kt
+++ b/services/rds/common/test/RdsAuthTokenGeneratorTest.kt
@@ -13,6 +13,7 @@ import aws.smithy.kotlin.runtime.time.ManualClock
 import kotlinx.coroutines.test.runTest
 import kotlin.test.Test
 import kotlin.test.assertContains
+import kotlin.test.assertEquals
 import kotlin.test.assertFalse
 import kotlin.time.Duration.Companion.seconds
 
@@ -45,5 +46,9 @@ class RdsAuthTokenGeneratorTest {
         listOf("http://", "https://").forEach {
             assertFalse(token.contains(it))
         }
+
+        val urlToken = Url.parse("https://$token")
+        val xAmzDate = urlToken.parameters.decodedParameters.getValue("X-Amz-Date").single()
+        assertEquals(clock.now(), Instant.fromIso8601(xAmzDate))
     }
 }

From 23146e76628eb1ea3a1425ddd997ba50a9753649 Mon Sep 17 00:00:00 2001
From: Matas Lauzadis <lauzmata@amazon.com>
Date: Fri, 10 Jan 2025 12:11:48 -0500
Subject: [PATCH 20/24] Update to latest version of smithy-kotlin

---
 gradle/libs.versions.toml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml
index b601621eb49..ab06432f1e1 100644
--- a/gradle/libs.versions.toml
+++ b/gradle/libs.versions.toml
@@ -11,8 +11,8 @@ coroutines-version = "1.9.0"
 atomicfu-version = "0.25.0"
 
 # smithy-kotlin codegen and runtime are versioned separately
-smithy-kotlin-runtime-version = "1.3.31"
-smithy-kotlin-codegen-version = "0.33.31"
+smithy-kotlin-runtime-version = "1.3.34"
+smithy-kotlin-codegen-version = "0.33.34"
 
 # codegen
 smithy-version = "1.51.0"

From 4156091fe284a8af581b63027838e74955918df0 Mon Sep 17 00:00:00 2001
From: Matas Lauzadis <lauzmata@amazon.com>
Date: Thu, 16 Jan 2025 10:39:11 -0500
Subject: [PATCH 21/24] Ensure `Host` header is signed

---
 .../aws/sdk/kotlin/services/dsql/DsqlAuthTokenGeneratorTest.kt  | 2 ++
 services/rds/common/test/RdsAuthTokenGeneratorTest.kt           | 1 +
 2 files changed, 3 insertions(+)

diff --git a/services/dsql/common/test/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGeneratorTest.kt b/services/dsql/common/test/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGeneratorTest.kt
index 0fa84a87d9b..4cfb51c8d6e 100644
--- a/services/dsql/common/test/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGeneratorTest.kt
+++ b/services/dsql/common/test/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGeneratorTest.kt
@@ -36,6 +36,7 @@ class DsqlAuthTokenGeneratorTest {
         assertContains(token, "peccy.dsql.us-east-1.on.aws?Action=DbConnect")
         assertContains(token, "X-Amz-Credential=akid%2F20240827%2Fus-east-1%2Fdsql%2Faws4_request")
         assertContains(token, "X-Amz-Expires=450")
+        assertContains(token, "X-Amz-SignedHeaders=host")
 
         // Token should not contain a scheme
         listOf("http://", "https://").forEach {
@@ -65,6 +66,7 @@ class DsqlAuthTokenGeneratorTest {
         assertContains(token, "peccy.dsql.us-east-1.on.aws?Action=DbConnectAdmin")
         assertContains(token, "X-Amz-Credential=akid%2F20240827%2Fus-east-1%2Fdsql%2Faws4_request")
         assertContains(token, "X-Amz-Expires=450")
+        assertContains(token, "X-Amz-SignedHeaders=host")
 
         // Token should not contain a scheme
         listOf("http://", "https://").forEach {
diff --git a/services/rds/common/test/RdsAuthTokenGeneratorTest.kt b/services/rds/common/test/RdsAuthTokenGeneratorTest.kt
index ceb34f77d23..9865a2ede25 100644
--- a/services/rds/common/test/RdsAuthTokenGeneratorTest.kt
+++ b/services/rds/common/test/RdsAuthTokenGeneratorTest.kt
@@ -41,6 +41,7 @@ class RdsAuthTokenGeneratorTest {
         assertContains(token, "prod-instance.us-east-1.rds.amazonaws.com:3306?Action=connect&DBUser=peccy")
         assertContains(token, "X-Amz-Credential=akid%2F20240827%2Fus-east-1%2Frds-db%2Faws4_request")
         assertContains(token, "X-Amz-Expires=450")
+        assertContains(token, "X-Amz-SignedHeaders=host")
 
         // Token should not contain a scheme
         listOf("http://", "https://").forEach {

From 755b7e2f45abdf04543423f64059905b7b986321 Mon Sep 17 00:00:00 2001
From: Matas Lauzadis <lauzmata@amazon.com>
Date: Thu, 16 Jan 2025 11:06:38 -0500
Subject: [PATCH 22/24] Upgrade to v4

---
 .github/workflows/continuous-integration.yml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/continuous-integration.yml b/.github/workflows/continuous-integration.yml
index 684d610607a..4000a61f4f2 100644
--- a/.github/workflows/continuous-integration.yml
+++ b/.github/workflows/continuous-integration.yml
@@ -42,7 +42,7 @@ jobs:
           ./gradlew -Ptest.java.version=${{ matrix.java-version }} jvmTest --stacktrace
       - name: Save Test Reports
         if: failure()
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
           name: test-reports
           path: '**/build/reports'
@@ -80,7 +80,7 @@ jobs:
           ./gradlew testAllProtocols
       - name: Save Test Reports
         if: failure()
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
           name: test-reports
           path: '**/build/reports'

From 13aea714eed63dcabdd9d06d360af426d2a3dade Mon Sep 17 00:00:00 2001
From: Matas Lauzadis <lauzmata@amazon.com>
Date: Thu, 16 Jan 2025 14:41:31 -0500
Subject: [PATCH 23/24] Bump to latest build plugin version

---
 gradle/libs.versions.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml
index 9ddae9e8179..08ee57fd278 100644
--- a/gradle/libs.versions.toml
+++ b/gradle/libs.versions.toml
@@ -4,7 +4,7 @@ ksp-version = "2.1.0-1.0.29" # Keep in sync with kotlin-version
 
 dokka-version = "1.9.10"
 
-aws-kotlin-repo-tools-version = "0.4.17"
+aws-kotlin-repo-tools-version = "0.4.18"
 
 # libs
 coroutines-version = "1.9.0"

From bcece78aeda87499727ec9d9bcea43620200757a Mon Sep 17 00:00:00 2001
From: Matas Lauzadis <lauzmata@amazon.com>
Date: Thu, 16 Jan 2025 15:21:19 -0500
Subject: [PATCH 24/24] Bump smithy-kotlin

---
 gradle/libs.versions.toml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml
index 08ee57fd278..2195cfc3a42 100644
--- a/gradle/libs.versions.toml
+++ b/gradle/libs.versions.toml
@@ -12,8 +12,8 @@ atomicfu-version = "0.25.0"
 binary-compatibility-validator-version = "0.16.3"
 
 # smithy-kotlin codegen and runtime are versioned separately
-smithy-kotlin-runtime-version = "1.4.0"
-smithy-kotlin-codegen-version = "0.34.0"
+smithy-kotlin-runtime-version = "1.4.1"
+smithy-kotlin-codegen-version = "0.34.1"
 
 # codegen
 smithy-version = "1.53.0"