chore: Adding mechanism to provide plugin extensions (#2852)

Author: sebaland

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/AWSCognitoAuthPlugin+Configure.swift

@@ -29,6 +29,8 @@ extension AWSCognitoAuthPlugin {
                 AuthPluginErrorConstants.decodeConfigurationError.recoverySuggestion)
         }
 
+        jsonConfiguration = jsonValueConfiguration
+
         let authConfiguration = try ConfigurationHelper.authConfiguration(jsonValueConfiguration)
 
         let credentialStoreResolver = CredentialStoreState.Resolver().eraseToAnyResolver()
@@ -89,6 +91,13 @@ extension AWSCognitoAuthPlugin {
                 frameworkMetadata: AmplifyAWSServiceConfiguration.frameworkMetaData(),
                 region: userPoolConfig.region
             )
+
+            if var httpClientEngineProxy = httpClientEngineProxy {
+                let sdkEngine = configuration.httpClientEngine
+                httpClientEngineProxy.target = sdkEngine
+                configuration.httpClientEngine = httpClientEngineProxy
+            }
+
             return CognitoIdentityProviderClient(config: configuration)
         default:
             fatalError()

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/AWSCognitoAuthPlugin+PluginExtension.swift

@@ -0,0 +1,19 @@
+//
+// Copyright Amazon.com Inc. or its affiliates.
+// All Rights Reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+@_spi(InternalAmplifyPluginExtension) import AWSPluginsCore
+import Foundation
+import ClientRuntime
+
+extension AWSCognitoAuthPlugin {
+    @_spi(InternalAmplifyPluginExtension)
+    public func add(pluginExtension: AWSPluginExtension) {
+        if let customHttpEngine = pluginExtension as? HttpClientEngineProxy {
+            self.httpClientEngineProxy = customHttpEngine
+        }
+    }
+}

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/AWSCognitoAuthPlugin.swift

@@ -7,6 +7,7 @@
 
 import Foundation
 import Amplify
+import AWSPluginsCore
 
 public final class AWSCognitoAuthPlugin: AWSCognitoAuthPluginBehavior {
 
@@ -29,6 +30,11 @@ public final class AWSCognitoAuthPlugin: AWSCognitoAuthPluginBehavior {
 
     var taskQueue: TaskQueue<Any>!
 
+    var httpClientEngineProxy: HttpClientEngineProxy?
+
+    @_spi(InternalAmplifyConfiguration)
+    internal(set) public var jsonConfiguration: JSONValue?
+
     /// The unique key of the plugin within the auth category.
     public var key: PluginKey {
         return "awsCognitoAuthPlugin"

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Support/Utils/HttpClientEngineProxy.swift

@@ -0,0 +1,16 @@
+//
+// Copyright Amazon.com Inc. or its affiliates.
+// All Rights Reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+@_spi(InternalHttpEngineProxy) @_spi(InternalAmplifyPluginExtension) import AWSPluginsCore
+import ClientRuntime
+import Foundation
+
+protocol HttpClientEngineProxy: HttpClientEngine {
+    var target: HttpClientEngine? { set get }
+}
+
+extension UserAgentSuffixAppender: HttpClientEngineProxy {}

AmplifyPlugins/Core/AWSPluginsCore/Utils/AWSPluginExtension.swift

@@ -0,0 +1,11 @@
+//
+// Copyright Amazon.com Inc. or its affiliates.
+// All Rights Reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+import Foundation
+
+@_spi(InternalAmplifyPluginExtension)
+public protocol AWSPluginExtension {}

AmplifyPlugins/Core/AWSPluginsCore/Utils/CustomHttpClientEngine/UserAgentSuffixAppender.swift

@@ -0,0 +1,41 @@
+//
+// Copyright Amazon.com Inc. or its affiliates.
+// All Rights Reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+import ClientRuntime
+
+@_spi(InternalAmplifyPluginExtension)
+public class UserAgentSuffixAppender: AWSPluginExtension {
+    @_spi(InternalHttpEngineProxy)
+    public var target: HttpClientEngine? = nil
+    public let suffix: String
+    private let userAgentHeader = "User-Agent"
+
+    public init(suffix: String) {
+        self.suffix = suffix
+    }
+}
+
+@_spi(InternalHttpEngineProxy)
+extension UserAgentSuffixAppender: HttpClientEngine {
+    public func execute(request: SdkHttpRequest) async throws -> HttpResponse {
+        guard let target = target  else {
+            throw ClientError.unknownError("HttpClientEngine is not set")
+        }
+        var headers = request.headers
+        let currentUserAgent = headers.value(for: userAgentHeader) ?? ""
+        headers.update(
+            name: userAgentHeader,
+            value: "\(currentUserAgent) \(suffix)"
+        )
+        request.headers = headers
+        return try await target.execute(request: request)
+    }
+
+    public func close() async {
+        await target?.close()
+    }
+}

AmplifyPlugins/Core/AWSPluginsCoreTests/Utils/UserAgentSuffixAppenderTests.swift

@@ -0,0 +1,102 @@
+//
+// Copyright Amazon.com Inc. or its affiliates.
+// All Rights Reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+@_spi(InternalAmplifyPluginExtension) @_spi(InternalHttpEngineProxy) import AWSPluginsCore
+import ClientRuntime
+import XCTest
+
+class UserAgentSuffixAppenderTests: XCTestCase {
+    private let userAgentKey = "User-Agent"
+    private let customSuffix = "myCustomSuffix"
+    private var appender: UserAgentSuffixAppender!
+    private var httpClientEngine: MockHttpClientEngine!
+
+    override func setUp() {
+        appender = UserAgentSuffixAppender(suffix: customSuffix)
+        httpClientEngine = MockHttpClientEngine()
+        appender.target = httpClientEngine
+    }
+
+    override func tearDown() {
+        appender = nil
+        httpClientEngine = nil
+    }
+
+    /// Given: A UserAgentSuffixAppender with a configured httpClientEngine
+    /// When: A request is invoked with an existing User-Agent
+    /// Then: The user agent suffix is appended
+    func testExecute_withExistingUserAgentHeader_shouldAppendSuffix() async throws {
+        let request = createRequest()
+        request.headers.add(name: userAgentKey, value: "existingUserAgent")
+
+        _ = try await appender.execute(request: request)
+        XCTAssertEqual(httpClientEngine.executeCount, 1)
+        XCTAssertNotNil(httpClientEngine.executeRequest)
+        let userAgent = try XCTUnwrap(request.headers.value(for: userAgentKey))
+        XCTAssertTrue(userAgent.hasSuffix(customSuffix))
+    }
+
+    /// Given: A UserAgentSuffixAppender with a configured httpClientEngine
+    /// When: A request is invoked with no User-Agent
+    /// Then: The user agent is created containing the suffix
+    func testExecute_withoutExistingUserAgentHeader_shouldCreateHeader() async throws {
+        let request = createRequest()
+
+        _ = try await appender.execute(request: request)
+        XCTAssertEqual(httpClientEngine.executeCount, 1)
+        XCTAssertNotNil(httpClientEngine.executeRequest)
+        let userAgent = try XCTUnwrap(request.headers.value(for: userAgentKey))
+        XCTAssertTrue(userAgent.hasSuffix(customSuffix))
+    }
+
+    /// Given: A UserAgentSuffixAppender with no httpClientEngine configured
+    /// When: A request is invoked
+    /// Then: An error is thrown
+    func testExecute_withoutHttpClientEngine_shouldThrowError() async {
+        appender = UserAgentSuffixAppender(suffix: customSuffix)
+        do {
+            _ = try await appender.execute(request: createRequest())
+            XCTFail("Should not succeed")
+        } catch {
+            guard case ClientError.unknownError(_) = error else {
+                XCTFail("Expected .unknownError, got \(error)")
+                return
+            }
+        }
+    }
+
+    /// Given: A UserAgentSuffixAppender with a configured httpClientEngine
+    /// When: close is invoked
+    /// Then: The httpClientEngine's close API should be called
+    func testClose_shouldInvokeClose() async {
+        await appender.close()
+        XCTAssertEqual(httpClientEngine.closeCount, 1)
+    }
+
+    private func createRequest() -> SdkHttpRequest {
+        return SdkHttpRequest(
+            method: .get,
+            endpoint: .init(host: "customHost"),
+            headers: .init()
+        )
+    }
+}
+
+private class MockHttpClientEngine: HttpClientEngine {
+    var executeCount = 0
+    var executeRequest: SdkHttpRequest?
+    func execute(request: SdkHttpRequest) async throws -> HttpResponse {
+        executeCount += 1
+        executeRequest = request
+        return .init(body: .empty, statusCode: .accepted)
+    }
+
+    var closeCount = 0
+    func close() async {
+        closeCount += 1
+    }
+}