feat: Auth token generator for AWS DSQL

Sources/Core/AWSClientRuntime/Sources/AWSClientRuntime/Customizations/RDS/AuthTokenGenerator.swift → Sources/Core/AWSClientRuntime/Sources/AWSClientRuntime/Customizations/AuthTokenGenerator.swift

@@ -16,7 +16,7 @@ import SmithyIdentity
 import struct Foundation.Date
 import struct Foundation.TimeInterval
 
-/// A utility class with a single utility method that generates IAM authentication token used for connecting to RDS.
+/// A utility class with utility methods that generate IAM authentication token used for connecting to RDS & Aurora DSQL.
 @_spi(AuthTokenGenerator)
 public class AuthTokenGenerator {
     public var awsCredentialIdentity: AWSCredentialIdentity
@@ -41,15 +41,15 @@ public class AuthTokenGenerator {
         self.awsCredentialIdentity = try await awsCredentialIdentityResolver.getIdentity()
     }
 
-    /// Generates authenetication token using given inputs to the method and credential identity instance variable.
+    /// Generates authenetication token for RDS using given inputs to the method and credential identity instance variable.
     ///
     /// - Parameters:
     ///   - endpoint: The endpoint of the RDS instance. E.g., `rdsmysql.123456789012.us-west-2.rds.amazonaws.com`
     ///   - port: The port of the RDS instance to connect to. E.g., `3306`
     ///   - region: The region that RDS instance is located in. E.g., `us-west-2`
     ///   - username: The username of the RDS database user. E.g., `admin`
     ///   - expiration: The expiration for the token in seconds. Default is 900 seconds (15 minutes).
-    public func generateAuthToken(
+    public func generateRDSAuthToken(
         endpoint: String,
         port: Int16,
         region: String,
@@ -97,4 +97,51 @@ public class AuthTokenGenerator {
 
         return rdsAuthToken
     }
+
+    public func generateDSQLAuthToken(
+        endpoint: String,
+        region: String,
+        expiration: TimeInterval = 900,
+        isForAdmin: Bool
+    ) async throws -> String {
+        CommonRuntimeKit.initialize()
+        let requestBuilder = HTTPRequestBuilder()
+        requestBuilder.withHost(endpoint)
+
+        // Add the Host header and the required query items for the desired presigned URL.
+        requestBuilder.withHeader(name: "Host", value: "\(endpoint)")
+        let actionQueryItemValue = isForAdmin ? "DbConnectAdmin" : "DbConnect"
+        requestBuilder.withQueryItem(URIQueryItem(name: "Action", value: actionQueryItemValue))
+
+        let signingConfig = AWSSigningConfig(
+            credentials: self.awsCredentialIdentity,
+            expiration: expiration,
+            signedBodyValue: .empty,
+            flags: SigningFlags(
+                useDoubleURIEncode: true,
+                shouldNormalizeURIPath: true,
+                omitSessionToken: false
+            ),
+            date: Date(),
+            service: "dsql",
+            region: region,
+            signatureType: .requestQueryParams,
+            signingAlgorithm: .sigv4
+        )
+
+        let signedRequest = await AWSSigV4Signer().sigV4SignedRequest(
+            requestBuilder: requestBuilder,
+            signingConfig: signingConfig
+        )
+
+        guard let presignedURL = signedRequest?.destination.url else {
+            throw ClientError.authError("Failed to generate auth token for RDS.")
+        }
+
+        // Remove https:// from the presigned URL to get final value for RDS auth token.
+        let startIndex = presignedURL.absoluteString.index(presignedURL.absoluteString.startIndex, offsetBy: 8)
+        let rdsAuthToken = String(presignedURL.absoluteString[startIndex...])
+
+        return rdsAuthToken
+    }
 }

Sources/Services/AWSDSQL/Sources/AWSDSQL/AuthTokenGenerator.swift

@@ -0,0 +1,77 @@
+//
+// Copyright Amazon.com Inc. or its affiliates.
+// All Rights Reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+// Code generated by smithy-swift-codegen. DO NOT EDIT!
+
+
+
+@_spi(AuthTokenGenerator) import class AWSClientRuntime.AuthTokenGenerator
+import SmithyIdentity
+import struct Foundation.TimeInterval
+
+/// A utility class with methods that generate IAM authentication token used for connecting to DSQL.
+public class AuthTokenGenerator {
+    private let generator: AWSClientRuntime.AuthTokenGenerator
+
+    /// The initializer that takes in AWSCredentialIdentity struct to use to generate the IAM authentication token.
+    public init(awsCredentialIdentity: AWSCredentialIdentity) {
+        self.generator = AWSClientRuntime.AuthTokenGenerator(awsCredentialIdentity: awsCredentialIdentity)
+    }
+
+    /// The initializer that takes in a specific AWSCredentialIdentityResolver, used to resolve the AWSCredentialIdentity used to generate the IAM authentication token.
+    public init(awsCredentialIdentityResolver: any AWSCredentialIdentityResolver) async throws {
+        self.generator = try await AWSClientRuntime.AuthTokenGenerator(awsCredentialIdentityResolver: awsCredentialIdentityResolver)
+    }
+
+    /// Updates the AWS credentials used to generate the IAM auth token.
+    public func updateCredentials(newAWSCredentialIdentity: AWSCredentialIdentity) {
+        generator.updateCredentials(newAWSCredentialIdentity: newAWSCredentialIdentity)
+    }
+
+    /// Updates the AWS credentials used to generate the IAM auth token by resolving credentials from passed in resolver.
+    public func updateCredentials(awsCredentialIdentityResolver: any AWSCredentialIdentityResolver) async throws {
+        try await generator.updateCredentials(awsCredentialIdentityResolver: awsCredentialIdentityResolver)
+    }
+
+    /// Generates authenetication token for non-admin connection using given inputs to the method and credential identity instance variable.
+    ///
+    /// - Parameters:
+    ///   - endpoint: The endpoint of the RDS instance. E.g., `peccy.dsql.us-east-1.on.aws`
+    ///   - region: The region that RDS instance is located in. E.g., `us-east-1`
+    ///   - expiration: The expiration for the token in seconds. Default is 900 seconds (15 minutes).
+    public func generateDBConnectAuthToken(
+        endpoint: String,
+        region: String,
+        expiration: TimeInterval = 900
+    ) async throws -> String {
+        return try await generator.generateDSQLAuthToken(
+            endpoint: endpoint,
+            region: region,
+            expiration: expiration,
+            isForAdmin: false
+        )
+    }
+
+    /// Generates authenetication token for admin connection using given inputs to the method and credential identity instance variable.
+    ///
+    /// - Parameters:
+    ///   - endpoint: The endpoint of the RDS instance. E.g., `peccy.dsql.us-east-1.on.aws`
+    ///   - region: The region that RDS instance is located in. E.g., `us-east-1`
+    ///   - expiration: The expiration for the token in seconds. Default is 900 seconds (15 minutes).
+    public func generateDBConnectAdminAuthToken(
+        endpoint: String,
+        region: String,
+        expiration: TimeInterval = 900
+    ) async throws -> String {
+        return try await generator.generateDSQLAuthToken(
+            endpoint: endpoint,
+            region: region,
+            expiration: expiration,
+            isForAdmin: true
+        )
+    }
+}

Sources/Services/AWSRDS/Sources/AWSRDS/AuthTokenGenerator.swift

@@ -52,7 +52,7 @@ public class AuthTokenGenerator {
         username: String,
         expiration: TimeInterval = 900
     ) async throws -> String {
-        return try await generator.generateAuthToken(
+        return try await generator.generateRDSAuthToken(
             endpoint: endpoint,
             port: port,
             region: region,

codegen/smithy-aws-swift-codegen/src/main/kotlin/software/amazon/smithy/aws/swift/codegen/customization/dsql/AuthTokenGeneratorIntegration.kt

@@ -0,0 +1,94 @@
+package software.amazon.smithy.aws.swift.codegen.customization.dsql
+
+import software.amazon.smithy.aws.swift.codegen.sdkId
+import software.amazon.smithy.model.Model
+import software.amazon.smithy.model.shapes.ServiceShape
+import software.amazon.smithy.swift.codegen.SwiftDelegator
+import software.amazon.smithy.swift.codegen.SwiftSettings
+import software.amazon.smithy.swift.codegen.core.SwiftCodegenContext
+import software.amazon.smithy.swift.codegen.integration.ProtocolGenerator
+import software.amazon.smithy.swift.codegen.integration.SwiftIntegration
+import software.amazon.smithy.swift.codegen.model.expectShape
+
+class AuthTokenGeneratorIntegration : SwiftIntegration {
+    override fun enabledForService(model: Model, settings: SwiftSettings): Boolean =
+        model.expectShape<ServiceShape>(settings.service).sdkId == "DSQL"
+
+    override fun writeAdditionalFiles(
+        ctx: SwiftCodegenContext,
+        protocolGenerationContext: ProtocolGenerator.GenerationContext,
+        delegator: SwiftDelegator
+    ) {
+        delegator.useFileWriter("Sources/${ctx.settings.moduleName}/AuthTokenGenerator.swift") { writer ->
+            val authTokenGeneratorWrapperClass = """
+            @_spi(AuthTokenGenerator) import class AWSClientRuntime.AuthTokenGenerator
+            import SmithyIdentity
+            import struct Foundation.TimeInterval
+
+            /// A utility class with methods that generate IAM authentication token used for connecting to DSQL.
+            public class AuthTokenGenerator {
+                private let generator: AWSClientRuntime.AuthTokenGenerator
+
+                /// The initializer that takes in AWSCredentialIdentity struct to use to generate the IAM authentication token.
+                public init(awsCredentialIdentity: AWSCredentialIdentity) {
+                    self.generator = AWSClientRuntime.AuthTokenGenerator(awsCredentialIdentity: awsCredentialIdentity)
+                }
+
+                /// The initializer that takes in a specific AWSCredentialIdentityResolver, used to resolve the AWSCredentialIdentity used to generate the IAM authentication token.
+                public init(awsCredentialIdentityResolver: any AWSCredentialIdentityResolver) async throws {
+                    self.generator = try await AWSClientRuntime.AuthTokenGenerator(awsCredentialIdentityResolver: awsCredentialIdentityResolver)
+                }
+
+                /// Updates the AWS credentials used to generate the IAM auth token.
+                public func updateCredentials(newAWSCredentialIdentity: AWSCredentialIdentity) {
+                    generator.updateCredentials(newAWSCredentialIdentity: newAWSCredentialIdentity)
+                }
+
+                /// Updates the AWS credentials used to generate the IAM auth token by resolving credentials from passed in resolver.
+                public func updateCredentials(awsCredentialIdentityResolver: any AWSCredentialIdentityResolver) async throws {
+                    try await generator.updateCredentials(awsCredentialIdentityResolver: awsCredentialIdentityResolver)
+                }
+
+                /// Generates authenetication token for non-admin connection using given inputs to the method and credential identity instance variable.
+                ///
+                /// - Parameters:
+                ///   - endpoint: The endpoint of the RDS instance. E.g., `peccy.dsql.us-east-1.on.aws`
+                ///   - region: The region that RDS instance is located in. E.g., `us-east-1`
+                ///   - expiration: The expiration for the token in seconds. Default is 900 seconds (15 minutes).
+                public func generateDBConnectAuthToken(
+                    endpoint: String,
+                    region: String,
+                    expiration: TimeInterval = 900
+                ) async throws -> String {
+                    return try await generator.generateDSQLAuthToken(
+                        endpoint: endpoint,
+                        region: region,
+                        expiration: expiration,
+                        isForAdmin: false
+                    )
+                }
+
+                /// Generates authenetication token for admin connection using given inputs to the method and credential identity instance variable.
+                ///
+                /// - Parameters:
+                ///   - endpoint: The endpoint of the RDS instance. E.g., `peccy.dsql.us-east-1.on.aws`
+                ///   - region: The region that RDS instance is located in. E.g., `us-east-1`
+                ///   - expiration: The expiration for the token in seconds. Default is 900 seconds (15 minutes).
+                public func generateDBConnectAdminAuthToken(
+                    endpoint: String,
+                    region: String,
+                    expiration: TimeInterval = 900
+                ) async throws -> String {
+                    return try await generator.generateDSQLAuthToken(
+                        endpoint: endpoint,
+                        region: region,
+                        expiration: expiration,
+                        isForAdmin: true
+                    )
+                }
+            }
+            """.trimIndent()
+            writer.write(authTokenGeneratorWrapperClass)
+        }
+    }
+}

codegen/smithy-aws-swift-codegen/src/main/kotlin/software/amazon/smithy/aws/swift/codegen/customization/rds/AuthTokenGeneratorIntegration.kt

@@ -68,7 +68,7 @@ class AuthTokenGeneratorIntegration : SwiftIntegration {
                     username: String,
                     expiration: TimeInterval = 900
                 ) async throws -> String {
-                    return try await generator.generateAuthToken(
+                    return try await generator.generateRDSAuthToken(
                         endpoint: endpoint,
                         port: port,
                         region: region,

codegen/smithy-aws-swift-codegen/src/main/resources/META-INF/services/software.amazon.smithy.swift.codegen.integration.SwiftIntegration

@@ -27,3 +27,4 @@ software.amazon.smithy.swift.codegen.swiftintegrations.InitialRequestIntegration
 software.amazon.smithy.aws.swift.codegen.swiftintegrations.RegistryConfigIntegration
 software.amazon.smithy.aws.swift.codegen.swiftintegrations.AmzSdkRetryHeadersIntegration
 software.amazon.smithy.aws.swift.codegen.customization.rds.AuthTokenGeneratorIntegration
+software.amazon.smithy.aws.swift.codegen.customization.dsql.AuthTokenGeneratorIntegration