Abstract AuthTokenGenerator and implement RDS AuthTokenGenerator

Author: lauzadis

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)
     }
 }

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)
+    }
+}

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))
+        }
+    }
+}