From 3b9b2b211ce298b906ed0fe49d50b9b9eb6140b7 Mon Sep 17 00:00:00 2001 From: Matas Lauzadis Date: Thu, 9 Jan 2025 11:05:34 -0500 Subject: [PATCH] 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 + + 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