Fix API Keys (#31)

* temporary disable API Keys until awslabs/aws-sdk-swift#1979 is fixed

* swift format

* add a note in the README

* pave the road towards using a bearer token provider

* remove debugging statements

* swift-format

* fix bearer token authentication

* fix tests

* fix format

* add Examples/api-key in the list of integration tests

Author: sebsto

.github/workflows/pull_request.yml

@@ -35,7 +35,7 @@ jobs:
       # We pass the list of examples here, but we can't pass an array as argument
       # Instead, we pass a String with a valid JSON array.
       # The workaround is mentioned here https://github.com/orgs/community/discussions/11692
-      examples: "[ 'converse', 'converse-stream', 'text_chat' ]"
+      examples: "[ 'api-key', 'converse', 'converse-stream', 'text_chat' ]"
 
   swift-6-language-mode:
     name: Swift 6 Language Mode

.gitignore

@@ -13,7 +13,7 @@ Package.resolved
 .vscode
 .env
 Makefile
-
+**/temp
 node_modules
 
 # **/backend

README.md

@@ -135,7 +135,9 @@ let bedrock = try await BedrockService(
 )
 ```
 
-### API Key Authentication
+### API Key Authentication (temporary unavailable)
+
+This capability will be available when issue [#1979](https://github.com/awslabs/aws-sdk-swift/issues/1979) from the AWS SDK for Swift will be fixed or a workaround provided.
 
 Use an API key for authentication. API keys are generated in the AWS console and provide a simpler authentication method for specific use cases.
 

Sources/BedrockAuthentication.swift

@@ -56,16 +56,15 @@ public enum BedrockAuthentication: Sendable, CustomStringConvertible {
     /// Creates an AWS credential identity resolver depending on the authentication parameter.
     /// - Parameters:
     ///     - authentication: The authentication type to use
-    /// - Returns: An optional AWS credential identity resolver. A nil return value means that the default AWS credential provider chain will be used.
+    /// - Returns: An optional AWS credential identity resolver. A nil return value means that the default AWS credential provider chain will be used or that the authentication type does not require a specific resolver (like `apiKey`).
     ///
     func getAWSCredentialIdentityResolver(
         logger: Logger
     ) async throws -> (any SmithyIdentity.AWSCredentialIdentityResolver)? {
 
         switch self {
-        case .default,
-            .apiKey(_):
-            return nil
+        case .default, .apiKey(_):
+            return nil  //TODO should we throw an error when apiKey is used ?
         case .profile(let profileName):
             return try? ProfileAWSCredentialIdentityResolver(profileName: profileName)
         case .sso(let profileName):
@@ -84,4 +83,17 @@ public enum BedrockAuthentication: Sendable, CustomStringConvertible {
             return StaticAWSCredentialIdentityResolver(creds)
         }
     }
+
+    /// Creates a BearerTokenIdentityResolver depending on the authentication parameter.
+    /// - Returns: An optional BearerTokenIdentityResolver. A nil return value means that the authentication type requires an AWSCredentialsProvider instead (like `default`, `profile`, `sso`, `webIdentity`, or `static`).
+    /// - Note: Only `apiKey` authentication uses BearerTokenIdentityResolver.
+    func getBearerTokenIdentityResolver(logger: Logger) -> (any SmithyIdentity.BearerTokenIdentityResolver)? {
+        guard case .apiKey(let key) = self else {
+            return nil  // Only apiKey authentication uses BearerTokenIdentityResolver
+        }
+
+        // Create a StaticBearerTokenIdentityResolver with the provided API key
+        let identity = BearerTokenIdentity(token: key)
+        return StaticBearerTokenIdentityResolver(token: identity)
+    }
 }

Sources/BedrockService.swift

@@ -52,7 +52,7 @@ public struct BedrockService: Sendable {
     ) async throws {
         self.logger = logger ?? BedrockService.createLogger("bedrock.service")
         self.logger.trace(
-            "Initializing SwiftBedrock",
+            "Initializing BedrockService",
             metadata: ["region": .string(region.rawValue)]
         )
         self.region = region
@@ -121,7 +121,7 @@ public struct BedrockService: Sendable {
         -> BedrockClient
     {
         let config: BedrockClient.BedrockClientConfiguration = try await prepareConfig(
-            region: region,
+            initialConfig: BedrockClient.BedrockClientConfiguration(region: region.rawValue),
             authentication: authentication,
             logger: logger
         )
@@ -143,7 +143,9 @@ public struct BedrockService: Sendable {
         -> BedrockRuntimeClient
     {
         let config: BedrockRuntimeClient.BedrockRuntimeClientConfiguration = try await prepareConfig(
-            region: region,
+            initialConfig: BedrockRuntimeClient.BedrockRuntimeClientConfiguration(
+                region: region.rawValue
+            ),
             authentication: authentication,
             logger: logger
         )
@@ -152,13 +154,18 @@ public struct BedrockService: Sendable {
 
     /// Generic function to create client configuration and avoid duplication code.
     internal static func prepareConfig<C: BedrockConfigProtocol>(
-        region: Region,
+        initialConfig: C,
         authentication: BedrockAuthentication,
         logger: Logging.Logger
     ) async throws -> C {
-        var config: C = try await .init()
 
-        config.region = region.rawValue
+        var config = initialConfig
+
+        if logger.logLevel == .trace {
+            // enable trace HTTP requests and responses for the SDK
+            // see https://github.com/smithy-lang/smithy-swift/blob/main/Sources/ClientRuntime/Telemetry/Logging/ClientLogMode.swift
+            config.clientLogMode = .requestAndResponse
+        }
 
         // support profile, SSO, web identity and static authentication
         if let awsCredentialIdentityResolver = try? await authentication.getAWSCredentialIdentityResolver(
@@ -168,11 +175,25 @@ public struct BedrockService: Sendable {
         }
 
         // support API keys
-        if case .apiKey(let key) = authentication {
-            config.httpClientConfiguration.defaultHeaders.add(
-                name: "Authorization",
-                value: "Bearer \(key)"
-            )
+        if case .apiKey(_) = authentication {
+            // config.httpClientConfiguration.defaultHeaders.add(
+            //     name: "Authorization",
+            //     value: "Bearer \(key)"
+            // )
+            if let bearerTokenIdentityresolver = authentication.getBearerTokenIdentityResolver(logger: logger) {
+                config.bearerTokenIdentityResolver = bearerTokenIdentityresolver
+
+                // force utilisation of a bearer token instead of AWS credentials + Signv4
+                // see https://github.com/awslabs/aws-sdk-swift/blob/15b8951d108968f767f4199a3c011e27ac519d61/Sources/Services/AWSBedrockRuntime/Sources/AWSBedrockRuntime/AuthSchemeResolver.swift#L58
+                config.authSchemeResolver = DefaultBedrockRuntimeAuthSchemeResolver(authSchemePreference: [
+                    "httpBearerAuth"
+                ])
+            } else {
+                // TODO: should we throw an error here ?
+                logger.error(
+                    "API Key authentication is used but no BearerTokenIdentityResolver is provided. This will lead to issues."
+                )
+            }
             logger.trace("Using API Key for authentication")
         } else {
             logger.trace("Using AWS credentials for authentication")

Sources/Protocols/BedrockConfigProtocol.swift

@@ -16,13 +16,23 @@
 import AWSBedrock
 import AWSBedrockRuntime
 import ClientRuntime
+import SmithyHTTPAuthAPI
 import SmithyIdentity
 
 protocol BedrockConfigProtocol {
-    init() async throws
+    // support regular AWS Credentials + Sigv4 authentication
     var awsCredentialIdentityResolver: any SmithyIdentity.AWSCredentialIdentityResolver { get set }
-    var httpClientConfiguration: ClientRuntime.HttpClientConfiguration { get set }
-    var region: String? { get set }
+
+    // support bearer token authentication (for API Keys)
+    var bearerTokenIdentityResolver: any SmithyIdentity.BearerTokenIdentityResolver { get set }
+    var authSchemeResolver: SmithyHTTPAuthAPI.AuthSchemeResolver { get set }
+
+    // not used at the moment, we use the bearer token instead
+    //var httpClientConfiguration: ClientRuntime.HttpClientConfiguration { get set }
+
+    // for debugging
+    var clientLogMode: ClientRuntime.ClientLogMode { get set }
+
 }
 extension BedrockClient.BedrockClientConfiguration: @retroactive @unchecked Sendable, BedrockConfigProtocol {}
 extension BedrockRuntimeClient.BedrockRuntimeClientConfiguration: @retroactive @unchecked Sendable,

Tests/AuthenticationTests.swift

@@ -14,8 +14,10 @@
 //===----------------------------------------------------------------------===//
 
 import AWSBedrock
+import AWSBedrockRuntime
 import AwsCommonRuntimeKit
 import Logging
+import SmithyIdentity
 import Testing
 
 @testable import BedrockService
@@ -88,16 +90,24 @@ extension BedrockServiceTests {
         // when
         // create bedrock configuration with API Key authentication
         let config: BedrockClient.BedrockClientConfiguration = try await BedrockService.prepareConfig(
-            region: .useast1,
+            initialConfig: BedrockClient.BedrockClientConfiguration(
+                region: "us-east-1"  // default region
+            ),
             authentication: auth,
             logger: Logger(label: "test.logger"),
         )
 
         // then
         #expect(config.region == Region.useast1.rawValue)  // default region
-        #expect(
-            config.httpClientConfiguration.defaultHeaders.value(for: "Authorization") == "Bearer test-api-key-12345"
-        )
+
+        // check token
+        let resolver = config.bearerTokenIdentityResolver as? StaticBearerTokenIdentityResolver
+        let token = try await resolver?.getIdentity(identityProperties: nil).token
+        #expect(token == testApiKey, "Expected token to match the API key")
+
+        // check bearer auth scheme
+        let authScheme = (config.authSchemeResolver as? DefaultBedrockRuntimeAuthSchemeResolver)?.authSchemePreference
+        #expect(authScheme?.contains("httpBearerAuth") == true, "Expected auth scheme to be HTTP Bearer")
 
     }
 
@@ -147,7 +157,9 @@ extension BedrockServiceTests {
 
         // when
         let _: BedrockClient.BedrockClientConfiguration = try await BedrockService.prepareConfig(
-            region: .useast1,
+            initialConfig: BedrockClient.BedrockClientConfiguration(
+                region: "us-east-1"  // default region
+            ),
             authentication: auth,
             logger: logger
         )