Handle OIDC cancellation and workaround missing progress when linking a new device. (#4935)

* Fix a bug where the link device flow wasn't dismissed when complete.

* Listen for cancellation of the WAS when linking a device (dismissing the QR screen when it happens).

* Add CustomStringConvertible conformances to QRCodeLoginScreen actions.

ElementX/Sources/FlowCoordinators/AuthenticationFlowCoordinator.swift

@@ -322,8 +322,8 @@ class AuthenticationFlowCoordinator: FlowCoordinatorProtocol {
                 DispatchQueue.main.async {
                     self.stateMachine.tryEvent(.signedIn, userInfo: userSession)
                 }
-            case .requestOIDCAuthorisation:
+            case .requestOIDCAuthorisation, .linkedDevice:
-                fatalError("QR code login shouldn't request an OIDC flow.")
+                fatalError("QR code login shouldn't request an OIDC flow or link a device.")
             }
         }
         .store(in: &cancellables)

ElementX/Sources/FlowCoordinators/LinkNewDeviceFlowCoordinator.swift

@@ -9,7 +9,7 @@ import Combine
 import Foundation
 
 enum LinkNewDeviceFlowCoordinatorAction {
-    case requestOIDCAuthorisation(URL)
+    case requestOIDCAuthorisation(URL, OIDCAccountSettingsPresenter.Continuation)
     case dismiss
 }
 
@@ -76,8 +76,10 @@ class LinkNewDeviceFlowCoordinator: FlowCoordinatorProtocol {
                     fatalError("QR linking shouldn't send sign-in actions.")
                 case .dismiss:
                     navigationStackCoordinator.pop()
-                case .requestOIDCAuthorisation(let url):
+                case .requestOIDCAuthorisation(let url, let continuation):
-                    actionsSubject.send(.requestOIDCAuthorisation(url))
+                    actionsSubject.send(.requestOIDCAuthorisation(url, continuation))
+                case .linkedDevice:
+                    actionsSubject.send(.dismiss)
                 }
             }
             .store(in: &cancellables)

ElementX/Sources/FlowCoordinators/SettingsFlowCoordinator.swift

@@ -182,8 +182,8 @@ class SettingsFlowCoordinator: FlowCoordinatorProtocol {
                 switch action {
                 case .dismiss:
                     navigationStackCoordinator.setSheetCoordinator(nil)
-                case .requestOIDCAuthorisation(let url):
+                case .requestOIDCAuthorisation(let url, let continuation):
-                    presentAccountManagementURL(url)
+                    presentAccountManagementURL(url, continuation: continuation)
                 }
             }
             .store(in: &cancellables)
@@ -292,12 +292,13 @@ class SettingsFlowCoordinator: FlowCoordinatorProtocol {
     // MARK: OIDC Account Management
     
     private var accountSettingsPresenter: OIDCAccountSettingsPresenter?
-    private func presentAccountManagementURL(_ url: URL) {
+    private func presentAccountManagementURL(_ url: URL, continuation: OIDCAccountSettingsPresenter.Continuation? = nil) {
         // Note to anyone in the future if you come back here to make this open in Safari instead of a WAS.
         // As of iOS 16, there is an issue on the simulator with accessing the cookie but it works on a device. 🤷‍♂️
         accountSettingsPresenter = OIDCAccountSettingsPresenter(accountURL: url,
                                                                 presentationAnchor: flowParameters.windowManager.mainWindow,
-                                                                appSettings: flowParameters.appSettings)
+                                                                appSettings: flowParameters.appSettings,
+                                                                continuation: continuation)
         accountSettingsPresenter?.start()
     }
 }

ElementX/Sources/Screens/Authentication/OIDCAuthenticationPresenter.swift

@@ -50,11 +50,7 @@ class OIDCAuthenticationPresenter: NSObject {
         
         guard let url else {
             // Check for user cancellation to avoid showing an alert in that instance.
-            if let nsError = error as? NSError,
+            if error?.isOIDCUserCancellation == true {
-               nsError.domain == ASWebAuthenticationSessionErrorDomain,
-               nsError.code == ASWebAuthenticationSessionError.canceledLogin.rawValue,
-               // If there's a failure reason then the cancellation wasn't made by the user.
-               nsError.localizedFailureReason == nil {
                 // No need to show an error here, just abort and return a failure.
                 await authenticationService.abortOIDCLogin(data: oidcData)
                 return .failure(.oidcError(.userCancellation))
@@ -121,3 +117,20 @@ extension ASWebAuthenticationSession.Callback {
         }
     }
 }
+
+// MARK: - Helpers
+
+extension Error {
+    var isOIDCUserCancellation: Bool {
+        let nsError = self as NSError
+        
+        if nsError.domain == ASWebAuthenticationSessionErrorDomain,
+           nsError.code == ASWebAuthenticationSessionError.canceledLogin.rawValue,
+           // If there's a failure reason then the cancellation wasn't made by the user.
+           nsError.localizedFailureReason == nil {
+            return true
+        }
+        
+        return false
+    }
+}

ElementX/Sources/Screens/QRCodeLoginScreen/QRCodeLoginScreenCoordinator.swift

@@ -16,11 +16,22 @@ struct QRCodeLoginScreenCoordinatorParameters {
     let appMediator: AppMediatorProtocol
 }
 
-enum QRCodeLoginScreenCoordinatorAction {
+enum QRCodeLoginScreenCoordinatorAction: CustomStringConvertible {
     case dismiss
     case signInManually
     case signedIn(userSession: UserSessionProtocol)
-    case requestOIDCAuthorisation(URL)
+    case requestOIDCAuthorisation(URL, OIDCAccountSettingsPresenter.Continuation)
+    case linkedDevice
+    
+    var description: String {
+        switch self {
+        case .dismiss: "dismiss"
+        case .signInManually: "signInManually"
+        case .signedIn: "signedIn"
+        case .requestOIDCAuthorisation: "requestOIDCAuthorisation"
+        case .linkedDevice: "linkedDevice"
+        }
+    }
 }
 
 final class QRCodeLoginScreenCoordinator: CoordinatorProtocol {
@@ -53,8 +64,10 @@ final class QRCodeLoginScreenCoordinator: CoordinatorProtocol {
                 actionsSubject.send(.dismiss)
             case .signedIn(let userSession):
                 actionsSubject.send(.signedIn(userSession: userSession))
-            case .requestOIDCAuthorisation(let url):
+            case .requestOIDCAuthorisation(let url, let continuation):
-                actionsSubject.send(.requestOIDCAuthorisation(url))
+                actionsSubject.send(.requestOIDCAuthorisation(url, continuation))
+            case .linkedDevice:
+                actionsSubject.send(.linkedDevice)
             }
         }
         .store(in: &cancellables)

ElementX/Sources/Screens/QRCodeLoginScreen/QRCodeLoginScreenModels.swift

@@ -8,11 +8,22 @@
 
 import SwiftUI
 
-enum QRCodeLoginScreenViewModelAction {
+enum QRCodeLoginScreenViewModelAction: CustomStringConvertible {
     case dismiss
     case signInManually
     case signedIn(userSession: UserSessionProtocol)
-    case requestOIDCAuthorisation(URL)
+    case requestOIDCAuthorisation(URL, OIDCAccountSettingsPresenter.Continuation)
+    case linkedDevice
+    
+    var description: String {
+        switch self {
+        case .dismiss: "dismiss"
+        case .signInManually: "signInManually"
+        case .signedIn: "signedIn"
+        case .requestOIDCAuthorisation: "requestOIDCAuthorisation"
+        case .linkedDevice: "linkedDevice"
+        }
+    }
 }
 
 enum QRCodeLoginScreenMode {

ElementX/Sources/Screens/QRCodeLoginScreen/QRCodeLoginScreenViewModel.swift

@@ -21,6 +21,7 @@ class QRCodeLoginScreenViewModel: QRCodeLoginScreenViewModelType, QRCodeLoginScr
     }
     
     private var currentTask: AnyCancellable?
+    private var oidcResultTask: AnyCancellable?
     
     init(mode: QRCodeLoginScreenMode,
          canSignInManually: Bool,
@@ -186,12 +187,12 @@ class QRCodeLoginScreenViewModel: QRCodeLoginScreenViewModelType, QRCodeLoginScr
                 case .establishingSecureChannel(let checkCodeString):
                     state.state = .displayCode(.deviceCode(checkCodeString))
                 case .waitingForAuthorisation(let url):
-                    actionsSubject.send(.requestOIDCAuthorisation(url))
+                    requestOIDCAuthorization(url: url)
                 case .syncingSecrets:
                     break // Nothing to do.
                 case .done:
                     MXLog.info("Link with QR code completed.")
-                    actionsSubject.send(.dismiss)
+                    actionsSubject.send(.linkedDevice)
                 }
             }
     }
@@ -226,12 +227,14 @@ class QRCodeLoginScreenViewModel: QRCodeLoginScreenViewModelType, QRCodeLoginScr
                 case .qrScanned(let checkCodeSender):
                     state.state = .confirmCode(.inputCode(checkCodeSender))
                 case .waitingForAuthorisation(let url):
-                    actionsSubject.send(.requestOIDCAuthorisation(url))
+                    requestOIDCAuthorization(url: url)
                 case .syncingSecrets:
-                    break // Nothing to do.
+                    // break // Nothing to do.
-                case .done:
+                    // .done is rarely received at the moment, so lets consider linking to be done here.
                     MXLog.info("Link with QR code completed.")
-                    actionsSubject.send(.dismiss)
+                    actionsSubject.send(.linkedDevice)
+                case .done:
+                    break // Not necessary right now with the workaround above in place.
                 }
             }
     }
@@ -261,6 +264,27 @@ class QRCodeLoginScreenViewModel: QRCodeLoginScreenViewModelType, QRCodeLoginScr
         }
     }
     
+    private func requestOIDCAuthorization(url: URL) {
+        let (stream, continuation) = AsyncStream<Result<Void, OIDCError>>.makeStream()
+        actionsSubject.send(.requestOIDCAuthorisation(url, continuation))
+        
+        oidcResultTask = Task { [weak self] in
+            for await result in stream {
+                guard let self else { return }
+                switch result {
+                case .success:
+                    break // The state will be updated by the status publisher.
+                case .failure(.userCancellation):
+                    MXLog.info("User cancelled the WAS, dismissing.")
+                    actionsSubject.send(.dismiss)
+                case .failure:
+                    handleError(.unknown)
+                }
+            }
+        }
+        .asCancellable()
+    }
+    
     private func handleError(_ error: QRCodeLoginError) {
         MXLog.error("Failed to scan the QR code: \(error)")
         switch error {

ElementX/Sources/Screens/QRCodeLoginScreen/View/QRCodeLoginScreen.swift

@@ -260,13 +260,14 @@ struct QRCodeLoginScreen: View {
                     }
                 }
             } bottomContent: {
-                if case .inputCode = confirmCode {
+                switch confirmCode {
+                case .inputCode, .sendingCode:
                     Button(L10n.actionContinue) {
                         context.send(viewAction: .sendCheckCode)
                     }
                     .buttonStyle(.compound(.primary))
                     .disabled(context.checkCodeInput.count < 2 || confirmCode.isSending)
-                } else {
+                case .invalidCode:
                     Button(L10n.actionStartOver) {
                         context.send(viewAction: .errorAction(.startOver))
                     }

ElementX/Sources/Screens/Settings/AccountSettings/OIDCAccountSettingsPresenter.swift

@@ -20,21 +20,39 @@ class OIDCAccountSettingsPresenter: NSObject {
     private let presentationAnchor: UIWindow
     private let oidcRedirectURL: URL
     
-    init(accountURL: URL, presentationAnchor: UIWindow, appSettings: AppSettings) {
+    typealias Continuation = AsyncStream<Result<Void, OIDCError>>.Continuation
+    private let continuation: Continuation?
+    
+    init(accountURL: URL, presentationAnchor: UIWindow, appSettings: AppSettings, continuation: Continuation? = nil) {
         self.accountURL = accountURL
         self.presentationAnchor = presentationAnchor
         oidcRedirectURL = appSettings.oidcRedirectURL
+        self.continuation = continuation
         super.init()
     }
     
     /// Presents a web authentication session for the supplied data.
     func start() {
-        let session = ASWebAuthenticationSession(url: accountURL, callback: .oidcRedirectURL(oidcRedirectURL)) { _, _ in }
+        let session = ASWebAuthenticationSession(url: accountURL, callback: .oidcRedirectURL(oidcRedirectURL)) { [continuation] _, error in
+            guard let continuation else { return }
+            
+            if error?.isOIDCUserCancellation == true {
+                continuation.yield(.failure(.userCancellation))
+            } else {
+                let errorDescription = error.map(String.init(describing:)) ?? "Unknown error"
+                MXLog.error("A web authentication session error occurred: \(errorDescription)")
+                continuation.yield(.failure(.unknown))
+            }
+            
+            continuation.finish()
+        }
+        
         session.prefersEphemeralWebBrowserSession = false
         session.presentationContextProvider = self
         session.additionalHeaderFields = [
             "X-Element-User-Agent": UserAgentBuilder.makeASCIIUserAgent()
         ]
+        
         session.start()
     }
 }

UnitTests/Sources/QRCodeLoginScreenViewModelTests.swift

@@ -150,7 +150,7 @@ final class QRCodeLoginScreenViewModelTests: XCTestCase {
         try await deferredFailure.fulfill()
         
         deferredAction = deferFulfillment(viewModel.actionsPublisher) { action in
-            guard case .dismiss = action else { return false }
+            guard case .linkedDevice = action else { return false }
             return true
         }
         linkDesktopProgressSubject.send(.done)
@@ -179,17 +179,18 @@ final class QRCodeLoginScreenViewModelTests: XCTestCase {
         linkMobileProgressSubject.send(.waitingForAuthorisation(verificationURL: .homeDirectory))
         try await deferredAction.fulfill()
         
-        let currentState = context.viewState.state
+        // Note: The SDK rarely sends the done action, so this test has been updated for the workaround of finishing early.
-        let deferredFailure = deferFailure(context.$viewState, timeout: 1) { $0.state != currentState }
-        linkMobileProgressSubject.send(.syncingSecrets)
-        try await deferredFailure.fulfill()
-        
         deferredAction = deferFulfillment(viewModel.actionsPublisher) { action in
-            guard case .dismiss = action else { return false }
+            guard case .linkedDevice = action else { return false }
             return true
         }
-        linkMobileProgressSubject.send(.done)
+        linkMobileProgressSubject.send(.syncingSecrets)
         try await deferredAction.fulfill()
+        
+        let currentState = context.viewState.state
+        let deferredFailure = deferFailure(context.$viewState, timeout: 1) { $0.state != currentState }
+        linkMobileProgressSubject.send(.done)
+        try await deferredFailure.fulfill()
     }
     
     // MARK: - Helpers