punchnet-macos/Tun/DNS/DNSLocalService.swift

import Foundation

actor DNSLocalService {
    private let publicDnsServers: [String]
    private var onEvent: DNSEventHandler = { _ in }

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

    init(publicDnsServers: [String], retryDelay: Duration = .seconds(5)) {
        self.publicDnsServers = publicDnsServers
        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 dnsServer = self.publicDnsServers.randomElement() ?? "223.5.5.5"
            let client = DNSLocalClient(host: dnsServer)
            self.currentClient = client

            SDLLogger.log("[DNSLocalService] dnsLocalClient started", category: .dns)

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

                guard !self.isStopping else {
                    break
                }

                SDLLogger.log("[DNSLocalService] dnsLocalClient 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("[DNSLocalService] dnsLocalClient failed: \(error.localizedDescription), will restart", category: .dns)
            }

            if self.consumeImmediateRestartRequest() {
                SDLLogger.log("[DNSLocalService] dnsLocalClient 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 query(tracker: DNSLocalClient.DNSTracker, dnsPayload: Data) async {
        await self.currentClient?.query(tracker: tracker, dnsPayload: dnsPayload)
    }

    private func run(client: DNSLocalClient) 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()
        }
    }

    private func clearCurrent(_ client: DNSLocalClient) {
        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
    }
}