feat(Auth): Adding network preferences (#3379)

* feat(Auth): Adding network preferences

* adding tests

* Update AWSCognitoNetworkPreferences.swift

* Update AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Models/AWSCognitoNetworkPreferences.swift

Co-authored-by: Sebastian Villena <[email protected]>

* worked on review comments

---------

Co-authored-by: Sebastian Villena <[email protected]>

Author: harsh62

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

@@ -102,6 +102,15 @@ extension AWSCognitoAuthPlugin {
                 configuration.httpClientEngine = .userAgentEngine(for: configuration)
             }
 
+            if let requestTimeout = networkPreferences?.timeoutIntervalForRequest {
+                let requestTimeOutMs = requestTimeout * 1_000
+                configuration.connectTimeoutMs = UInt32(requestTimeOutMs)
+            }
+
+            if let maxRetryUnwrapped = networkPreferences?.maxRetryCount {
+                configuration.retryStrategyOptions = RetryStrategyOptions(maxRetriesBase:  Int(maxRetryUnwrapped))
+            }
+
             return CognitoIdentityProviderClient(config: configuration)
         default:
             fatalError()
@@ -116,6 +125,15 @@ extension AWSCognitoAuthPlugin {
             )
             configuration.httpClientEngine = .userAgentEngine(for: configuration)
 
+            if let requestTimeout = networkPreferences?.timeoutIntervalForRequest {
+                let requestTimeOutMs = requestTimeout * 1_000
+                configuration.connectTimeoutMs = UInt32(requestTimeOutMs)
+            }
+
+            if let maxRetryUnwrapped = networkPreferences?.maxRetryCount {
+                configuration.retryStrategyOptions = RetryStrategyOptions(maxRetriesBase:  Int(maxRetryUnwrapped))
+            }
+
             return CognitoIdentityClient(config: configuration)
         default:
             fatalError()
@@ -129,6 +147,15 @@ extension AWSCognitoAuthPlugin {
     private func makeURLSession() -> URLSession {
         let configuration = URLSessionConfiguration.default
         configuration.urlCache = nil
+
+        if let timeoutIntervalForRequest = networkPreferences?.timeoutIntervalForRequest {
+            configuration.timeoutIntervalForRequest = timeoutIntervalForRequest
+        }
+
+        if let timeoutIntervalForResource = networkPreferences?.timeoutIntervalForResource {
+            configuration.timeoutIntervalForResource = timeoutIntervalForResource
+        }
+
         return URLSession(configuration: configuration)
     }
 

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/AWSCognitoAuthPlugin.swift

@@ -32,6 +32,9 @@ public final class AWSCognitoAuthPlugin: AWSCognitoAuthPluginBehavior {
 
     var httpClientEngineProxy: HttpClientEngineProxy?
 
+    /// The user network preferences for timeout and retry
+    let networkPreferences: AWSCognitoNetworkPreferences?
+
     @_spi(InternalAmplifyConfiguration)
     internal(set) public var jsonConfiguration: JSONValue?
 
@@ -42,5 +45,13 @@ public final class AWSCognitoAuthPlugin: AWSCognitoAuthPluginBehavior {
 
     /// Instantiates an instance of the AWSCognitoAuthPlugin.
     public init() {
+        self.networkPreferences = nil
+    }
+
+    /// Instantiates an instance of the AWSCognitoAuthPlugin with custom network preferences
+    /// - Parameters:
+    ///   - networkPreferences: network preferences
+    public init(networkPreferences: AWSCognitoNetworkPreferences) {
+        self.networkPreferences = networkPreferences
     }
 }

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Models/AWSCognitoNetworkPreferences.swift

@@ -0,0 +1,31 @@
+//
+// Copyright Amazon.com Inc. or its affiliates.
+// All Rights Reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+import Foundation
+
+public struct AWSCognitoNetworkPreferences {
+
+    /// The maximum number of retries for failed requests.
+    public let maxRetryCount: UInt32
+
+    /// The timeout interval to use when waiting for additional data.
+    public let timeoutIntervalForRequest: TimeInterval
+
+    /// The maximum amount of time that a resource request should be allowed to take.
+    /// 
+    /// NOTE: This value is only applicable to HostedUI because the underlying Swift SDK does
+    /// not support resource timeouts
+    public let timeoutIntervalForResource: TimeInterval?
+
+    public init(maxRetryCount: UInt32,
+                timeoutIntervalForRequest: TimeInterval,
+                timeoutIntervalForResource: TimeInterval? = nil) {
+        self.maxRetryCount = maxRetryCount
+        self.timeoutIntervalForRequest = timeoutIntervalForRequest
+        self.timeoutIntervalForResource = timeoutIntervalForResource
+    }
+}

AmplifyPlugins/Auth/Tests/AWSCognitoAuthPluginUnitTests/ConfigurationTests/AWSCognitoAuthPluginConfigTests.swift

@@ -190,4 +190,50 @@ class AWSCognitoAuthPluginConfigTests: XCTestCase {
         }
     }
 
+    /// Test Auth configuration with valid config for user pool and identity pool, with network preferences
+    ///
+    /// - Given: Given valid config for user pool and identity pool, and network preferences
+    /// - When:
+    ///    - I configure auth with the given configuration and network preferences
+    /// - Then:
+    ///    - I should not get any error while configuring auth
+    ///
+    func testConfigWithUserPoolAndIdentityPoolWithNetworkPreferences() throws {
+        let plugin = AWSCognitoAuthPlugin(
+            networkPreferences: .init(
+                maxRetryCount: 2,
+                timeoutIntervalForRequest: 60,
+                timeoutIntervalForResource: 60))
+        try Amplify.add(plugin: plugin)
+
+        let categoryConfig = AuthCategoryConfiguration(plugins: [
+            "awsCognitoAuthPlugin": [
+                "CredentialsProvider": ["CognitoIdentity": ["Default":
+                                                                ["PoolId": "xx",
+                                                                 "Region": "us-east-1"]
+                                                           ]],
+                "CognitoUserPool": ["Default": [
+                    "PoolId": "xx",
+                    "Region": "us-east-1",
+                    "AppClientId": "xx",
+                    "AppClientSecret": "xx"]]
+            ]
+        ])
+        let amplifyConfig = AmplifyConfiguration(auth: categoryConfig)
+        do {
+            try Amplify.configure(amplifyConfig)
+            
+            let escapeHatch = plugin.getEscapeHatch()
+            guard case .userPoolAndIdentityPool(let userPoolClient, let identityPoolClient) = escapeHatch else {
+                XCTFail("Expected .userPool, got \(escapeHatch)")
+                return
+            }
+            XCTAssertNotNil(userPoolClient)
+            XCTAssertNotNil(identityPoolClient)
+
+        } catch {
+            XCTFail("Should not throw error. \(error)")
+        }
+    }
+
 }