punchnet-macos/Tun/Punchnet/DNS/DNSCloudClient.swift

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

final class DNSCloudClient {
    
    enum DNSCloudError: Error {
        case failed(Error)
        case cancelled
        case sendFailed(Error)
    }
    
    private enum State {
        case idle
        case running
        case stopped
    }
    
    private var state: State = .idle
    
    private var connection: NWConnection?
    private var receiveTask: Task<Void, Never>?
    private let dnsServerAddress: NWEndpoint
    
    // 用于对外输出收到的 DNS 响应包
    public let packetFlow: AsyncThrowingStream<Data, Error>
    private let packetContinuation: AsyncThrowingStream<Data, Error>.Continuation
    private var isPacketContinuationFinished: Bool = false
    
    /// - Parameter host: 你的 sn-server 地址 (如 "8.8.8.8")
    /// - Parameter port: 端口 (如 53)
    init(host: String, port: UInt16) {
        self.dnsServerAddress = .hostPort(host: NWEndpoint.Host(host), port: NWEndpoint.Port(integerLiteral: port))
        
        let packetPair = AsyncThrowingStream.makeStream(of: Data.self)
        self.packetFlow = packetPair.stream
        self.packetContinuation = packetPair.continuation
    }
    
    func start() {
        // 1. 配置参数：这是解决环路的关键
        let parameters = NWParameters.udp
        // 禁止此连接走 TUN 网卡（在 NE 中 TUN 通常被归类为 .other）
        parameters.prohibitedInterfaceTypes = [.other]
        // 2. 增强健壮性：启用多路径切换（替代 pathSelectionOptions 的意图）
        parameters.multipathServiceType = .handover
        
        // 2. 创建连接
        let connection = NWConnection(to: self.dnsServerAddress, using: parameters)
        connection.stateUpdateHandler = { [weak self] state in
            self?.handleConnectionStateUpdate(state, for: connection)
        }
        // 启动连接队列
        connection.start(queue: .global())
        
        self.connection = connection
    }
    
    /// 发送 DNS 查询包（由 TUN 拦截到的原始 IP 包数据）
    func forward(ipPacketData: Data) {
        guard let connection = self.connection, connection.state == .ready else {
            return
        }
        
        connection.send(content: ipPacketData, completion: .contentProcessed { error in
            if let error = error {
                self.finishPacketContinuationIfNeed(throwing: .sendFailed(error))
            }
        })
    }
    
    func stop() {
        guard self.state != .stopped else {
            return
        }
        
        self.state = .stopped
        
        self.receiveTask?.cancel()
        self.receiveTask = nil
        
        self.connection?.cancel()
        self.connection = nil
        
        self.finishPacketContinuationIfNeed(throwing: nil)
    }
    
    private func handleConnectionStateUpdate(_ state: NWConnection.State, for connection: NWConnection) {
        switch state {
        case .ready:
            SDLLogger.log("[DNSClient] Connection ready", for: .debug)
            self.startReceiveTask(for: connection)
            self.state = .running
        case .failed(let error):
            self.finishPacketContinuationIfNeed(throwing: .failed(error))
        case .cancelled:
            self.finishPacketContinuationIfNeed(throwing: .cancelled)
        default:
            break
        }
    }
    
    private func startReceiveTask(for connection: NWConnection) {
        guard self.receiveTask == nil else {
            return
        }
        
        let stream = Self.makeReceiveStream(for: connection)
        self.receiveTask = Task { [weak self] in
            for await data in stream {
                if Task.isCancelled {
                    break
                }
                self?.packetContinuation.yield(data)
            }
        }
    }
    
    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 static func makeReceiveStream(for connection: NWConnection) -> AsyncStream<Data> {
        return AsyncStream(bufferingPolicy: .bufferingNewest(256)) { continuation in
            func receiveNext() {
                connection.receiveMessage { content, _, _, error in
                    if let data = content, !data.isEmpty {
                        // 将收到的 DNS 响应写回 AsyncStream
                        continuation.yield(data)
                    }
                    
                    if error == nil && connection.state == .ready {
                        receiveNext() // 继续监听下一个包
                    } else {
                        continuation.finish()
                    }
                }
            }
            
            receiveNext()
        }
    }
    
}