Safer connectUser() by logging out the user first if needed (#2577)

* Safer connectUser() by logging out the user first if needed

* Clear active channels on log out

* Update connect guest user in E2E Test to use a different ID when connecting a guest user

* Add comment back

* Add test coverage

* Update CHANGELOG.md

* Add test coverage to clearing out active controllers when logging out

* Update CHANGELOG.md

* Improve authentication flow to log out when needed

* Add EnvironmentState tests

* Add tests for prepareEnvironment assertion on id mismatch

* Update AuthenticationRepository tests

* Fix log in flow test error

---------

Co-authored-by: Pol Quintana <[email protected]>

Author: nuno-vieira

CHANGELOG.md

@@ -15,6 +15,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 - Fix crash when accessing FetchCache with an unexecuted NSFetchRequest [#2572](https://github.com/GetStream/stream-chat-swift/pull/2572)
 - Fix an issue which was blocking a Guest Authentication operation to retrieve a connection token [#2574](https://github.com/GetStream/stream-chat-swift/pull/2574)
 - Make connect/disconnect safer when network is offline [#2571](https://github.com/GetStream/stream-chat-swift/pull/2571)
+- Make connect safer by logging out the user first if needed [#2577](https://github.com/GetStream/stream-chat-swift/pull/2577)
 
 ## StreamChatUI
 ### ✅ Added

Sources/StreamChat/ChatClient.swift

@@ -369,9 +369,7 @@ public class ChatClient {
     /// Connects anonymous user
     /// - Parameter completion: The completion that will be called once the **first** user session for the given token is setup.
     public func connectAnonymousUser(completion: ((Error?) -> Void)? = nil) {
-        authenticationRepository.connectUser(
-            userInfo: nil,
-            tokenProvider: { $0(.success(.anonymous)) },
+        authenticationRepository.connectAnonymousUser(
             completion: { completion?($0) }
         )
     }
@@ -403,6 +401,10 @@ public class ChatClient {
     public func logout(completion: @escaping () -> Void) {
         authenticationRepository.logOutUser()
 
+        // Stop tracking active components
+        activeChannelControllers.removeAllObjects()
+        activeChannelListControllers.removeAllObjects()
+
         let group = DispatchGroup()
         group.enter()
         disconnect {
@@ -473,21 +475,8 @@ public class ChatClient {
 }
 
 extension ChatClient: AuthenticationRepositoryDelegate {
-    /// Clears state related to the current user to leave the client ready for another user
-    /// Will clear:
-    ///     - Background workers
-    ///     - References to active controllers
-    ///     - Database
-    /// - Parameter completion: A block to be executed when the process is completed. Contains an error if something went wrong
-    func clearCurrentUserData(completion: @escaping (Error?) -> Void) {
-        createBackgroundWorkers()
-
-        // Stop tracking active components
-        activeChannelControllers.removeAllObjects()
-        activeChannelListControllers.removeAllObjects()
-
-        // Reset all existing local data.
-        databaseContainer.removeAllData(force: true, completion: completion)
+    func logOutUser(completion: @escaping () -> Void) {
+        logout(completion: completion)
     }
 
     func didFinishSettingUpAuthenticationEnvironment(for state: EnvironmentState) {

Sources/StreamChat/Repositories/AuthenticationRepository.swift

@@ -10,11 +10,21 @@ enum EnvironmentState {
     case firstConnection
     case newToken
     case newUser
+
+    init(currentUserId: UserId?, newUserId: UserId) {
+        if currentUserId == nil {
+            self = .firstConnection
+        } else if currentUserId == newUserId {
+            self = .newToken
+        } else {
+            self = .newUser
+        }
+    }
 }
 
 protocol AuthenticationRepositoryDelegate: AnyObject {
     func didFinishSettingUpAuthenticationEnvironment(for state: EnvironmentState)
-    func clearCurrentUserData(completion: @escaping (Error?) -> Void)
+    func logOutUser(completion: @escaping () -> Void)
 }
 
 class AuthenticationRepository {
@@ -130,26 +140,57 @@ class AuthenticationRepository {
         updateToken(token: token, notifyTokenWaiters: completeTokenWaiters)
     }
 
-    /// Establishes a connection for a  user.
+    /// Establishes a connection for a non anonymous user.
     /// - Parameters:
     ///   - userInfo:       The user information that will be created OR updated if it exists.
     ///   - tokenProvider:  The block to be used to get a token.
-    func connectUser(userInfo: UserInfo?, tokenProvider: @escaping TokenProvider, completion: @escaping (Error?) -> Void) {
-        self.tokenProvider = tokenProvider
-        scheduleTokenFetch(isRetry: false, userInfo: userInfo, tokenProvider: tokenProvider, completion: completion)
+    func connectUser(userInfo: UserInfo, tokenProvider: @escaping TokenProvider, completion: @escaping (Error?) -> Void) {
+        var logOutFirst: Bool {
+            if let currentUserId = currentUserId, currentUserId.isGuest {
+                return true
+            }
+
+            let state = EnvironmentState(currentUserId: currentUserId, newUserId: userInfo.id)
+            return state == .newUser
+        }
+
+        executeTokenFetch(logOutFirst: logOutFirst, userInfo: userInfo, tokenProvider: tokenProvider, completion: completion)
     }
 
     /// Establishes a connection for a guest user.
     /// - Parameters:
     ///   - userInfo: The user information that will be created OR updated if it exists.
     func connectGuestUser(userInfo: UserInfo, completion: @escaping (Error?) -> Void) {
-        connectUser(
-            userInfo: userInfo,
-            tokenProvider: { [weak self] completion in
-                self?.fetchGuestToken(userInfo: userInfo, completion: completion)
-            },
-            completion: completion
-        )
+        let tokenProvider: TokenProvider = { [weak self] completion in
+            self?.fetchGuestToken(userInfo: userInfo, completion: completion)
+        }
+        executeTokenFetch(logOutFirst: true, userInfo: userInfo, tokenProvider: tokenProvider, completion: completion)
+    }
+
+    /// Establishes a connection for an anonymous user.
+    func connectAnonymousUser(completion: @escaping (Error?) -> Void) {
+        let tokenProvider: TokenProvider = { $0(.success(.anonymous)) }
+        executeTokenFetch(logOutFirst: true, userInfo: nil, tokenProvider: tokenProvider, completion: completion)
+    }
+
+    private func executeTokenFetch(logOutFirst: Bool, userInfo: UserInfo?, tokenProvider: @escaping TokenProvider, completion: @escaping (Error?) -> Void) {
+        log.assert(delegate != nil, "Delegate should not be nil at this point")
+
+        let handleTokenFetch = { [weak self] in
+            self?.tokenProvider = tokenProvider
+            self?.scheduleTokenFetch(isRetry: false, userInfo: userInfo, tokenProvider: tokenProvider, completion: completion)
+        }
+
+        guard logOutFirst else {
+            handleTokenFetch()
+            return
+        }
+
+        if let delegate = delegate {
+            delegate.logOutUser(completion: handleTokenFetch)
+        } else {
+            handleTokenFetch()
+        }
     }
 
     func clearTokenProvider() {
@@ -176,42 +217,30 @@ class AuthenticationRepository {
 
     func prepareEnvironment(
         userInfo: UserInfo?,
-        newToken: Token,
-        completion: @escaping (Error?) -> Void
+        newToken: Token
     ) {
-        let state: EnvironmentState
-        if currentUserId == nil {
-            state = .firstConnection
-        } else if newToken.userId == currentUserId {
-            state = .newToken
-        } else {
-            state = .newUser
-        }
+        let state = EnvironmentState(currentUserId: currentUserId, newUserId: newToken.userId)
 
         log.assert(delegate != nil, "Delegate should not be nil at this point")
 
+        if let userInfo = userInfo, !newToken.userId.isGuest {
+            log.assert(
+                userInfo.id == newToken.userId,
+                "The id of the retrieved token should match the user information passed to connect"
+            )
+        }
+
         switch state {
         case .firstConnection, .newToken:
             connectionRepository.updateWebSocketEndpoint(with: newToken, userInfo: userInfo)
             setToken(token: newToken, completeTokenWaiters: true)
             delegate?.didFinishSettingUpAuthenticationEnvironment(for: state)
-            completion(nil)
 
         case .newUser:
             completeTokenWaiters(token: nil)
+            connectionRepository.updateWebSocketEndpoint(with: newToken, userInfo: userInfo)
             setToken(token: newToken, completeTokenWaiters: false)
             delegate?.didFinishSettingUpAuthenticationEnvironment(for: state)
-
-            // Setting a new connection is not possible in connectionless mode.
-            guard connectionRepository.isClientInActiveMode else {
-                completion(ClientError.ClientIsNotInActiveMode())
-                return
-            }
-
-            connectionRepository.disconnect(source: .userInitiated) { [weak delegate, weak connectionRepository] in
-                connectionRepository?.updateWebSocketEndpoint(with: newToken, userInfo: userInfo)
-                delegate?.clearCurrentUserData(completion: completion)
-            }
         }
     }
 
@@ -294,19 +323,12 @@ class AuthenticationRepository {
         }
 
         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
-                if let error = error {
-                    onCompletion(error)
-                    return
-                }
-
-                // We manually change the `connectionStatus` for passive client
-                // to `disconnected` when environment was prepared correctly
-                // (e.g. current user session is successfully restored).
-                connectionRepository?.forceConnectionStatusForInactiveModeIfNeeded()
-                connectionRepository?.connect(userInfo: userInfo, completion: onCompletion)
-            }
+            self?.prepareEnvironment(userInfo: userInfo, newToken: token)
+            // We manually change the `connectionStatus` for passive client
+            // to `disconnected` when environment was prepared correctly
+            // (e.g. current user session is successfully restored).
+            connectionRepository?.forceConnectionStatusForInactiveModeIfNeeded()
+            connectionRepository?.connect(userInfo: userInfo, completion: onCompletion)
         }
 
         let retryFetchIfPossible: (Error?) -> Void = { [weak self] error in
@@ -377,3 +399,9 @@ extension ClientError {
         }
     }
 }
+
+private extension UserId {
+    var isGuest: Bool {
+        hasPrefix(UserRole.guest.rawValue)
+    }
+}

StreamChatUITestsApp/ViewController.swift

@@ -205,10 +205,8 @@ extension StreamChatWrapper {
     }
 
     func connectGuestUser(completion: @escaping (Error?) -> Void) {
-        let userCredentials = UserCredentials.default
-        let tokenProvider = mockTokenProvider(for: userCredentials)
         client?.connectGuestUser(
-            userInfo: userCredentials.userInfo,
+            userInfo: .init(id: "123"),
             completion: completion
         )
     }

TestTools/StreamChatTestTools/Mocks/StreamChat/Repositories/AuthenticationRepository_Mock.swift

@@ -10,6 +10,7 @@ class AuthenticationRepository_Mock: AuthenticationRepository, Spy {
     enum Signature {
         static let connectTokenProvider = "connectUser(userInfo:tokenProvider:completion:)"
         static let connectGuest = "connectGuestUser(userInfo:completion:)"
+        static let connectAnon = "connectAnonymousUser(completion:)"
         static let refreshToken = "refreshToken(completion:)"
         static let clearTokenProvider = "clearTokenProvider()"
         static let logOut = "logOutUser()"
@@ -24,6 +25,7 @@ class AuthenticationRepository_Mock: AuthenticationRepository, Spy {
 
     var connectUserResult: Result<Void, Error>?
     var connectGuestResult: Result<Void, Error>?
+    var connectAnonResult: Result<Void, Error>?
     var refreshTokenResult: Result<Void, Error>?
     var completeWaitersToken: Token?
 
@@ -73,6 +75,13 @@ class AuthenticationRepository_Mock: AuthenticationRepository, Spy {
         }
     }
 
+    override func connectAnonymousUser(completion: @escaping (Error?) -> Void) {
+        record()
+        if let result = connectAnonResult {
+            completion(result.error)
+        }
+    }
+
     override func setToken(token: Token, completeTokenWaiters: Bool) {
         record()
         setMockToken(token)

TestTools/StreamChatTestTools/SpyPattern/Spy/Logger_Spy.swift

@@ -8,6 +8,7 @@ import Foundation
 final class Logger_Spy: Logger, Spy {
     var originalLogger: Logger?
     var recordedFunctions: [String] = []
+    var failedAsserts: Int = 0
 
     func injectMock() {
         let logger = LogConfig.logger
@@ -39,6 +40,7 @@ final class Logger_Spy: Logger, Spy {
         lineNumber: UInt = #line
     ) {
         record()
+        failedAsserts += condition() ? 0 : 1
     }
 
     override func assertionFailure(
@@ -49,5 +51,6 @@ final class Logger_Spy: Logger, Spy {
         lineNumber: UInt = #line
     ) {
         record()
+        failedAsserts += 1
     }
 }