dcmaw-16581-diff.patch

diff --git a/Sources/Extensions/CustomTypes/LoginSessionConfiguration+OneLogin.swift b/Sources/Extensions/CustomTypes/LoginSessionConfiguration+OneLogin.swift
index e8272e1..85cb231 100644
--- a/Sources/Extensions/CustomTypes/LoginSessionConfiguration+OneLogin.swift
+++ b/Sources/Extensions/CustomTypes/LoginSessionConfiguration+OneLogin.swift
@@ -21,10 +21,8 @@ extension LoginSessionConfiguration {
             locale: env.isLocaleWelsh ? .cy : .en,
             persistentSessionId: persistentSessionID,
             tokenHeaders: shouldAttestIntegrity ?
-            try await Task {
-                try await OneLoginAppIntegrityService()
-                    .integrityAssertions(integrityService)
-            }.value : nil
+            try await OneLoginAppIntegrityService(integrityService: integrityService)
+            .integrityAssertions(): nil
         )
     }
 }
diff --git a/Sources/Extensions/URLRequest+OneLogin.swift b/Sources/Extensions/URLRequest+OneLogin.swift
index fedf0f5..0118967 100644
--- a/Sources/Extensions/URLRequest+OneLogin.swift
+++ b/Sources/Extensions/URLRequest+OneLogin.swift
@@ -9,13 +9,9 @@ extension URLRequest {
         var request = URLRequest(url: AppEnvironment.stsToken)
         
         request.asXWWWFormURLEncoded()
-        
-        let headers = try await Task {
-            try await OneLoginAppIntegrityService()
-                .integrityAssertions(appIntegrityProvider)
-        }.value
-        
-        for (key, value) in headers {
+                
+        for (key, value) in try await OneLoginAppIntegrityService(integrityService: appIntegrityProvider)
+            .integrityAssertions() {
             request.setValue(
                 value,
                 forHTTPHeaderField: key
diff --git a/Sources/Utilities/OneLoginAppIntegrityService.swift b/Sources/Utilities/OneLoginAppIntegrityService.swift
index 8eb9988..43851a4 100644
--- a/Sources/Utilities/OneLoginAppIntegrityService.swift
+++ b/Sources/Utilities/OneLoginAppIntegrityService.swift
@@ -1,41 +1,93 @@
 import AppIntegrity
+import Combine
+import Dispatch
 
-class OneLoginAppIntegrityService {
+
+/// Use this type to automatically retry calls to `AppIntegrityProvider/integrityAssertions` every time an error is thrown.
+///
+/// e.g. the following code will make two (2) attempts at fetching the integrity assertions before giving up and throwing an error.
+/// ```
+///     let service = OneLoginAppIntegrityService(integrityService: integrityService)
+///     service.integrityAssertions(attempts: 2)
+/// ```
+///
+/// After that point the instance to `OneLoginAppIntegrityService` should be discarded.
+/// Should you wish to make another number of attempts, you **must** create a new `OneLoginAppIntegrityService` instance.
+///
+/// - seealso: ``integrityAssertions(attempts:)`` on making an attempt at fetching the integrity assertions
+actor OneLoginAppIntegrityService {
     private(set) var errorRetries = 0
+    private let integrityService: AppIntegrityProvider
     
-    func integrityAssertions(
-        _ integrityService: AppIntegrityProvider
-    ) async throws -> [String: String] {
+    init(integrityService: AppIntegrityProvider) {
+        self.integrityService = integrityService
+    }
+
+    private func integrityAssertions(after interval: DispatchTimeInterval) async throws -> [String: String] {
+        try await withCheckedThrowingContinuation { continuation in
+            DispatchQueue.global().asyncAfter(deadline: .now() + interval) {
+                Task {
+                    do {
+                        continuation.resume(returning: try await self.integrityService.integrityAssertions)
+                    } catch {
+                        continuation.resume(throwing: error)
+                    }
+                }
+            }
+        }        
+    }
+    
+    private func attempt(_ retry: Int = 0, max maxAttempts: Int) async throws -> [String: String] {
         do {
-            return try await integrityService.integrityAssertions
+            return try await self.integrityAssertions(after: .milliseconds(100 * retry))
         } catch let error as FirebaseAppCheckError {
             switch error.kind {
             case .network:
-                return try await retryIntegrityAssertions(error, integrityService)
+                self.errorRetries = retry + 1
+                if self.errorRetries == maxAttempts {
+                    throw error
+                }
+                return try await self.attempt(self.errorRetries, max: maxAttempts)
             case .unknown, .invalidConfiguration, .keychainAccess, .notSupported, .generic:
                 throw error
             }
         } catch let error as ClientAssertionError {
             switch error.kind {
             case .invalidToken, .serverError, .cantDecodeClientAssertion:
-                return try await retryIntegrityAssertions(error, integrityService)
+                self.errorRetries = retry + 1
+                if self.errorRetries == maxAttempts {
+                    throw error
+                }
+                return try await self.attempt(self.errorRetries, max: maxAttempts)
             case .invalidPublicKey:
                 throw error
             }
         }
     }
     
-    private func retryIntegrityAssertions(
-        _ error: Error,
-        _ integrityService: AppIntegrityProvider
-    ) async throws -> [String: String] {
-        errorRetries += 1
-        
-        guard errorRetries < 3 else {
-            throw error
-        }
-        
-        try await Task.sleep(nanoseconds: 100_000_000 * UInt64(errorRetries))
-        return try await integrityAssertions(integrityService)
+    /// This method will attempt to return the integrity assertions as returned by calling `AppIntegrityProvider/integrityAssertions` on the `AppIntegrityProvider` used to consturct this instance.
+    ///
+    /// For specific [1] errors thrown, the method will try again, up to a maximum the number of attempts given, before giving up and throwing the error as thrown by `AppIntegrityProvider/integrityAssertions`
+    ///
+    /// [1]: Only the following errors will be retried:
+    /// - `FirebaseAppCheckError/network`
+    /// - `ClientAssertionError/invalidToken`
+    /// - `ClientAssertionError/serverError`
+    /// - `ClientAssertionError/cantDecodeClientAssertion`
+    ///
+    /// The first attempt is performed now; each follow-up attempt will incure an aditional 100ms delay per attempt over time.
+    ///     e.g.
+    ///     - attempt 1  → now
+    ///     - attempt 2  → 100ms
+    ///     - attempt 3  → 200ms
+    ///     - attempt 4  → 300ms
+    ///
+    /// - parameter attempts: the number of attempts before giving up; **default** is 3.
+    /// - returns the integrity assertions as returned by `AppIntegrityProvider/integrityAssertions`
+    /// - throws the error as thrown by `AppIntegrityProvider/integrityAssertions` on the **last**  attempt.
+    /// - remark: this function is not deisnged to be called in parallel, in which case its behaviour is undefined. No strong guarantees are provided when making concurrent calls to this function, e.g. 2 parallel calls may both lead to an attempt now.
+    ///
+    public func integrityAssertions(attempts maxAttempts: Int = 3) async throws -> [String: String] {
+        return try await self.attempt(max: maxAttempts)
     }
 }
diff --git a/Tests/UnitTests/Mocks/Sessions+Services/MockAppIntegrityProvider.swift b/Tests/UnitTests/Mocks/Sessions+Services/MockAppIntegrityProvider.swift
index 360faf7..e44b355 100644
--- a/Tests/UnitTests/Mocks/Sessions+Services/MockAppIntegrityProvider.swift
+++ b/Tests/UnitTests/Mocks/Sessions+Services/MockAppIntegrityProvider.swift
@@ -1,10 +1,12 @@
 import AppIntegrity
 
 final class MockAppIntegrityProvider: AppIntegrityProvider {
+    var attempts = 0
     var errorThrownAssertingIntegrity: Error?
     
     var integrityAssertions: [String: String] {
         get throws {
+            self.attempts += 1
             if let errorThrownAssertingIntegrity {
                 throw errorThrownAssertingIntegrity
             }
diff --git a/Tests/UnitTests/Utilities/OneLoginAppIntegrityServiceTests.swift b/Tests/UnitTests/Utilities/OneLoginAppIntegrityServiceTests.swift
index f681bc5..6af6f4d 100644
--- a/Tests/UnitTests/Utilities/OneLoginAppIntegrityServiceTests.swift
+++ b/Tests/UnitTests/Utilities/OneLoginAppIntegrityServiceTests.swift
@@ -3,134 +3,145 @@ import AppIntegrity
 import Testing
 
 struct OneLoginAppIntegrityServiceTests {
-    let mockInterityService = MockAppIntegrityProvider()
     
     @Test("Integrity assertions are retried for network error")
-    func integrityAssertionsNetworkError() async throws {
+    func integrityAssertionNetworkError() async throws {
+        let mockInterityService = MockAppIntegrityProvider()
         mockInterityService.errorThrownAssertingIntegrity = FirebaseAppCheckError(.network)
-        let sut = OneLoginAppIntegrityService()
+        let sut = OneLoginAppIntegrityService(integrityService: mockInterityService)
         
         do {
-            _ = try await sut.integrityAssertions(mockInterityService)
+            _ = try await sut.integrityAssertions()
         } catch let error as FirebaseAppCheckError {
             #expect(error.kind == .network)
-            #expect(sut.errorRetries == 3)
+            #expect(await sut.errorRetries == 3)
+            #expect(mockInterityService.attempts == 3)
         }
     }
     
     @Test("Integrity assertions are not retried for unknown error")
     func integrityAssertionsUnknownError() async throws {
+        let mockInterityService = MockAppIntegrityProvider()
         mockInterityService.errorThrownAssertingIntegrity = FirebaseAppCheckError(.unknown)
-        let sut = OneLoginAppIntegrityService()
+        let sut = OneLoginAppIntegrityService(integrityService: mockInterityService)
         
         do {
-            _ = try await sut.integrityAssertions(mockInterityService)
+            _ = try await sut.integrityAssertions()
         } catch let error as FirebaseAppCheckError {
             #expect(error.kind == .unknown)
-            #expect(sut.errorRetries == 0)
+            #expect(await sut.errorRetries == 0)
+            #expect(mockInterityService.attempts == 1)
         }
     }
     
     @Test("Integrity assertions are not retried for invalid configuration error")
     func integrityAssertionsInvalidConfigurationError() async throws {
+        let mockInterityService = MockAppIntegrityProvider()
         mockInterityService.errorThrownAssertingIntegrity = FirebaseAppCheckError(.invalidConfiguration)
-        let sut = OneLoginAppIntegrityService()
+        let sut = OneLoginAppIntegrityService(integrityService: mockInterityService)
         
         do {
-            _ = try await sut.integrityAssertions(mockInterityService)
+            _ = try await sut.integrityAssertions()
         } catch let error as FirebaseAppCheckError {
             #expect(error.kind == .invalidConfiguration)
-            #expect(sut.errorRetries == 0)
+            #expect(await sut.errorRetries == 0)
         }
     }
     
     @Test("Integrity assertions are not retried for keychain access error")
     func integrityAssertionsKeychainAccessError() async throws {
+        let mockInterityService = MockAppIntegrityProvider()
         mockInterityService.errorThrownAssertingIntegrity = FirebaseAppCheckError(.keychainAccess)
-        let sut = OneLoginAppIntegrityService()
+        let sut = OneLoginAppIntegrityService(integrityService: mockInterityService)
         
         do {
-            _ = try await sut.integrityAssertions(mockInterityService)
+            _ = try await sut.integrityAssertions()
         } catch let error as FirebaseAppCheckError {
             #expect(error.kind == .keychainAccess)
-            #expect(sut.errorRetries == 0)
+            #expect(await sut.errorRetries == 0)
         }
     }
     
     @Test("Integrity assertions are not retried for keychain access error")
     func integrityAssertionsNotSupportedError() async throws {
+        let mockInterityService = MockAppIntegrityProvider()
         mockInterityService.errorThrownAssertingIntegrity = FirebaseAppCheckError(.notSupported)
-        let sut = OneLoginAppIntegrityService()
+        let sut = OneLoginAppIntegrityService(integrityService: mockInterityService)
         
         do {
-            _ = try await sut.integrityAssertions(mockInterityService)
+            _ = try await sut.integrityAssertions()
         } catch let error as FirebaseAppCheckError {
             #expect(error.kind == .notSupported)
-            #expect(sut.errorRetries == 0)
+            #expect(await sut.errorRetries == 0)
         }
     }
     
     @Test("Integrity assertions are not retried for generic error")
     func integrityAssertionsGenericError() async throws {
+        let mockInterityService = MockAppIntegrityProvider()
         mockInterityService.errorThrownAssertingIntegrity = FirebaseAppCheckError(.generic)
-        let sut = OneLoginAppIntegrityService()
+        let sut = OneLoginAppIntegrityService(integrityService: mockInterityService)
         
         do {
-            _ = try await sut.integrityAssertions(mockInterityService)
+            _ = try await sut.integrityAssertions()
         } catch let error as FirebaseAppCheckError {
             #expect(error.kind == .generic)
-            #expect(sut.errorRetries == 0)
+            #expect(await sut.errorRetries == 0)
         }
     }
     
     @Test("Integrity assertions are retried for invalid token error")
-    func integrityAssertionsInvalidTokenError() async throws {
+    func integrityAssertionInvalidTokenError() async throws {
+        let mockInterityService = MockAppIntegrityProvider()
         mockInterityService.errorThrownAssertingIntegrity = ClientAssertionError(.invalidToken)
-        let sut = OneLoginAppIntegrityService()
+        let sut = OneLoginAppIntegrityService(integrityService: mockInterityService)
         
         do {
-            _ = try await sut.integrityAssertions(mockInterityService)
+            _ = try await sut.integrityAssertions()
         } catch let error as ClientAssertionError {
             #expect(error.kind == .invalidToken)
-            #expect(sut.errorRetries == 3)
+            #expect(await sut.errorRetries == 3)
         }
     }
     
     @Test("Integrity assertions are retried for server error")
-    func integrityAssertionsServerError() async throws {
+    func integrityAssertionServerError() async throws {
+        let mockInterityService = MockAppIntegrityProvider()
         mockInterityService.errorThrownAssertingIntegrity = ClientAssertionError(.serverError)
-        let sut = OneLoginAppIntegrityService()
+        let sut = OneLoginAppIntegrityService(integrityService: mockInterityService)
         
         do {
-            _ = try await sut.integrityAssertions(mockInterityService)
+            _ = try await sut.integrityAssertions()
         } catch let error as ClientAssertionError {
             #expect(error.kind == .serverError)
-            #expect(sut.errorRetries == 3)
+            #expect(await sut.errorRetries == 3)
         }
     }
     
     @Test("Integrity assertions are retried for cant decode client assertion error")
-    func integrityAssertionsCantDecodeClientAssertionError() async throws {
+    func integrityAssertionCantDecodeClientAssertionError() async throws {
+        let mockInterityService = MockAppIntegrityProvider()
         mockInterityService.errorThrownAssertingIntegrity = ClientAssertionError(.cantDecodeClientAssertion)
-        let sut = OneLoginAppIntegrityService()
+        let sut = OneLoginAppIntegrityService(integrityService: mockInterityService)
         
         do {
-            _ = try await sut.integrityAssertions(mockInterityService)
+            _ = try await sut.integrityAssertions()
         } catch let error as ClientAssertionError {
             #expect(error.kind == .cantDecodeClientAssertion)
-            #expect(sut.errorRetries == 3)
+            #expect(await sut.errorRetries == 3)
         }
     }
     
     @Test("Integrity assertions are not retried for any other ClientAssertionError")
     func integrityAssertionsClientAssertionError() async throws {
+        let mockInterityService = MockAppIntegrityProvider()
         mockInterityService.errorThrownAssertingIntegrity = ClientAssertionError(.invalidPublicKey)
-        let sut = OneLoginAppIntegrityService()
+        let sut = OneLoginAppIntegrityService(integrityService: mockInterityService)
         
         await #expect(
             throws: ClientAssertionError(.invalidPublicKey)
         ) {
-            try await sut.integrityAssertions(mockInterityService)
+            try await sut.integrityAssertions()
         }
     }
 }