From f819c07ea4748987e4cf1dbe5f8d70ae4db58e61 Mon Sep 17 00:00:00 2001
From: Dmitry Gachkovsky <sacethi0@gmail.com>
Date: Tue, 11 May 2021 10:03:09 +0300
Subject: [PATCH] [CIS-224] Implement token refresh

---
 Sources/StreamChat/ChatClient.swift           | 34 ++++++++++++++--
 Sources/StreamChat/ChatClient_Tests.swift     | 40 +++++++++++++++++++
 .../ChatClientUpdater_Mock.swift              |  9 ++++-
 3 files changed, 78 insertions(+), 5 deletions(-)

diff --git a/Sources/StreamChat/ChatClient.swift b/Sources/StreamChat/ChatClient.swift
index 67dca20160c..4796064cd84 100644
--- a/Sources/StreamChat/ChatClient.swift
+++ b/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) {
diff --git a/Sources/StreamChat/ChatClient_Tests.swift b/Sources/StreamChat/ChatClient_Tests.swift
index b164f03f286..de7af46cf8f 100644
--- a/Sources/StreamChat/ChatClient_Tests.swift
+++ b/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 {
diff --git a/Sources/StreamChatTestTools/ChatClientUpdater_Mock.swift b/Sources/StreamChatTestTools/ChatClientUpdater_Mock.swift
index cecd98000b1..63542d4ebfe 100644
--- a/Sources/StreamChatTestTools/ChatClientUpdater_Mock.swift
+++ b/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