punchnet-macos/Tun/DNS/DNSService.swift

import Foundation
import NetworkExtension

actor DNSService {
    enum Event {
        case packet(Data)
    }

    typealias EventHandler = @Sendable (Event) async -> Void

    private let serverIP: String
    private let publicDnsServers: [String]
    private let onEvent: EventHandler

    private var dnsClient: DNSCloudClient?
    private var dnsLocalClient: DNSLocalClient?

    init(serverIP: String, publicDnsServers: [String], onEvent: @escaping EventHandler) {
        self.serverIP = serverIP
        self.publicDnsServers = publicDnsServers
        self.onEvent = onEvent
    }

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

            group.addTask {
                try await Self.runRestarting(name: "dnsServiceCloud") {
                    try await self.runCloud()
                }
            }

            group.addTask {
                try await Self.runRestarting(name: "dnsServiceLocal") {
                    try await self.runLocal()
                }
            }

            try await group.waitForAll()
        }
    }

    func stop() async {
        let dnsClient = self.dnsClient
        self.dnsClient = nil

        let dnsLocalClient = self.dnsLocalClient
        self.dnsLocalClient = nil

        dnsClient?.stop()
        await dnsLocalClient?.stop()
    }

    func forward(ipPacketData: Data) {
        self.dnsClient?.forward(ipPacketData: ipPacketData)
    }

    func queryLocal(tracker: DNSLocalClient.DNSTracker, dnsPayload: Data) async {
        await self.dnsLocalClient?.query(tracker: tracker, dnsPayload: dnsPayload)
    }

    private func runCloud() async throws {
        let dnsClient = DNSCloudClient(serverIP: self.serverIP, port: 15353)
        self.dnsClient = dnsClient

        defer {
            dnsClient.stop()
            if self.dnsClient === dnsClient {
                self.dnsClient = nil
            }
        }

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

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

            group.addTask {
                try await withTaskCancellationHandler {
                    for try await packet in dnsClient.packetFlow {
                        try Task.checkCancellation()
                        await onEvent(.packet(packet))
                    }
                } onCancel: {
                    dnsClient.stop()
                }
            }

            try await group.next()
        }
    }

    private func runLocal() async throws {
        let dnsServer = self.publicDnsServers.randomElement() ?? "223.5.5.5"
        let dnsLocalClient = DNSLocalClient(host: dnsServer)
        self.dnsLocalClient = dnsLocalClient
        SDLLogger.log("[DNSService] dnsLocalClient started")

        defer {
            if self.dnsLocalClient === dnsLocalClient {
                self.dnsLocalClient = nil
            }
        }

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

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

                group.addTask {
                    try await withTaskCancellationHandler {
                        for try await packet in dnsLocalClient.packetFlow {
                            try Task.checkCancellation()
                            await onEvent(.packet(packet))
                        }
                    } onCancel: {
                        Task {
                            await dnsLocalClient.stop()
                        }
                    }
                }

                try await group.next()
            }
            await dnsLocalClient.stop()
        } catch {
            await dnsLocalClient.stop()
            throw error
        }
    }

    private static func runRestarting(
        name: String,
        retryDelay: Duration = .seconds(5),
        operation: @escaping @Sendable () async throws -> Void
    ) async throws {
        while !Task.isCancelled {
            do {
                try Task.checkCancellation()
                try await operation()
                SDLLogger.log("[DNSService] worker \(name) ended, will restart", for: .debug)
            } catch is CancellationError {
                SDLLogger.log("[DNSService] worker \(name) cancelled", for: .debug)
                throw CancellationError()
            } catch {
                SDLLogger.log("[DNSService] worker \(name) crashed: \(error.localizedDescription), will restart", for: .debug)
            }

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