[CIS-224] Implement token refresh

Sources/StreamChat/ChatClient.swift

@@ -186,6 +186,7 @@ public class _ChatClient<ExtraData: ExtraDataTypes> {
     }()
 
     private(set) lazy var internetConnection = environment.internetConnection()
+    private(set) lazy var clientUpdater = environment.clientUpdaterBuilder(self)
 
     /// The environment object containing all dependencies of this `Client` instance.
     private let environment: Environment
@@ -288,8 +289,7 @@ public class _ChatClient<ExtraData: ExtraDataTypes> {
 
         currentUserId = fetchCurrentUserIdFromDatabase()
 
-        let updater = environment.clientUpdaterBuilder(self)
-        updater.reloadUserIfNeeded(completion: completion)
+        clientUpdater.reloadUserIfNeeded(completion: completion)
     }
 
     deinit {
@@ -410,7 +410,17 @@ extension ClientError {
 extension _ChatClient: ConnectionStateDelegate {
     func webSocketClient(_ client: WebSocketClient, didUpdateConectionState state: WebSocketConnectionState) {
         connectionStatus = .init(webSocketConnectionState: state)
-
+
+        // Reaquire a new token if current one has expired
+        if case let .disconnected(error) = state,
+           let error = error,
+           error.isTokenExpiredError {
+            clientUpdater.reloadUserIfNeeded()
+        }
+        norifyConnectionIdWaitersIfNeeded(for: state)
+    }
+
+    private func norifyConnectionIdWaitersIfNeeded(for state: WebSocketConnectionState) {
         _connectionId.mutate { mutableConnectionId in
             _connectionIdWaiters.mutate { connectionIdWaiters in
 
@@ -422,7 +432,13 @@ extension _ChatClient: ConnectionStateDelegate {
                 } else {
                     mutableConnectionId = nil
 
-                    if case .disconnected = state {
+                    if case let .disconnected(error) = state {
+                        // Do not notify waiters in case of token expired error
+                        if let error = error,
+                           error.isTokenExpiredError {
+                            return
+                        }
+
                         // No reconnection attempt schedule, we should fail all existing connectionId waiters.
                         connectionIdWaiters.forEach { $0(nil) }
                         connectionIdWaiters.removeAll()
@@ -433,6 +449,16 @@ extension _ChatClient: ConnectionStateDelegate {
     }
 }
 
+private extension ClientError {
+    var isTokenExpiredError: Bool {
+        if let error = underlyingError as? ErrorPayload,
+           ErrorPayload.tokenInvadlidErrorCodes ~= error.code {
+            return true
+        }
+        return false
+    }
+}
+
 /// `Client` provides connection details for the `RequestEncoder`s it creates.
 extension _ChatClient: ConnectionDetailsProviderDelegate {
     func provideToken(completion: @escaping (_ token: Token?) -> Void) {

Sources/StreamChat/ChatClient_Tests.swift

@@ -392,6 +392,46 @@ class ChatClient_Tests: StressTestCase {
         XCTAssertNil(providedConnectionId)
     }
 
+    func test_client_webSocketIsDisconnected_becauseTokenExpired_callsReloadUserIfNeeded() {
+        // Create a new chat client
+        let client = ChatClient(
+            config: inMemoryStorageConfig,
+            tokenProvider: .anonymous,
+            workerBuilders: workerBuilders,
+            eventWorkerBuilders: [],
+            environment: testEnv.environment
+        )
+
+        // Simulate access to `webSocketClient` so it is initialized
+        _ = client.webSocketClient
+
+        // Set a connection Id waiter and set `providedConnectionId` to a non-nil value
+        var providedConnectionId: ConnectionId? = .unique
+        client.provideConnectionId {
+            providedConnectionId = $0
+        }
+        XCTAssertNotNil(providedConnectionId)
+
+        // Was called on ChatClient init
+        XCTAssertEqual(testEnv.clientUpdater!.reloadUserIfNeeded_callsCount, 1)
+
+        // Simulate WebSocketConnection change to "disconnected"
+        let error = ClientError(with: ErrorPayload(code: 40, message: "", statusCode: 200))
+        testEnv.webSocketClient?
+            .connectionStateDelegate?
+            .webSocketClient(
+                testEnv.webSocketClient!,
+                didUpdateConectionState: .disconnected(error: error)
+            )
+
+        // Was called one more time on receiving token expired error
+        XCTAssertEqual(testEnv.clientUpdater!.reloadUserIfNeeded_callsCount, 2)
+
+        // Assert the provided connection id is not `nil`, because we don't reset connectionId
+        // on receiving token expiration error
+        XCTAssertNotNil(providedConnectionId)
+    }
+
     // MARK: - APIClient tests
 
     func test_apiClientConfiguration() throws {

Sources/StreamChatTestTools/ChatClientUpdater_Mock.swift

@@ -9,7 +9,13 @@ class ChatClientUpdaterMock<ExtraData: ExtraDataTypes>: ChatClientUpdater<ExtraD
     @Atomic var prepareEnvironment_newToken: Token?
     var prepareEnvironment_called: Bool { prepareEnvironment_newToken != nil }
 
-    @Atomic var reloadUserIfNeeded_called = false
+    @Atomic var reloadUserIfNeeded_called = false {
+        didSet {
+            reloadUserIfNeeded_callsCount += 1
+        }
+    }
+
+    var reloadUserIfNeeded_callsCount = 0
     @Atomic var reloadUserIfNeeded_completion: ((Error?) -> Void)?
     @Atomic var reloadUserIfNeeded_callSuper: (() -> Void)?
 
@@ -47,6 +53,7 @@ class ChatClientUpdaterMock<ExtraData: ExtraDataTypes>: ChatClientUpdater<ExtraD
         prepareEnvironment_newToken = nil
 
         reloadUserIfNeeded_called = false
+        reloadUserIfNeeded_callsCount = 0
         reloadUserIfNeeded_completion = nil
         reloadUserIfNeeded_callSuper = nil