Add a WebRegistrationScreen (not included in the flow yet). (#3281)

2026-11-03 16:10:20 +00:00 · 2024-09-16 11:03:29 +01:00
parent 2745eabc0a
commit 6cfe09b96d
23 changed files with 743 additions and 34 deletions
@@ -341,6 +341,81 @@ open class ClientSDKMock: MatrixRustSDK.Client {
        }
    }

+    //MARK: - awaitRoomRemoteEcho
+
+    open var awaitRoomRemoteEchoRoomIdThrowableError: Error?
+    var awaitRoomRemoteEchoRoomIdUnderlyingCallsCount = 0
+    open var awaitRoomRemoteEchoRoomIdCallsCount: Int {
+        get {
+            if Thread.isMainThread {
+                return awaitRoomRemoteEchoRoomIdUnderlyingCallsCount
+            } else {
+                var returnValue: Int? = nil
+                DispatchQueue.main.sync {
+                    returnValue = awaitRoomRemoteEchoRoomIdUnderlyingCallsCount
+                }
+
+                return returnValue!
+            }
+        }
+        set {
+            if Thread.isMainThread {
+                awaitRoomRemoteEchoRoomIdUnderlyingCallsCount = newValue
+            } else {
+                DispatchQueue.main.sync {
+                    awaitRoomRemoteEchoRoomIdUnderlyingCallsCount = newValue
+                }
+            }
+        }
+    }
+    open var awaitRoomRemoteEchoRoomIdCalled: Bool {
+        return awaitRoomRemoteEchoRoomIdCallsCount > 0
+    }
+    open var awaitRoomRemoteEchoRoomIdReceivedRoomId: String?
+    open var awaitRoomRemoteEchoRoomIdReceivedInvocations: [String] = []
+
+    var awaitRoomRemoteEchoRoomIdUnderlyingReturnValue: Room!
+    open var awaitRoomRemoteEchoRoomIdReturnValue: Room! {
+        get {
+            if Thread.isMainThread {
+                return awaitRoomRemoteEchoRoomIdUnderlyingReturnValue
+            } else {
+                var returnValue: Room? = nil
+                DispatchQueue.main.sync {
+                    returnValue = awaitRoomRemoteEchoRoomIdUnderlyingReturnValue
+                }
+
+                return returnValue!
+            }
+        }
+        set {
+            if Thread.isMainThread {
+                awaitRoomRemoteEchoRoomIdUnderlyingReturnValue = newValue
+            } else {
+                DispatchQueue.main.sync {
+                    awaitRoomRemoteEchoRoomIdUnderlyingReturnValue = newValue
+                }
+            }
+        }
+    }
+    open var awaitRoomRemoteEchoRoomIdClosure: ((String) async throws -> Room)?
+
+    open override func awaitRoomRemoteEcho(roomId: String) async throws -> Room {
+        if let error = awaitRoomRemoteEchoRoomIdThrowableError {
+            throw error
+        }
+        awaitRoomRemoteEchoRoomIdCallsCount += 1
+        awaitRoomRemoteEchoRoomIdReceivedRoomId = roomId
+        DispatchQueue.main.async {
+            self.awaitRoomRemoteEchoRoomIdReceivedInvocations.append(roomId)
+        }
+        if let awaitRoomRemoteEchoRoomIdClosure = awaitRoomRemoteEchoRoomIdClosure {
+            return try await awaitRoomRemoteEchoRoomIdClosure(roomId)
+        } else {
+            return awaitRoomRemoteEchoRoomIdReturnValue
+        }
+    }
+
    //MARK: - cachedAvatarUrl

    open var cachedAvatarUrlThrowableError: Error?
@@ -2662,6 +2737,71 @@ open class ClientSDKMock: MatrixRustSDK.Client {
        }
    }

+    //MARK: - server
+
+    var serverUnderlyingCallsCount = 0
+    open var serverCallsCount: Int {
+        get {
+            if Thread.isMainThread {
+                return serverUnderlyingCallsCount
+            } else {
+                var returnValue: Int? = nil
+                DispatchQueue.main.sync {
+                    returnValue = serverUnderlyingCallsCount
+                }
+
+                return returnValue!
+            }
+        }
+        set {
+            if Thread.isMainThread {
+                serverUnderlyingCallsCount = newValue
+            } else {
+                DispatchQueue.main.sync {
+                    serverUnderlyingCallsCount = newValue
+                }
+            }
+        }
+    }
+    open var serverCalled: Bool {
+        return serverCallsCount > 0
+    }
+
+    var serverUnderlyingReturnValue: String?
+    open var serverReturnValue: String? {
+        get {
+            if Thread.isMainThread {
+                return serverUnderlyingReturnValue
+            } else {
+                var returnValue: String?? = nil
+                DispatchQueue.main.sync {
+                    returnValue = serverUnderlyingReturnValue
+                }
+
+                return returnValue!
+            }
+        }
+        set {
+            if Thread.isMainThread {
+                serverUnderlyingReturnValue = newValue
+            } else {
+                DispatchQueue.main.sync {
+                    serverUnderlyingReturnValue = newValue
+                }
+            }
+        }
+    }
+    open var serverClosure: (() -> String?)?
+
+    open override func server() -> String? {
+        serverCallsCount += 1
+        if let serverClosure = serverClosure {
+            return serverClosure()
+        } else {
+            return serverReturnValue
+        }
+    }
+
    //MARK: - session

    open var sessionThrowableError: Error?
@@ -9974,6 +10114,82 @@ open class NotificationSettingsSDKMock: MatrixRustSDK.NotificationSettings {
        try await unmuteRoomRoomIdIsEncryptedIsOneToOneClosure?(roomId, isEncrypted, isOneToOne)
    }
 }
+open class OidcAuthorizationDataSDKMock: MatrixRustSDK.OidcAuthorizationData {
+    init() {
+        super.init(noPointer: .init())
+    }
+
+    public required init(unsafeFromRawPointer pointer: UnsafeMutableRawPointer) {
+        fatalError("init(unsafeFromRawPointer:) has not been implemented")
+    }
+
+    fileprivate var pointer: UnsafeMutableRawPointer!
+
+    //MARK: - loginUrl
+
+    var loginUrlUnderlyingCallsCount = 0
+    open var loginUrlCallsCount: Int {
+        get {
+            if Thread.isMainThread {
+                return loginUrlUnderlyingCallsCount
+            } else {
+                var returnValue: Int? = nil
+                DispatchQueue.main.sync {
+                    returnValue = loginUrlUnderlyingCallsCount
+                }
+
+                return returnValue!
+            }
+        }
+        set {
+            if Thread.isMainThread {
+                loginUrlUnderlyingCallsCount = newValue
+            } else {
+                DispatchQueue.main.sync {
+                    loginUrlUnderlyingCallsCount = newValue
+                }
+            }
+        }
+    }
+    open var loginUrlCalled: Bool {
+        return loginUrlCallsCount > 0
+    }
+
+    var loginUrlUnderlyingReturnValue: String!
+    open var loginUrlReturnValue: String! {
+        get {
+            if Thread.isMainThread {
+                return loginUrlUnderlyingReturnValue
+            } else {
+                var returnValue: String? = nil
+                DispatchQueue.main.sync {
+                    returnValue = loginUrlUnderlyingReturnValue
+                }
+
+                return returnValue!
+            }
+        }
+        set {
+            if Thread.isMainThread {
+                loginUrlUnderlyingReturnValue = newValue
+            } else {
+                DispatchQueue.main.sync {
+                    loginUrlUnderlyingReturnValue = newValue
+                }
+            }
+        }
+    }
+    open var loginUrlClosure: (() -> String)?
+
+    open override func loginUrl() -> String {
+        loginUrlCallsCount += 1
+        if let loginUrlClosure = loginUrlClosure {
+            return loginUrlClosure()
+        } else {
+            return loginUrlReturnValue
+        }
+    }
+}
 open class QrCodeDataSDKMock: MatrixRustSDK.QrCodeData {
    init() {
        super.init(noPointer: .init())
@@ -13,6 +13,8 @@ struct LoginHomeserver: Equatable {
    let address: String
    /// The types login supported by the homeserver.
    var loginMode: LoginMode
+    /// A temporary helper URL that can be used for registration.
+    var registrationHelperURL: URL?
    
    /// Creates a new homeserver value.
    init(address: String, loginMode: LoginMode) {
@@ -23,6 +23,7 @@ enum LoginScreenCoordinatorAction {
    case signedIn(UserSessionProtocol)
 }

+// Note: This code was brought over from Riot, we should move the authentication service logic into the view model.
 final class LoginScreenCoordinator: CoordinatorProtocol {
    private let parameters: LoginScreenCoordinatorParameters
    private var viewModel: LoginScreenViewModelProtocol
@@ -66,9 +66,9 @@ struct ServerConfirmationScreen: View {
 // MARK: - Previews

 struct ServerConfirmationScreen_Previews: PreviewProvider, TestablePreview {
-    static let loginViewModel = ServerConfirmationScreenViewModel(authenticationService: MockAuthenticationServiceProxy(),
+    static let loginViewModel = ServerConfirmationScreenViewModel(authenticationService: MockAuthenticationService(),
                                                                  authenticationFlow: .login)
-    static let registerViewModel = ServerConfirmationScreenViewModel(authenticationService: MockAuthenticationServiceProxy(),
+    static let registerViewModel = ServerConfirmationScreenViewModel(authenticationService: MockAuthenticationService(),
                                                                     authenticationFlow: .register)
    
    static var previews: some View {
@@ -21,6 +21,7 @@ enum ServerSelectionScreenCoordinatorAction {
    case dismiss
 }

+// Note: This code was brought over from Riot, we should move the authentication service logic into the view model.
 final class ServerSelectionScreenCoordinator: CoordinatorProtocol {
    private let parameters: ServerSelectionScreenCoordinatorParameters
    private let userIndicatorController: UserIndicatorControllerProtocol
@@ -32,6 +32,7 @@ enum SoftLogoutScreenCoordinatorResult: CustomStringConvertible {
    }
 }

+// Note: This code was brought over from Riot, we should move the authentication service logic into the view model.
 final class SoftLogoutScreenCoordinator: CoordinatorProtocol {
    private let parameters: SoftLogoutScreenCoordinatorParameters
    private var viewModel: SoftLogoutScreenViewModelProtocol
@@ -0,0 +1,139 @@
+//
+// Copyright 2022-2024 New Vector Ltd.
+//
+// SPDX-License-Identifier: AGPL-3.0-only
+// Please see LICENSE in the repository root for full details.
+//
+
+import Compound
+import SwiftUI
+import WebKit
+
+struct WebRegistrationScreen: View {
+    @ObservedObject var context: WebRegistrationScreenViewModel.Context
+    
+    var body: some View {
+        NavigationStack {
+            WebRegistrationWebView(url: context.viewState.url, viewModelContext: context)
+                .navigationTitle(L10n.screenCreateAccountTitle)
+                .navigationBarTitleDisplayMode(.inline)
+                .toolbarBackground(.visible, for: .navigationBar)
+                .toolbar {
+                    ToolbarItem(placement: .cancellationAction) {
+                        Button(L10n.actionCancel) {
+                            context.send(viewAction: .cancel)
+                        }
+                    }
+                }
+        }
+    }
+}
+
+struct WebRegistrationWebView: UIViewRepresentable {
+    let url: URL
+    let viewModelContext: WebRegistrationScreenViewModel.Context
+    
+    func makeUIView(context: Context) -> WKWebView {
+        context.coordinator.webView
+    }
+    
+    func updateUIView(_ webView: WKWebView, context: Context) { }
+    
+    func makeCoordinator() -> Coordinator {
+        Coordinator(url: url, viewModelContext: viewModelContext)
+    }
+    
+    class Coordinator: NSObject, WKUIDelegate {
+        private let url: URL
+        private let viewModelContext: WebRegistrationScreenViewModel.Context
+        
+        private(set) var webView: WKWebView!
+        
+        init(url: URL, viewModelContext: WebRegistrationScreenViewModel.Context) {
+            self.url = url
+            self.viewModelContext = viewModelContext
+            
+            super.init()
+            
+            let eventHandlerName = "elementx"
+            let userContentController = WKUserContentController()
+            userContentController.add(WKScriptMessageHandlerWrapper(self), name: eventHandlerName)
+            
+            let eventHandlerScript = """
+            window.addEventListener(
+                "mobileregistrationresponse",
+                (event) => {
+                    window.webkit.messageHandlers.\(eventHandlerName).postMessage(JSON.stringify(event.detail));
+                },
+                false,
+              );
+            """
+            
+            let userScript = WKUserScript(source: eventHandlerScript, injectionTime: .atDocumentEnd, forMainFrameOnly: false)
+            userContentController.addUserScript(userScript)
+            
+            let configuration = WKWebViewConfiguration()
+            configuration.userContentController = userContentController
+            configuration.preferences.javaScriptCanOpenWindowsAutomatically = true
+            
+            webView = WKWebView(frame: .zero, configuration: configuration)
+            webView.uiDelegate = self
+            webView.load(URLRequest(url: url))
+        }
+        
+        nonisolated func userContentController(_ userContentController: WKUserContentController,
+                                               didReceive message: WKScriptMessage) {
+            guard let jsonString = message.body as? String, let jsonData = jsonString.data(using: .utf8) else {
+                MXLog.error("Unexpected response.")
+                return
+            }
+            
+            guard let credentials = try? JSONDecoder().decode(WebRegistrationCredentials.self, from: jsonData) else {
+                MXLog.error("Invalid response.")
+                return
+            }
+            
+            MXLog.info("Received login credentials.")
+            Task { await viewModelContext.send(viewAction: .signedIn(credentials)) }
+        }
+        
+        // MARK: WKUIDelegate
+        
+        func webView(_ webView: WKWebView,
+                     createWebViewWith configuration: WKWebViewConfiguration,
+                     for navigationAction: WKNavigationAction,
+                     windowFeatures: WKWindowFeatures) -> WKWebView? {
+            if let url = navigationAction.request.url, UIApplication.shared.canOpenURL(url) {
+                UIApplication.shared.open(url)
+            }
+            return nil
+        }
+    }
+    
+    /// Avoids retain loops between the configuration and webView coordinator
+    private class WKScriptMessageHandlerWrapper: NSObject, WKScriptMessageHandler {
+        private weak var coordinator: Coordinator?
+        
+        init(_ coordinator: Coordinator) {
+            self.coordinator = coordinator
+        }
+        
+        // MARK: WKScriptMessageHandler
+        
+        nonisolated func userContentController(_ userContentController: WKUserContentController,
+                                               didReceive message: WKScriptMessage) {
+            coordinator?.userContentController(userContentController, didReceive: message)
+        }
+    }
+}
+
+// MARK: - Previews
+
+struct WebRegistrationScreen_Previews: PreviewProvider {
+    static let viewModel = WebRegistrationScreenViewModel(registrationHelperURL: "https://develop.element.io/#/mobile_register")
+    static var previews: some View {
+        NavigationStack {
+            WebRegistrationScreen(context: viewModel.context)
+        }
+    }
+}
@@ -0,0 +1,83 @@
+//
+// Copyright 2022-2024 New Vector Ltd.
+//
+// SPDX-License-Identifier: AGPL-3.0-only
+// Please see LICENSE in the repository root for full details.
+//
+
+import Combine
+import MatrixRustSDK
+import SwiftUI
+
+struct WebRegistrationScreenCoordinatorParameters {
+    /// The service used to authenticate the user.
+    let authenticationService: AuthenticationServiceProtocol
+    
+    let userIndicatorController: UserIndicatorControllerProtocol
+}
+
+enum WebRegistrationScreenCoordinatorAction: CustomStringConvertible {
+    case cancel
+    case signedIn(UserSessionProtocol)
+    
+    var description: String {
+        switch self {
+        case .cancel: "cancel"
+        case .signedIn: "signedIn"
+        }
+    }
+}
+
+// Note: This code was based on the LoginScreen, we should move the authentication service logic into the view model.
+final class WebRegistrationScreenCoordinator: CoordinatorProtocol {
+    private let parameters: WebRegistrationScreenCoordinatorParameters
+    private let viewModel: WebRegistrationScreenViewModelProtocol
+    
+    private var cancellables = Set<AnyCancellable>()
+ 
+    private let actionsSubject: PassthroughSubject<WebRegistrationScreenCoordinatorAction, Never> = .init()
+    var actionsPublisher: AnyPublisher<WebRegistrationScreenCoordinatorAction, Never> {
+        actionsSubject.eraseToAnyPublisher()
+    }
+    
+    init(parameters: WebRegistrationScreenCoordinatorParameters) {
+        self.parameters = parameters
+        
+        guard let registrationHelperURL = parameters.authenticationService.homeserver.value.registrationHelperURL else {
+            MXLog.error("Attempted registration without a helper URL.")
+            fatalError("A helper URL is required.")
+        }
+        viewModel = WebRegistrationScreenViewModel(registrationHelperURL: registrationHelperURL)
+    }
+    
+    func start() {
+        viewModel.actionsPublisher.sink { [weak self] action in
+            MXLog.info("Coordinator: received view model action: \(action)")
+            
+            guard let self else { return }
+            switch action {
+            case .cancel:
+                actionsSubject.send(.cancel)
+            case .signedIn(let credentials):
+                Task { await self.completeRegistration(using: credentials) }
+            }
+        }
+        .store(in: &cancellables)
+    }
+        
+    func toPresentable() -> AnyView {
+        AnyView(WebRegistrationScreen(context: viewModel.context))
+    }
+    
+    // MARK: - Private
+    
+    private func completeRegistration(using credentials: WebRegistrationCredentials) async {
+        switch await parameters.authenticationService.completeWebRegistration(using: credentials) {
+        case .success(let userSession):
+            actionsSubject.send(.signedIn(userSession))
+        case .failure(let error):
+            MXLog.error("Failed registration: \(error)")
+            parameters.userIndicatorController.alertInfo = .init(id: UUID(), title: L10n.errorUnknown, message: String(describing: error))
+        }
+    }
+}
@@ -0,0 +1,51 @@
+//
+// Copyright 2022-2024 New Vector Ltd.
+//
+// SPDX-License-Identifier: AGPL-3.0-only
+// Please see LICENSE in the repository root for full details.
+//
+
+import Foundation
+
+enum WebRegistrationScreenViewModelAction: CustomStringConvertible {
+    case cancel
+    case signedIn(WebRegistrationCredentials)
+    
+    var description: String {
+        switch self {
+        case .cancel: "cancel"
+        case .signedIn: "signedIn"
+        }
+    }
+}
+
+struct WebRegistrationScreenViewState: BindableState {
+    var url: URL
+    var bindings = WebRegistrationScreenViewStateBindings()
+}
+
+struct WebRegistrationScreenViewStateBindings { }
+
+enum WebRegistrationScreenViewAction: CustomStringConvertible {
+    case cancel
+    case signedIn(WebRegistrationCredentials)
+    
+    var description: String {
+        switch self {
+        case .cancel: "cancel"
+        case .signedIn: "signedIn"
+        }
+    }
+}
+
+struct WebRegistrationCredentials: Decodable {
+    let userID: String
+    let accessToken: String
+    let deviceID: String
+    
+    enum CodingKeys: String, CodingKey {
+        case userID = "user_id"
+        case accessToken = "access_token"
+        case deviceID = "device_id"
+    }
+}
@@ -0,0 +1,34 @@
+//
+// Copyright 2022-2024 New Vector Ltd.
+//
+// SPDX-License-Identifier: AGPL-3.0-only
+// Please see LICENSE in the repository root for full details.
+//
+
+import Combine
+import MatrixRustSDK
+import SwiftUI
+
+typealias WebRegistrationScreenViewModelType = StateStoreViewModel<WebRegistrationScreenViewState, WebRegistrationScreenViewAction>
+
+class WebRegistrationScreenViewModel: WebRegistrationScreenViewModelType, WebRegistrationScreenViewModelProtocol {
+    private let actionsSubject: PassthroughSubject<WebRegistrationScreenViewModelAction, Never> = .init()
+    var actionsPublisher: AnyPublisher<WebRegistrationScreenViewModelAction, Never> {
+        actionsSubject.eraseToAnyPublisher()
+    }
+
+    init(registrationHelperURL: URL) {
+        super.init(initialViewState: WebRegistrationScreenViewState(url: registrationHelperURL))
+    }
+    
+    override func process(viewAction: WebRegistrationScreenViewAction) {
+        MXLog.info("View model: received view action: \(viewAction)")
+        
+        switch viewAction {
+        case .cancel:
+            actionsSubject.send(.cancel)
+        case .signedIn(let credentials):
+            actionsSubject.send(.signedIn(credentials))
+        }
+    }
+}
@@ -0,0 +1,14 @@
+//
+// Copyright 2022-2024 New Vector Ltd.
+//
+// SPDX-License-Identifier: AGPL-3.0-only
+// Please see LICENSE in the repository root for full details.
+//
+
+import Combine
+
+@MainActor
+protocol WebRegistrationScreenViewModelProtocol {
+    var actionsPublisher: AnyPublisher<WebRegistrationScreenViewModelAction, Never> { get }
+    var context: WebRegistrationScreenViewModelType.Context { get }
+}
@@ -40,15 +40,21 @@ class AuthenticationService: AuthenticationServiceProtocol {
            
            let client = try await makeClientBuilder().build(homeserverAddress: homeserverAddress)
            let loginDetails = await client.homeserverLoginDetails()
+            let elementWellKnown = await client.getElementWellKnown()
            
            MXLog.info("Sliding sync: \(client.slidingSyncVersion())")
            
-            if loginDetails.supportsOidcLogin() {
-                homeserver.loginMode = .oidc
+            homeserver.loginMode = if loginDetails.supportsOidcLogin() {
+                .oidc
            } else if loginDetails.supportsPasswordLogin() {
-                homeserver.loginMode = .password
+                .password
            } else {
-                homeserver.loginMode = .unsupported
+                .unsupported
+            }
+            
+            homeserver.registrationHelperURL = switch elementWellKnown {
+            case .success(let wellKnown): wellKnown.registrationHelperUrl.flatMap(URL.init)
+            case .failure: nil
            }
            
            self.client = client
@@ -125,6 +131,25 @@ class AuthenticationService: AuthenticationServiceProtocol {
        }
    }
    
+    func completeWebRegistration(using credentials: WebRegistrationCredentials) async -> Result<any UserSessionProtocol, AuthenticationServiceError> {
+        guard let client else { return .failure(.failedLoggingIn) }
+        let session = Session(accessToken: credentials.accessToken,
+                              refreshToken: nil,
+                              userId: credentials.userID,
+                              deviceId: credentials.deviceID,
+                              homeserverUrl: client.homeserver(),
+                              oidcData: nil,
+                              slidingSyncVersion: client.slidingSyncVersion())
+        
+        do {
+            try await client.restoreSession(session: session)
+            return await userSession(for: client)
+        } catch {
+            MXLog.error("Failed restoring the client using the provided credentials.")
+            return .failure(.failedUsingWebCredentials)
+        }
+    }
+    
    // MARK: - Private
    
    private func makeClientBuilder() -> AuthenticationClientBuilder {
@@ -27,6 +27,7 @@ enum AuthenticationServiceError: Error {
    case accountDeactivated
    case failedLoggingIn
    case sessionTokenRefreshNotSupported
+    case failedUsingWebCredentials
 }

 protocol AuthenticationServiceProtocol {
@@ -43,6 +44,8 @@ protocol AuthenticationServiceProtocol {
    func loginWithOIDCCallback(_ callbackURL: URL, data: OIDCAuthorizationDataProxy) async -> Result<UserSessionProtocol, AuthenticationServiceError>
    /// Performs a password login using the current homeserver.
    func login(username: String, password: String, initialDeviceName: String?, deviceID: String?) async -> Result<UserSessionProtocol, AuthenticationServiceError>
+    /// Completes registration using the credentials obtained via the helper URL.
+    func completeWebRegistration(using credentials: WebRegistrationCredentials) async -> Result<UserSessionProtocol, AuthenticationServiceError>
 }

 // MARK: - OIDC
@@ -9,7 +9,7 @@ import Combine
 import Foundation
 import MatrixRustSDK

-class MockAuthenticationServiceProxy: AuthenticationServiceProtocol {
+class MockAuthenticationService: AuthenticationServiceProtocol {
    let validCredentials = (username: "alice", password: "12345678")
    
    private let homeserverSubject: CurrentValueSubject<LoginHomeserver, Never>
@@ -58,4 +58,8 @@ class MockAuthenticationServiceProxy: AuthenticationServiceProtocol {
        let userSession = UserSessionMock(.init(clientProxy: ClientProxyMock(.init(userID: username))))
        return .success(userSession)
    }
+    
+    func completeWebRegistration(using credentials: WebRegistrationCredentials) async -> Result<any UserSessionProtocol, AuthenticationServiceError> {
+        .failure(.failedLoggingIn)
+    }
 }
@@ -0,0 +1,32 @@
+//
+// Copyright 2024 New Vector Ltd.
+//
+// SPDX-License-Identifier: AGPL-3.0-only
+// Please see LICENSE in the repository root for full details.
+//
+
+import Foundation
+import MatrixRustSDK
+
+extension ClientProtocol {
+    func getElementWellKnown() async -> Result<MatrixRustSDK.ElementWellKnown, ClientProxyError> {
+        do {
+            let serverName = if let userIDServerName = try? userIdServerName() {
+                "https://\(userIDServerName)"
+            } else {
+                server()
+            }
+            
+            guard let serverName,
+                  let url = URL(string: serverName)?.appending(path: "/.well-known/element/element.json") else {
+                return .failure(.invalidServerName)
+            }
+            
+            let response = try await getUrl(url: url.absoluteString)
+            let wellKnown = try makeElementWellKnown(string: response)
+            return .success(wellKnown)
+        } catch {
+            return .failure(.sdkError(error))
+        }
+    }
+}
@@ -594,20 +594,7 @@ class ClientProxy: ClientProxyProtocol {
    }
    
    func getElementWellKnown() async -> Result<ElementWellKnown?, ClientProxyError> {
-        guard let userIDServerName,
-              var url = URL(string: "https://\(userIDServerName)") else {
-            return .failure(.invalidUserIDServerName)
-        }
-        
-        url.append(path: "/.well-known/element/element.json")
-        
-        do {
-            let response = try await client.getUrl(url: url.absoluteString)
-            let sdkWellKnown = try makeElementWellKnown(string: response)
-            return .success(ElementWellKnown(sdkWellKnown))
-        } catch {
-            return .failure(.sdkError(error))
-        }
+        await client.getElementWellKnown().map(ElementWellKnown.init)
    }
        
    // MARK: Ignored users
@@ -32,7 +32,7 @@ enum ClientProxyError: Error {
    case sdkError(Error)
    
    case invalidMedia
-    case invalidUserIDServerName
+    case invalidServerName
    case failedUploadingMedia(Error, MatrixErrorCode)
    case roomPreviewIsPrivate
 }
@@ -19,8 +19,10 @@ struct ElementWellKnown {
    }
    
    let call: Call?
+    let registrationHelperURL: URL?
    
    init?(_ wellKnown: MatrixRustSDK.ElementWellKnown) {
-        call = Call(wellKnown.call)
+        call = wellKnown.call.flatMap(Call.init)
+        registrationHelperURL = wellKnown.registrationHelperUrl.flatMap(URL.init)
    }
 }
@@ -111,20 +111,20 @@ class MockScreen: Identifiable {
        switch id {
        case .login:
            let navigationStackCoordinator = NavigationStackCoordinator()
-            let coordinator = LoginScreenCoordinator(parameters: .init(authenticationService: MockAuthenticationServiceProxy(),
+            let coordinator = LoginScreenCoordinator(parameters: .init(authenticationService: MockAuthenticationService(),
                                                                       analytics: ServiceLocator.shared.analytics,
                                                                       userIndicatorController: ServiceLocator.shared.userIndicatorController))
            navigationStackCoordinator.setRootCoordinator(coordinator)
            return navigationStackCoordinator
        case .serverSelection:
            let navigationStackCoordinator = NavigationStackCoordinator()
-            let coordinator = ServerSelectionScreenCoordinator(parameters: .init(authenticationService: MockAuthenticationServiceProxy(),
+            let coordinator = ServerSelectionScreenCoordinator(parameters: .init(authenticationService: MockAuthenticationService(),
                                                                                 userIndicatorController: ServiceLocator.shared.userIndicatorController,
                                                                                 isModallyPresented: true))
            navigationStackCoordinator.setRootCoordinator(coordinator)
            return navigationStackCoordinator
        case .authenticationFlow:
-            let flowCoordinator = AuthenticationFlowCoordinator(authenticationService: MockAuthenticationServiceProxy(),
+            let flowCoordinator = AuthenticationFlowCoordinator(authenticationService: MockAuthenticationService(),
                                                                qrCodeLoginService: QRCodeLoginServiceMock(),
                                                                bugReportService: BugReportServiceMock(),
                                                                navigationRootCoordinator: navigationRootCoordinator,