punchnet-macos/Tun/DNS/DNSCloudService.swift

import Foundation

actor DNSCloudService {
    private let serverIP: String
    private var onEvent: DNSEventHandler = { _ in }

    private var currentClient: DNSCloudClient?
    private var isRunning = false
    private var isStopping = false
    private var needsImmediateRestart = false
    private let retryDelay: Duration

    init(serverIP: String, retryDelay: Duration = .seconds(5)) {
        self.serverIP = serverIP
        self.retryDelay = retryDelay
    }

    func updateEventHandler(_ onEvent: @escaping DNSEventHandler) {
        self.onEvent = onEvent
    }

    func run() async throws {
        guard !self.isRunning else {
            return
        }

        self.isRunning = true
        self.isStopping = false

        defer {
            self.isRunning = false
            self.currentClient = nil
        }

        while !Task.isCancelled, !self.isStopping {
            let client = DNSCloudClient(serverIP: self.serverIP, port: 15353)
            self.currentClient = client

            do {
                try await self.run(client: client)
                self.clearCurrent(client)
                await client.stop()

                guard !self.isStopping else {
                    break
                }

                SDLLogger.log("[DNSCloudService] dnsCloudClient ended, will restart", category: .dns)
            } catch is CancellationError {
                self.clearCurrent(client)
                await client.stop()
                throw CancellationError()
            } catch {
                self.clearCurrent(client)
                await client.stop()

                guard !self.isStopping else {
                    break
                }

                SDLLogger.log("[DNSCloudService] dnsCloudClient failed: \(error.localizedDescription), will restart", category: .dns)
            }

            if self.consumeImmediateRestartRequest() {
                SDLLogger.log("[DNSCloudService] dnsCloudClient invalidated after wakeup, will restart immediately", category: .dns)
                continue
            }

            try await Task.sleep(for: self.retryDelay)
        }
    }

    func stop() async {
        self.isStopping = true
        self.needsImmediateRestart = false
        await self.invalidateCurrentClient()
    }

    func recoverAfterWake() async {
        guard !self.isStopping else {
            return
        }

        self.needsImmediateRestart = self.currentClient != nil
        await self.invalidateCurrentClient()
    }

    func forward(ipPacketData: Data) async {
        await self.currentClient?.forward(ipPacketData: ipPacketData)
    }

    private func clearCurrent(_ client: DNSCloudClient) {
        if self.currentClient === client {
            self.currentClient = nil
        }
    }

    private func invalidateCurrentClient() async {
        let client = self.currentClient
        self.currentClient = nil

        await client?.stop()
    }

    private func consumeImmediateRestartRequest() -> Bool {
        let needsImmediateRestart = self.needsImmediateRestart
        self.needsImmediateRestart = false
        return needsImmediateRestart
    }

    private func run(client: DNSCloudClient) async throws {
        let onEvent = self.onEvent

        try await withThrowingTaskGroup(of: Void.self) { group in
            defer {
                group.cancelAll()
            }

            group.addTask {
                try await client.run()
            }

            group.addTask {
                for try await packet in client.packetFlow {
                    try Task.checkCancellation()
                    await onEvent(.packet(packet))
                }
            }

            _ = try await group.next()
        }
    }
}