element-x-ios/ElementX/Sources/Services/Authentication/LinkNewDeviceService.swift

//
// Copyright 2025 Element Creations Ltd.
//
// SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-Element-Commercial.
// Please see LICENSE files in the repository root for full details.
//

import Combine
import CoreImage.CIFilterBuiltins
import MatrixRustSDK
import SwiftUI

// sourcery: AutoMockable
protocol LinkNewDeviceServiceProtocol {
    /// Links a new device by showing it a QR code.
    func linkMobileDevice() -> LinkNewDeviceService.LinkMobileProgressPublisher
    /// Links a new device using a QR code generated by said device.
    func linkDesktopDevice(with scannedQRData: Data) -> LinkNewDeviceService.LinkDesktopProgressPublisher
}

class LinkNewDeviceService: LinkNewDeviceServiceProtocol {
    /// Publishes the progress of linking a new device by showing it a QR code.
    typealias LinkMobileProgressPublisher = CurrentValuePublisher<LinkMobileProgress, QRCodeLoginError>
    /// Publishes the progress of linking a new device by scanning a QR code generated by said device.
    typealias LinkDesktopProgressPublisher = CurrentValuePublisher<LinkDesktopProgress, QRCodeLoginError>
    
    /// The progress of linking a new device by showing it a QR code.
    enum LinkMobileProgress: Equatable {
        case starting
        case qrReady(UIImage)
        case qrScanned(CheckCodeSenderProxy)
        case waitingForAuthorisation(verificationURL: URL)
        case syncingSecrets
        case done
    }
    
    /// The progress of linking a new device by scanning a QR code generated by said device.
    enum LinkDesktopProgress: Equatable {
        case starting
        case establishingSecureChannel(checkCodeString: String)
        case waitingForAuthorisation(verificationURL: URL)
        case syncingSecrets
        case done
    }
    
    private let grantLoginHandler: GrantLoginWithQrCodeHandlerProtocol
    
    init(handler: GrantLoginWithQrCodeHandlerProtocol) {
        grantLoginHandler = handler
    }
    
    /// Links a new device by showing it a QR code.
    func linkMobileDevice() -> LinkMobileProgressPublisher {
        let progressSubject = CurrentValueSubject<LinkMobileProgress, QRCodeLoginError>(.starting)
        
        // Bridge from the SDK's synchronous callback into Swift Concurrency. Yielding is safe from
        // any thread; the single long-lived `for await` consumer forwards the progress updates on
        // the main actor in FIFO order.
        let (progressStream, progressContinuation) = AsyncStream<GrantGeneratedQrLoginProgress>.makeStream()
        
        let listener = SDKListener {
            progressContinuation.yield($0)
        }
        
        Task {
            for await rustProgress in progressStream {
                do {
                    try progressSubject.send(.init(rustProgress: rustProgress))
                } catch {
                    MXLog.error("Invalid GenerateProgress")
                    progressSubject.send(completion: .failure(.unknown))
                }
            }
        }
        
        Task {
            defer { progressContinuation.finish() }
            do {
                // Note: The SDK doesn't provide us with a way to cancel the grant if the user hit the cancel button 🤷‍♂️
                try await grantLoginHandler.generate(progressListener: listener) // The success state is handled by the listener.
                // We send the .done progress in case the listener didn't get a chance to pass it on from the SDK before being deallocated
                progressSubject.send(LinkMobileProgress.done)
            } catch let error as HumanQrGrantLoginError {
                MXLog.error("QR code reciprocate error: \(error)")
                progressSubject.send(completion: .failure(.init(rustError: error)))
            } catch {
                MXLog.error("QR code reciprocate unknown error: \(error)")
                progressSubject.send(completion: .failure(.unknown))
            }
        }
        
        return progressSubject.asCurrentValuePublisher()
    }
    
    /// Links a new device using a QR code generated by said device.
    func linkDesktopDevice(with scannedQRData: Data) -> LinkDesktopProgressPublisher {
        let progressSubject = CurrentValueSubject<LinkDesktopProgress, QRCodeLoginError>(.starting)
        
        // Same callback bridging as `linkMobileDevice()` above.
        let (progressStream, progressContinuation) = AsyncStream<GrantQrLoginProgress>.makeStream()
        
        let listener = SDKListener {
            progressContinuation.yield($0)
        }
        
        Task {
            for await rustProgress in progressStream {
                do {
                    try progressSubject.send(.init(rustProgress: rustProgress))
                } catch {
                    MXLog.error("Invalid ScanProgress")
                    progressSubject.send(completion: .failure(.unknown))
                }
            }
        }
        
        let qrCodeData: QrCodeData
        do {
            qrCodeData = try QrCodeData.fromBytes(bytes: scannedQRData)
        } catch {
            MXLog.error("QR code decode error: \(error)")
            progressSubject.send(completion: .failure(.invalidQRCode))
            progressContinuation.finish()
            return progressSubject.asCurrentValuePublisher()
        }
        
        // At some stage the SDK will have a `qrCodeData.intent` which we should check before continuing here.
        // Note the equivalent check will also happen for sign in with QR code in the AuthenticationService.
        
        Task {
            defer { progressContinuation.finish() }
            do {
                // Note: The SDK doesn't provide us with a way to cancel the grant if the user hit the cancel button 🤷‍♂️
                try await grantLoginHandler.scan(qrCodeData: qrCodeData, progressListener: listener) // The success state is handled by the listener.
                // We send the .done progress in case the listener didn't get a chance to pass it on from the SDK before being deallocated
                progressSubject.send(LinkDesktopProgress.done)
            } catch let error as HumanQrGrantLoginError {
                MXLog.error("QR code reciprocate error: \(error)")
                progressSubject.send(completion: .failure(.init(rustError: error)))
            } catch {
                MXLog.error("QR code reciprocate unknown error: \(error)")
                progressSubject.send(completion: .failure(.unknown))
            }
        }
        
        return progressSubject.asCurrentValuePublisher()
    }
}

extension LinkNewDeviceService.LinkMobileProgress: CustomStringConvertible {
    enum Error: Swift.Error {
        case invalidQRCodeData
        case invalidVerificationURI(String)
    }
    
    init(rustProgress: GrantGeneratedQrLoginProgress) throws {
        self = switch rustProgress {
        case .starting: .starting
        case .qrReady(let qrCode):
            if let image = UIImage(qrCodeData: qrCode.toBytes()) {
                .qrReady(image)
            } else {
                throw Error.invalidQRCodeData
            }
        case .qrScanned(let checkCodeSender): .qrScanned(.init(underlyingSender: checkCodeSender))
        case .waitingForAuth(let verificationURI):
            // verificationURI is a String; ASWebAuthenticationSession requires a URL.
            if let url = URL(string: verificationURI) {
                .waitingForAuthorisation(verificationURL: url)
            } else {
                throw Error.invalidVerificationURI(verificationURI)
            }
        case .syncingSecrets: .syncingSecrets
        case .done: .done
        }
    }
    
    var description: String {
        switch self {
        case .starting: "starting"
        case .qrReady: "qrReady"
        case .qrScanned: "qrScanned"
        case .waitingForAuthorisation: "waitingForAuthorisation"
        case .syncingSecrets: "syncingSecrets"
        case .done: "done"
        }
    }
}

extension LinkNewDeviceService.LinkDesktopProgress: CustomStringConvertible {
    enum Error: Swift.Error { case invalidVerificationURI(String) }
    
    init(rustProgress: GrantQrLoginProgress) throws {
        self = switch rustProgress {
        case .starting: .starting
        case .establishingSecureChannel(_, let checkCodeString): .establishingSecureChannel(checkCodeString: checkCodeString)
        case .waitingForAuth(let verificationURI):
            // verificationURI is a String; ASWebAuthenticationSession requires a URL.
            if let url = URL(string: verificationURI) {
                .waitingForAuthorisation(verificationURL: url)
            } else {
                throw Error.invalidVerificationURI(verificationURI)
            }
        case .syncingSecrets: .syncingSecrets
        case .done: .done
        }
    }
    
    var description: String {
        switch self {
        case .starting: "starting"
        case .establishingSecureChannel: "establishingSecureChannel"
        case .waitingForAuthorisation: "waitingForAuthorisation"
        case .syncingSecrets: "syncingSecrets"
        case .done: "done"
        }
    }
}

private extension QRCodeLoginError {
    init(rustError: HumanQrGrantLoginError) {
        self = switch rustError {
        case .ConnectionInsecure:
            .connectionInsecure
        case .InvalidCheckCode:
            .invalidCheckCode
        case .UnsupportedProtocol:
            .linkingNotSupported
        case .Expired, .NotFound, .DeviceNotFound:
            .expired
        case .Cancelled:
            .cancelled
        case .OtherDeviceAlreadySignedIn:
            .deviceAlreadySignedIn
        case .UnsupportedQrCodeType:
            .invalidQRCode
        case .Unknown, .MissingSecretsBackup, .DeviceIdAlreadyInUse:
            .unknown
        }
    }
}

nonisolated class CheckCodeSenderProxy: Equatable {
    static func == (lhs: CheckCodeSenderProxy, rhs: CheckCodeSenderProxy) -> Bool {
        lhs.underlyingSender === rhs.underlyingSender
    }
    
    let underlyingSender: CheckCodeSenderProtocol
    
    init(underlyingSender: CheckCodeSenderProtocol) {
        self.underlyingSender = underlyingSender
    }
    
    func send(code: UInt8) async throws {
        try await underlyingSender.send(code: code)
    }
}

extension UIImage {
    convenience init?(qrCodeData: Data) {
        let qrContext = CIContext()
        let qrFilter = CIFilter.qrCodeGenerator()
        
        qrFilter.message = qrCodeData
        qrFilter.correctionLevel = "Q"
        
        guard let outputImage = qrFilter.outputImage,
              let cgImage = qrContext.createCGImage(outputImage, from: outputImage.extent) else {
            MXLog.error("Failed to generate an image from the supplied QR code data.")
            return nil
        }
        
        self.init(cgImage: cgImage)
    }
}