punchnet-macos/Tun/DNS/DNSCloudClient.swift

//
//  SDLDNSClient 2.swift
//  punchnet
//
//  Created by 安礼成 on 2026/4/9.
//
import Foundation
import Network

actor DNSCloudClient {
    
    enum DNSCloudError: Error {
        case failed(Error)
        case cancelled
        case sendFailed(Error)
        case invalidData
    }
    
    private let queue = DispatchQueue(label: "com.sdl.DNSCloudClient.queue")
    private let connection: NWConnection
    
    // 用于对外输出收到的 DNS 响应包
    nonisolated let packetFlow: AsyncThrowingStream<Data, Error>
    private let packetContinuation: AsyncThrowingStream<Data, Error>.Continuation
    
    private let readySignal = AsyncOneShot<Void>()
    
    private var isStopped: Bool = false
    private var isPacketContinuationFinished: Bool = false
    
    /// - Parameter serverIP: 你的 sn-server IP 地址 (如 "8.8.8.8")
    /// - Parameter port: 端口 (如 53)
    init(serverIP: String, port: UInt16) {
        let dnsServerAddress = NWEndpoint.hostPort(host: Self.makeEndpointHost(address: serverIP), port: NWEndpoint.Port(integerLiteral: port))
        
        let packetPair = AsyncThrowingStream.makeStream(of: Data.self)
        self.packetFlow = packetPair.stream
        self.packetContinuation = packetPair.continuation
        
        // 1. 配置参数：这是解决环路的关键
        let parameters = NWParameters.udp
        // 禁止此连接走 TUN 网卡（在 NE 中 TUN 通常被归类为 .other）
        parameters.prohibitedInterfaceTypes = [.other]
        // 2. 增强健壮性：启用多路径切换（替代 pathSelectionOptions 的意图）
        parameters.multipathServiceType = .handover
        
        // 2. 创建连接
        self.connection = NWConnection(to: dnsServerAddress, using: parameters)
    }

    private static func makeEndpointHost(address ip: String) -> NWEndpoint.Host {
        if let ipv4Address = IPv4Address(ip) {
            return .ipv4(ipv4Address)
        }

        if let ipv6Address = IPv6Address(ip) {
            return .ipv6(ipv6Address)
        }

        preconditionFailure("invalid DNS cloud server IP: \(ip)")
    }
    
    func run() async throws {
        self.connection.stateUpdateHandler = { [weak self] state in
            Task {
                await self?.handleConnectionStateUpdate(state)
            }
        }
        self.connection.start(queue: self.queue)
        
        try await withTaskCancellationHandler {
            try await self.readySignal.wait()

            while true {
                try Task.checkCancellation()
                let data = try await self.readOnce()
                self.packetContinuation.yield(data)
            }
        } onCancel: {
            self.connection.cancel()
        }
    }
    
    /// 发送 DNS 查询包（由 TUN 拦截到的原始 IP 包数据）
    func forward(ipPacketData: Data) async {
        do {
            try await self.readySignal.wait(timeout: .seconds(3))
        } catch {
            SDLLogger.log("[DNSCloudClient] drop query before ready: \(error)", category: .dns)
            return
        }

        guard !self.isStopped, connection.state == .ready else {
            return
        }
        
        connection.send(content: ipPacketData, completion: .contentProcessed { [weak self] error in
            if let error {
                Task {
                    await self?.finishPacketContinuationIfNeed(throwing: .sendFailed(error))
                }
            }
        })
    }
    
    func stop() async {
        guard !self.isStopped else {
            return
        }
        self.isStopped = true
        
        self.connection.cancel()
        await self.readySignal.fail(DNSCloudError.cancelled)
        self.finishPacketContinuationIfNeed(throwing: nil)
        
        SDLLogger.log("[SDLCloudClient] stopped", category: .dns)
    }
    
    private func handleConnectionStateUpdate(_ state: NWConnection.State) async {
        switch state {
        case .ready:
            guard !self.isStopped else {
                return
            }

            SDLLogger.log("[DNSClient] Connection ready", category: .dns)
            await self.readySignal.succeed(())
        case .failed(let error):
            await self.readySignal.fail(DNSCloudError.failed(error))
            self.finishPacketContinuationIfNeed(throwing: .failed(error))
        case .cancelled:
            await self.readySignal.fail(DNSCloudError.cancelled)
            self.finishPacketContinuationIfNeed(throwing: .cancelled)
        default:
            break
        }
    }

    private func finishPacketContinuationIfNeed(throwing error: DNSCloudError?) {
        guard !self.isPacketContinuationFinished else {
            return
        }

        self.isPacketContinuationFinished = true

        if let error {
            self.packetContinuation.finish(throwing: error)
        } else {
            self.packetContinuation.finish()
        }
    }
    
    private func readOnce() async throws -> Data {
        guard self.connection.state == .ready else {
            throw DNSCloudError.cancelled
        }
        
        let readContinuation = OnceContinuation<Data, Error>()
        return try await withTaskCancellationHandler {
            try await withCheckedThrowingContinuation { cont in
                readContinuation.set(cont)
                self.connection.receiveMessage { content, _, _, error in
                    if let error {
                        readContinuation.resume(throwing: error)
                    } else if let data = content, !data.isEmpty {
                        readContinuation.resume(returning: data)
                    } else {
                        readContinuation.resume(throwing: DNSCloudError.invalidData)
                    }
                }
            }
        } onCancel: {
            readContinuation.resume(throwing: CancellationError())
        }
    }
    
    deinit {
        SDLLogger.log("[DNSCloudClient] deinit", category: .dns)
    }
    
}