[Internal] Add retries to Authentication Repository (#2414)

* Wait until connect completes to proceed with tests * Fix bugs in Sinatra and update navBar accId * Replace status with halt for better logs * Fix * Add retries mechanism to Authentication Repository * Show channel list while connecting * Update specs * Update CHANGELOG * Address PR comments Co-authored-by: Pol Quintana <pol.quintana1@gmail.com>
GetStream · Dec 16, 2022 · 2f0cb97 · 2f0cb97
1 parent d0d3489
commit 2f0cb97
Show file tree

Hide file tree

Showing 14 changed files with 174 additions and 121 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -11,6 +11,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 - Add support for hiding connection status with `isInvisible` [#2373](https://github.com/GetStream/stream-chat-swift/pull/2373)
 - Add `.withAttachments` in `MessageSearchFilterScope` to filter messages with attachments only [#2417](https://github.com/GetStream/stream-chat-swift/pull/2417)
 - Add `.withoutAttachments` in `MessageSearchFilterScope` to filter messages without any attachments [#2417](https://github.com/GetStream/stream-chat-swift/pull/2417)
+- Add retries mechanism to AuthenticationRepository [#2414](https://github.com/GetStream/stream-chat-swift/pull/2414)
 
 ### 🐞 Fixed
 - Fix connecting user with non-expiring tokens (ex: development token) [#2393](https://github.com/GetStream/stream-chat-swift/pull/2393)

diff --git a/Sources/StreamChat/APIClient/APIClient.swift b/Sources/StreamChat/APIClient/APIClient.swift
@@ -14,7 +14,7 @@ class APIClient {
 
     /// `APIClient` uses this object to decode the results of network requests.
     let decoder: RequestDecoder
-    
+
     /// Used for reobtaining tokens when they expire and API client receives token expiration error
     let tokenRefresher: (@escaping () -> Void) -> Void
 
@@ -45,12 +45,6 @@ class APIClient {
 
     /// Shows whether the token is being refreshed at the moment
     @Atomic private var isRefreshingToken: Bool = false
-
-    /// Amount of consecutive token refresh attempts
-    @Atomic private var tokenRefreshConsecutiveFailures: Int = 0
-
-    /// Maximum amount of consecutive token refresh attempts before failing
-    let maximumTokenRefreshAttempts = 10
 
     /// Maximum amount of times a request can be retried
     private let maximumRequestRetries = 3
@@ -166,11 +160,6 @@ class APIClient {
         endpoint: Endpoint<Response>,
         completion: @escaping (Result<Response, Error>) -> Void
     ) {
-        if tokenRefreshConsecutiveFailures > maximumTokenRefreshAttempts {
-            completion(.failure(ClientError.TooManyTokenRefreshAttempts()))
-            return
-        }
-
         guard !isRefreshingToken else {
             completion(.failure(ClientError.RefreshingToken()))
             return
@@ -207,7 +196,6 @@ class APIClient {
                         response: response,
                         error: error
                     )
-                    self.tokenRefreshConsecutiveFailures = 0
                     completion(.success(decodedResponse))
                 } catch {
                     if error is ClientError.ExpiredToken == false {
@@ -236,18 +224,11 @@ class APIClient {
             completion(ClientError.RefreshingToken())
             return
         }
-        isRefreshingToken = true
 
-        // We stop the queue so no more operations are triggered during the refresh
-        operationQueue.isSuspended = true
-
-        // Increase the amount of consecutive failures
-        _tokenRefreshConsecutiveFailures.mutate { $0 += 1 }
+        enterTokenFetchMode()
 
         tokenRefresher { [weak self] in
-            self?.isRefreshingToken = false
-            // We restart the queue now that token refresh is completed
-            self?.operationQueue.isSuspended = false
+            self?.exitTokenFetchMode()
             completion(ClientError.TokenRefreshed())
         }
     }
@@ -299,6 +280,18 @@ class APIClient {
         isInRecoveryMode = false
         operationQueue.isSuspended = false
     }
+
+    func enterTokenFetchMode() {
+        // We stop the queue so no more operations are triggered during the refresh
+        isRefreshingToken = true
+        operationQueue.isSuspended = true
+    }
+
+    func exitTokenFetchMode() {
+        // We restart the queue now that token refresh is completed
+        isRefreshingToken = false
+        operationQueue.isSuspended = false
+    }
 }
 
 extension URLRequest {

diff --git a/Sources/StreamChat/APIClient/RequestDecoder.swift b/Sources/StreamChat/APIClient/RequestDecoder.swift
@@ -90,12 +90,6 @@ extension ClientError {
     class RefreshingToken: ClientError {}
     class TokenRefreshed: ClientError {}
     class ConnectionError: ClientError {}
-    class TooManyTokenRefreshAttempts: ClientError {
-        override var localizedDescription: String {
-            "Authentication failed on expired tokens after too many refresh attempts, please check that your user tokens are created correctly."
-        }
-    }
-
     class ResponseBodyEmpty: ClientError {
         override var localizedDescription: String { "Response body is empty." }
     }

diff --git a/Sources/StreamChat/Repositories/AuthenticationRepository.swift b/Sources/StreamChat/Repositories/AuthenticationRepository.swift
@@ -18,8 +18,20 @@ protocol AuthenticationRepositoryDelegate: AnyObject {
 }
 
 class AuthenticationRepository {
+    private enum Constants {
+        /// Maximum amount of consecutive token refresh attempts before failing
+        static let maximumTokenRefreshAttempts = 10
+    }
+
     private let tokenQueue: DispatchQueue = DispatchQueue(label: "io.getstream.auth-repository", attributes: .concurrent)
-    private var _isGettingToken: Bool = false
+    private var _isGettingToken: Bool = false {
+        didSet {
+            guard oldValue != _isGettingToken else { return }
+            _isGettingToken ? apiClient.enterTokenFetchMode() : apiClient.exitTokenFetchMode()
+        }
+    }
+
+    private var _consecutiveRefreshFailures: Int = 0
     private var _currentUserId: UserId?
     private var _currentToken: Token?
     private var _tokenProvider: TokenProvider?
@@ -31,6 +43,11 @@ class AuthenticationRepository {
         set { tokenQueue.async(flags: .barrier) { self._isGettingToken = newValue }}
     }
 
+    private var consecutiveRefreshFailures: Int {
+        get { tokenQueue.sync { _consecutiveRefreshFailures } }
+        set { tokenQueue.async(flags: .barrier) { self._consecutiveRefreshFailures = newValue }}
+    }
+
     private(set) var currentUserId: UserId? {
         get { tokenQueue.sync { _currentUserId } }
         set { tokenQueue.async(flags: .barrier) { self._currentUserId = newValue }}
@@ -65,8 +82,6 @@ class AuthenticationRepository {
     private let apiClient: APIClient
     private let databaseContainer: DatabaseContainer
     private let connectionRepository: ConnectionRepository
-    /// A timer that runs token refreshing job
-    private var tokenRetryTimer: TimerControl?
     /// Retry timing strategy for refreshing an expired token
     private var tokenExpirationRetryStrategy: RetryStrategy
     private let timerType: Timer.Type
@@ -121,7 +136,7 @@ class AuthenticationRepository {
     ///   - tokenProvider:  The block to be used to get a token.
     func connectUser(userInfo: UserInfo?, tokenProvider: @escaping TokenProvider, completion: @escaping (Error?) -> Void) {
         self.tokenProvider = tokenProvider
-        getToken(userInfo: userInfo, tokenProvider: tokenProvider, completion: completion)
+        scheduleTokenFetch(isRetry: false, userInfo: userInfo, tokenProvider: tokenProvider, completion: completion)
     }
 
     /// Establishes a connection for a guest user.
@@ -156,16 +171,7 @@ class AuthenticationRepository {
             return
         }
 
-        let tokenProviderCheckingSuccess: TokenProvider = { [weak self] completion in
-            tokenProvider { result in
-                if case .success = result {
-                    self?.tokenExpirationRetryStrategy.resetConsecutiveFailures()
-                }
-                completion(result)
-            }
-        }
-
-        scheduleTokenFetch(userInfo: nil, tokenProvider: tokenProviderCheckingSuccess, completion: completion)
+        scheduleTokenFetch(isRetry: false, userInfo: nil, tokenProvider: tokenProvider, completion: completion)
     }
 
     func prepareEnvironment(
@@ -240,24 +246,24 @@ class AuthenticationRepository {
         waiters.forEach { $0.value(token) }
     }
 
-    private func scheduleTokenFetch(userInfo: UserInfo?, tokenProvider: @escaping TokenProvider, completion: @escaping (Error?) -> Void) {
-        guard !isGettingToken else {
+    private func scheduleTokenFetch(isRetry: Bool, userInfo: UserInfo?, tokenProvider: @escaping TokenProvider, completion: @escaping (Error?) -> Void) {
+        guard !isGettingToken || isRetry else {
             tokenRequestCompletions.append(completion)
             return
         }
 
-        tokenRetryTimer = timerType.schedule(
+        timerType.schedule(
             timeInterval: tokenExpirationRetryStrategy.getDelayAfterTheFailure(),
             queue: .main
         ) { [weak self] in
             log.debug("Firing timer for a new token request", subsystems: .authentication)
-            self?.getToken(userInfo: nil, tokenProvider: tokenProvider, completion: completion)
+            self?.getToken(isRetry: isRetry, userInfo: nil, tokenProvider: tokenProvider, completion: completion)
         }
     }
 
-    private func getToken(userInfo: UserInfo?, tokenProvider: @escaping TokenProvider, completion: @escaping (Error?) -> Void) {
+    private func getToken(isRetry: Bool, userInfo: UserInfo?, tokenProvider: @escaping TokenProvider, completion: @escaping (Error?) -> Void) {
         tokenRequestCompletions.append(completion)
-        guard !isGettingToken else {
+        guard !isGettingToken || isRetry else {
             log.debug("Trying to get a token while already getting one", subsystems: .authentication)
             return
         }
@@ -276,11 +282,17 @@ class AuthenticationRepository {
                 self._isGettingToken = false
                 let completions = self._tokenRequestCompletions
                 self._tokenRequestCompletions = []
+                self._consecutiveRefreshFailures = 0
                 return completions
             }
             completionBlocks?.forEach { $0(error) }
         }
 
+        guard consecutiveRefreshFailures < Constants.maximumTokenRefreshAttempts else {
+            onCompletion(ClientError.TooManyFailedTokenRefreshAttempts())
+            return
+        }
+
         let onTokenReceived: (Token) -> Void = { [weak self, weak connectionRepository] token in
             self?.prepareEnvironment(userInfo: userInfo, newToken: token) { error in
                 // Errors thrown during `prepareEnvironment` cannot be recovered
@@ -297,13 +309,29 @@ class AuthenticationRepository {
             }
         }
 
+        let retryFetchIfPossible: (Error?) -> Void = { [weak self] error in
+            guard let self = self else { return }
+            self.consecutiveRefreshFailures += 1
+            guard self.consecutiveRefreshFailures < Constants.maximumTokenRefreshAttempts else {
+                onCompletion(error ?? ClientError.TooManyFailedTokenRefreshAttempts())
+                return
+            }
+
+            // We don't need to pass the completion again, as it is already present in `tokenRequestCompletions`
+            self.scheduleTokenFetch(isRetry: true, userInfo: userInfo, tokenProvider: tokenProvider, completion: { _ in })
+        }
+
         log.debug("Requesting a new token", subsystems: .authentication)
-        tokenProvider { result in
+        tokenProvider { [weak self] result in
             switch result {
-            case let .success(newToken):
+            case let .success(newToken) where !newToken.isExpired:
                 onTokenReceived(newToken)
+                self?.tokenExpirationRetryStrategy.resetConsecutiveFailures()
+            case .success:
+                retryFetchIfPossible(nil)
             case let .failure(error):
-                onCompletion(error)
+                log.info("Failed fetching token with error: \(error)")
+                retryFetchIfPossible(error)
             }
         }
     }
@@ -334,3 +362,14 @@ class AuthenticationRepository {
         tokenWaiters[waiter] = nil
     }
 }
+
+extension ClientError {
+    public class TooManyFailedTokenRefreshAttempts: ClientError {
+        override public var localizedDescription: String {
+            """
+                Token fetch has failed more than 10 times.
+                Please make sure that your `tokenProvider` is correctly functioning.
+            """
+        }
+    }
+}
diff --git a/Sources/StreamChat/Repositories/ConnectionRepository.swift b/Sources/StreamChat/Repositories/ConnectionRepository.swift
@@ -149,14 +149,14 @@ class ConnectionRepository {
         case let .connected(connectionId: id):
             shouldNotifyConnectionIdWaiters = true
             connectionId = id
-        case let .disconnected(source):
-            if let error = source.serverError,
-               error.isInvalidTokenError {
-                onInvalidToken()
-                shouldNotifyConnectionIdWaiters = false
-            } else {
-                shouldNotifyConnectionIdWaiters = true
-            }
+
+        case let .disconnecting(source) where source.serverError?.isInvalidTokenError == true,
+             let .disconnected(source) where source.serverError?.isInvalidTokenError == true:
+            onInvalidToken()
+            shouldNotifyConnectionIdWaiters = false
+            connectionId = nil
+        case .disconnected:
+            shouldNotifyConnectionIdWaiters = true
             connectionId = nil
         case .initialized,
              .connecting,

diff --git a/StreamChatUITestsApp/StreamChat/ChannelList.swift b/StreamChatUITestsApp/StreamChat/ChannelList.swift
@@ -48,5 +48,6 @@ final class ChannelList: ChatChannelListVC, ChatConnectionControllerDelegate {
         case .disconnected:
             title = "disconnected"
         }
+        navigationItem.titleView?.accessibilityIdentifier = title
     }
 }
diff --git a/StreamChatUITestsApp/ViewController.swift b/StreamChatUITestsApp/ViewController.swift
@@ -36,12 +36,15 @@ final class ViewController: UIViewController {
     @objc func didTap() {
         // Setup chat client
         streamChat.setUpChat()
-        streamChat.connectUser(completion: { _ in})
+        streamChat.connectUser(completion: { _ in })
+        showChannelList()
+    }
 
+    private func showChannelList() {
         // create UI
         let channelList = streamChat.makeChannelListViewController()
         router = channelList.router as? CustomChannelListRouter
-        
+
         // create connection switch if needed
         let switchControl = self.createIsConnectedSwitchIfNeeded()
 

diff --git a/StreamChatUITestsAppUITests/Pages/ChannelListPage.swift b/StreamChatUITestsAppUITests/Pages/ChannelListPage.swift
@@ -23,12 +23,8 @@ enum ChannelListPage {
             format: "identifier LIKE 'titleLabel' AND label LIKE '\(withName)'")).firstMatch
     }
 
-    static var connectionStatus: XCUIElement {
-        if ProcessInfo().operatingSystemVersion.majorVersion == 12 {
-            return app.navigationBars.otherElements.firstMatch
-        } else {
-            return app.navigationBars.staticTexts.firstMatch
-        }
+    static func connectionLabel(withStatus: ChannelListPage.ConnectionStatus) -> XCUIElement {
+        app.navigationBars.matching(NSPredicate(format: "identifier LIKE '\(withStatus.rawValue)'")).firstMatch
     }
 
     enum ConnectionStatus: String {

diff --git a/StreamChatUITestsAppUITests/Robots/UserRobot+Asserts.swift b/StreamChatUITestsAppUITests/Robots/UserRobot+Asserts.swift
@@ -151,9 +151,8 @@ extension UserRobot {
         file: StaticString = #filePath,
         line: UInt = #line
     ) -> Self {
-        let expectedStatus = status.rawValue
-        let actualStatus = ChannelListPage.connectionStatus.waitForText(expectedStatus, timeout: timeout).text
-        XCTAssertEqual(actualStatus, expectedStatus, file: file, line: line)
+        let correctStatus = ChannelListPage.connectionLabel(withStatus: status).wait(timeout: timeout).exists
+        XCTAssertTrue(correctStatus, file: file, line: line)
         return self
     }
 }

diff --git a/StreamChatUITestsAppUITests/Tests/Authentication_Tests.swift b/StreamChatUITestsAppUITests/Tests/Authentication_Tests.swift
@@ -11,7 +11,6 @@ final class Authentication_Tests: StreamTestCase {
         mockServerEnabled = false
         app.setLaunchArguments(.jwt)
         try super.setUpWithError()
-        throw XCTSkip("[CIS-2309] JWT authentication issues")
     }
 
     func test_tokenExpiriesBeforeUserLogsIn() {
@@ -102,7 +101,7 @@ final class Authentication_Tests: StreamTestCase {
         AND("server returns an error") {}
         AND("JWT generation recovers on server side") {}
         THEN("app requests a token refresh a second time") {
-            userRobot.assertConnectionStatus(.connected, timeout: 30)
+            userRobot.assertConnectionStatus(.connected)
         }
     }
 }