fix(datastore): using URLProtocol monitor auth request headers (#3221)

* fix(datastore): using URLProtocol monitor multiAuth request headers

* test(datastore): update integration test xcode version

* test(datastore): update integration test for auth IAM

* remove redundant semicolons

Author: Di Wu

.github/workflows/integ_test_datastore_auth_iam.yml

@@ -5,12 +5,12 @@ on:
 
 permissions:
     id-token: write
-    contents: read 
+    contents: read
 
 jobs:
   datastore-integration-auth-iam-test-iOS:
     timeout-minutes: 30
-    runs-on: macos-12
+    runs-on: macos-13
     environment: IntegrationTest
     steps:
       - uses: actions/checkout@2541b1294d2704b0964813337f33b291d3f8596b
@@ -33,6 +33,8 @@ jobs:
         with:
           project_path: ./AmplifyPlugins/DataStore/Tests/DataStoreHostApp
           scheme: AWSDataStorePluginAuthIAMTests
+          destination: 'platform=iOS Simulator,name=iPhone 14,OS=latest'
+          xcode_path: '/Applications/Xcode_14.3.app'
 
   datastore-integration-auth-iam-test-tvOS:
     timeout-minutes: 30

.github/workflows/integ_test_datastore_multi_auth.yml

@@ -10,7 +10,7 @@ permissions:
 jobs:
   datastore-integration-multi-auth-test-iOS:
     timeout-minutes: 30
-    runs-on: macos-12
+    runs-on: macos-13
     environment: IntegrationTest
     steps:
       - uses: actions/checkout@2541b1294d2704b0964813337f33b291d3f8596b
@@ -33,6 +33,8 @@ jobs:
         with:
           project_path: ./AmplifyPlugins/DataStore/Tests/DataStoreHostApp
           scheme: AWSDataStorePluginMultiAuthTests
+          destination: 'platform=iOS Simulator,name=iPhone 14,OS=latest'
+          xcode_path: '/Applications/Xcode_14.3.app'
 
   datastore-integration-multi-auth-test-tvOS:
     timeout-minutes: 30

AmplifyPlugins/DataStore/Tests/DataStoreHostApp/AWSDataStorePluginAuthIAMTests/AWSDataStoreAuthBaseTest.swift

@@ -23,46 +23,77 @@ struct TestUser {
     let password: String
 }
 
-class AuthRecorderInterceptor: URLRequestInterceptor {
-    let awsAuthService: AWSAuthService = AWSAuthService()
-    var consumedAuthTypes: Set<AWSAuthorizationType> = []
-    private let accessQueue = DispatchQueue(label: "com.amazon.AuthRecorderInterceptor.consumedAuthTypes")
-
-    private func recordAuthType(_ authType: AWSAuthorizationType) {
-        accessQueue.async {
-            self.consumedAuthTypes.insert(authType)
-        }
-    }
+class DataStoreAuthBaseTestURLSessionFactory: URLSessionBehaviorFactory {
+    static let testIdHeaderKey = "x-amplify-test"
 
-    func intercept(_ request: URLRequest) throws -> URLRequest {
-        guard let headers = request.allHTTPHeaderFields else {
-            fatalError("No headers found in request \(request)")
-        }
+    static let subject = PassthroughSubject<(String, Set<AWSAuthorizationType>), Never>()
 
-        let authHeaderValue = headers["Authorization"]
-        let apiKeyHeaderValue = headers["x-api-key"]
+    class Sniffer: URLProtocol {
 
-        if apiKeyHeaderValue != nil {
-            recordAuthType(.apiKey)
-        }
+        override class func canInit(with request: URLRequest) -> Bool {
+            guard let headers = request.allHTTPHeaderFields else {
+                fatalError("No headers found in request \(request)")
+            }
+
+            guard let testId = headers[DataStoreAuthBaseTestURLSessionFactory.testIdHeaderKey] else {
+                return false
+            }
+
+            var result: Set<AWSAuthorizationType> = []
+            let authHeaderValue = headers["Authorization"]
+            let apiKeyHeaderValue = headers["x-api-key"]
+
+            if apiKeyHeaderValue != nil {
+                result.insert(.apiKey)
+            }
+
+            if let authHeaderValue = authHeaderValue,
+               case let .success(claims) = AWSAuthService().getTokenClaims(tokenString: authHeaderValue),
+               let cognitoIss = claims["iss"] as? String, cognitoIss.contains("cognito") {
+                result.insert(.amazonCognitoUserPools)
+            }
 
-        if let authHeaderValue = authHeaderValue,
-           case let .success(claims) = awsAuthService.getTokenClaims(tokenString: authHeaderValue),
-           let cognitoIss = claims["iss"] as? String, cognitoIss.contains("cognito") {
-            recordAuthType(.amazonCognitoUserPools)
+            if let authHeaderValue = authHeaderValue,
+               authHeaderValue.starts(with: "AWS4-HMAC-SHA256") {
+                result.insert(.awsIAM)
+            }
+
+            DataStoreAuthBaseTestURLSessionFactory.subject.send((testId, result))
+            return false
         }
 
-        if let authHeaderValue = authHeaderValue,
-           authHeaderValue.starts(with: "AWS4-HMAC-SHA256") {
-            recordAuthType(.awsIAM)
+    }
+
+    class Interceptor: URLRequestInterceptor {
+        let testId: String?
+
+        init(testId: String?) {
+            self.testId = testId
         }
 
-        return request
+        func intercept(_ request: URLRequest) async throws -> URLRequest {
+            if let testId {
+                var mutableRequest = request
+                mutableRequest.setValue(testId, forHTTPHeaderField: DataStoreAuthBaseTestURLSessionFactory.testIdHeaderKey)
+                return mutableRequest
+            }
+            return request
+        }
     }
 
-    func reset() {
-        consumedAuthTypes = []
+    func makeSession(withDelegate delegate: URLSessionBehaviorDelegate?) -> URLSessionBehavior {
+        let urlSessionDelegate = delegate?.asURLSessionDelegate
+        let configuration = URLSessionConfiguration.default
+        configuration.tlsMinimumSupportedProtocolVersion = .TLSv12
+        configuration.tlsMaximumSupportedProtocolVersion = .TLSv13
+        configuration.protocolClasses?.insert(Sniffer.self, at: 0)
+
+        let session = URLSession(configuration: configuration,
+                                 delegate: urlSessionDelegate,
+                                 delegateQueue: nil)
+        return AmplifyURLSession(session: session)
     }
+
 }
 
 class AWSDataStoreAuthBaseTest: XCTestCase {
@@ -71,7 +102,6 @@ class AWSDataStoreAuthBaseTest: XCTestCase {
     var amplifyConfig: AmplifyConfiguration!
     var user1: TestUser?
     var user2: TestUser?
-    var authRecorderInterceptor: AuthRecorderInterceptor!
 
     override func setUp() {
         continueAfterFailure = false
@@ -138,8 +168,6 @@ class AWSDataStoreAuthBaseTest: XCTestCase {
             self.user1 = TestUser(username: user1, password: passwordUser1)
             self.user2 = TestUser(username: user2, password: passwordUser2)
 
-            authRecorderInterceptor = AuthRecorderInterceptor()
-
             amplifyConfig = try TestConfigHelper.retrieveAmplifyConfiguration(forResource: configFile)
 
         } catch {
@@ -161,7 +189,8 @@ class AWSDataStoreAuthBaseTest: XCTestCase {
     func setup(
         withModels models: AmplifyModelRegistration,
         testType: DataStoreAuthTestType,
-        apiPluginFactory: () -> AWSAPIPlugin = { AWSAPIPlugin(sessionFactory: AmplifyURLSessionFactory()) }
+        testId: String? = nil,
+        apiPluginFactory: () -> AWSAPIPlugin = { AWSAPIPlugin(sessionFactory: DataStoreAuthBaseTestURLSessionFactory()) }
     ) async {
         do {
             setupCredentials(forAuthStrategy: testType)
@@ -182,7 +211,10 @@ class AWSDataStoreAuthBaseTest: XCTestCase {
 
             // register auth recorder interceptor
             let apiName = try apiEndpointName()
-            try apiPlugin.add(interceptor: authRecorderInterceptor, for: apiName)
+            try apiPlugin.add(
+                interceptor: DataStoreAuthBaseTestURLSessionFactory.Interceptor(testId: testId),
+                for: apiName
+            )
 
             await signOut()
         } catch {
@@ -486,13 +518,27 @@ extension AWSDataStoreAuthBaseTest {
         await waitForExpectations([expectations.mutationDelete, expectations.mutationDeleteProcessed], timeout: 60)
     }
 
-    func assertUsedAuthTypes(_ authTypes: [AWSAuthorizationType],
-                             file: StaticString = #file,
-                             line: UInt = #line) {
-        XCTAssertEqual(authRecorderInterceptor.consumedAuthTypes,
-                       Set(authTypes),
-                       file: file,
-                       line: line)
+    func assertUsedAuthTypes(
+        testId: String,
+        authTypes: [AWSAuthorizationType],
+        file: StaticString = #file,
+        line: UInt = #line
+    ) -> XCTestExpectation {
+        let expectation = expectation(description: "Should have expected auth types")
+        expectation.assertForOverFulfill = false
+        DataStoreAuthBaseTestURLSessionFactory.subject
+        .filter { $0.0 == testId }
+        .map { $0.1 }
+        .collect(.byTime(DispatchQueue.global(), .milliseconds(3500)))
+        .sink {
+            let result = $0.reduce(Set<AWSAuthorizationType>()) { partialResult, data in
+                partialResult.union(data)
+            }
+            XCTAssertEqual(result, Set(authTypes), file: file, line: line)
+            expectation.fulfill()
+        }
+        .store(in: &requests)
+        return expectation
     }
 }
 

AmplifyPlugins/DataStore/Tests/DataStoreHostApp/AWSDataStorePluginAuthIAMTests/AWSDataStoreCategoryPluginIAMAuthIntegrationTests.swift

@@ -16,15 +16,19 @@ class AWSDataStoreCategoryPluginIAMAuthIntegrationTests: AWSDataStoreAuthBaseTes
     /// Then: DataStore is successfully initialized, query returns a result,
     ///      mutation is processed for authenticated users
     func testIAMAllowPrivate() async {
+        let testId = UUID().uuidString
         await setup(withModels: IAMPrivateModelRegistration(),
-              testType: .defaultAuthIAM)
+                    testType: .defaultAuthIAM,
+                    testId: testId)
 
         await signIn(user: user1)
 
         let expectations = makeExpectations()
 
         await assertDataStoreReady(expectations)
 
+        let authTypeExpectation = assertUsedAuthTypes(testId: testId, authTypes: [.awsIAM])
+
         // Query
         await assertQuerySuccess(modelType: TodoIAMPrivate.self,
                            expectations) { error in
@@ -38,21 +42,25 @@ class AWSDataStoreCategoryPluginIAMAuthIntegrationTests: AWSDataStoreAuthBaseTes
             XCTFail("Error mutation \(error)")
         }
 
-        assertUsedAuthTypes([.awsIAM])
+        await fulfillment(of: [authTypeExpectation], timeout: 5)
     }
 
     /// Given: a guest user,  a model with `allow public` auth rule with IAM as provider
     /// When: DataStore query/mutation operations are sent with IAM
     /// Then: DataStore is successfully initialized, query returns a result,
     ///      mutation is processed for unauthenticated users
-    func testIAMAllowPublic() async{
+    func testIAMAllowPublic() async {
+        let testId = UUID().uuidString
         await setup(withModels: IAMPublicModelRegistration(),
-              testType: .defaultAuthIAM)
+                    testType: .defaultAuthIAM,
+                    testId: testId)
 
         let expectations = makeExpectations()
 
         await assertDataStoreReady(expectations)
 
+        let authTypeExpectation = assertUsedAuthTypes(testId: testId, authTypes: [.awsIAM])
+
         // Query
         await assertQuerySuccess(modelType: TodoIAMPublic.self,
                            expectations) { error in
@@ -66,7 +74,7 @@ class AWSDataStoreCategoryPluginIAMAuthIntegrationTests: AWSDataStoreAuthBaseTes
             XCTFail("Error mutation \(error)")
         }
 
-        assertUsedAuthTypes([.awsIAM])
+        await fulfillment(of: [authTypeExpectation], timeout: 5)
     }
 }