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

import Foundation
import Network

final class DNSLocalClient {
    // 需要保存DNS请求的追踪信息
    struct DNSTracker {
        let transactionID: UInt16
        let clientIP: UInt32   // 原始包的源 IP (大端序)
        let clientPort: UInt16 // 原始包的源端口 (大端序)
        let createdAt: Date    // 用于超时清理
    }
    
    private var connections: [NWConnection] = []
    
    // 阿里云 + 腾讯云
    private let dnsServers = ["223.5.5.5", "119.29.29.29"]

    public let packetFlow: AsyncStream<Data>
    private let packetContinuation: AsyncStream<Data>.Continuation
    
    private let locker = NSLock()
    private var trackers: [UInt16: [DNSTracker]] = [:]
    
    // 定期的任务清理
    private var cleanupTask: Task<Void, Never>?
    private let timeoutInterval: TimeInterval = 10.0 // 超过10秒认为丢包

    init() {
        let (stream, continuation) = AsyncStream.makeStream(of: Data.self, bufferingPolicy: .unbounded)
        self.packetFlow = stream
        self.packetContinuation = continuation
    }

    func start() {
        for server in dnsServers {
            let endpoint = NWEndpoint.hostPort(host: NWEndpoint.Host(server), port: 53)
            let parameters = NWParameters.udp
            parameters.prohibitedInterfaceTypes = [.other]
            
            let conn = NWConnection(to: endpoint, using: parameters)
            
            conn.stateUpdateHandler = { [weak self] state in
                switch state {
                case .ready:
                    self?.receiveLoop(for: conn)
                case .failed(let error):
                    SDLLogger.log("[DNSLocalClient] failed with error: \(error.localizedDescription)", for: .debug)
                    self?.stop()
                case .cancelled:
                    self?.packetContinuation.finish()
                default:
                    ()
                }
            }
            conn.start(queue: .global())
            
            connections.append(conn)
        }
        
        // 启动清理循环
        self.cleanupTask = Task { [weak self] in
            while !Task.isCancelled {
                // 每隔 cleanupTick 秒运行一次
                try? await Task.sleep(nanoseconds: 5 * 1_000_000_000)
                self?.performCleanup()
            }
        }
    }

    /// 并发查询：对所有服务器广播
    func query(tracker: DNSTracker, dnsPayload: Data) {
        locker.lock()
        self.trackers[tracker.transactionID, default: []].append(tracker)
        locker.unlock()
        
        for conn in connections where conn.state == .ready {
            conn.send(content: dnsPayload, completion: .contentProcessed({ _ in }))
        }
    }

    private func receiveLoop(for conn: NWConnection) {
        conn.receiveMessage { [weak self] content, _, _, error in
            if let data = content {
                // ！！！核心：由于 AsyncStream 是流式的
                // 谁先 yield，上层就先收到谁。
                // 只要上层收到了第一个有效响应并回填给系统，
                self?.handleResponse(data: data)
            }
            
            if error == nil && conn.state == .ready {
                self?.receiveLoop(for: conn)
            }
        }
    }
    
    private func handleResponse(data: Data) {
        guard data.count > 2 else {
            return
        }
        
        let tranId = UInt16(data[0]) << 8 | UInt16(data[1])
        
        locker.lock()
        let items = self.trackers.removeValue(forKey: tranId)
        locker.unlock()
        
        items?.forEach { tracker in
            let packet = Self.createDNSResponse(
                payload: data,
                srcIP: DNSHelper.dnsDestIpAddr,
                srcPort: 53,
                destIP: tracker.clientIP,
                destPort: tracker.clientPort
            )
            self.packetContinuation.yield(packet)
        }
    }
    
    private func performCleanup() {
        locker.lock()
        defer {
            locker.unlock()
        }
        
        // 遍历所有 ID，过滤掉过期的 tracker
        let now = Date()
        for (id, list) in trackers {
            let validItems = list.filter { now.timeIntervalSince($0.createdAt) < timeoutInterval }
            
            if validItems.isEmpty {
                trackers.removeValue(forKey: id)
            } else {
                trackers[id] = validItems
            }
        }
    }
    
    func stop() {
        connections.forEach { conn in
            conn.cancel()
        }
        self.connections.removeAll()
        
        self.cleanupTask?.cancel()
        self.cleanupTask = nil
    }
    
}

extension DNSLocalClient {
    /// 构造发回 TUN 的完整 UDP/IPv4 数据包
    static func createDNSResponse(payload: Data, srcIP: UInt32, srcPort: UInt16, destIP: UInt32, destPort: UInt16) -> Data {
        let udpLen = 8 + payload.count
        let ipLen = 20 + udpLen
        
        // --- 1. IPv4 Header (20 字节) ---
        var ipHeader = Data(count: 20)
        ipHeader[0] = 0x45 // Version 4, IHL 5
        ipHeader[2...3] = withUnsafeBytes(of: UInt16(ipLen).bigEndian) { Data($0) }
        ipHeader[8] = 64   // TTL
        ipHeader[9] = 17   // Protocol UDP
        
        // 填充 IP 地址
        ipHeader[12...15] = withUnsafeBytes(of: srcIP.bigEndian) { Data($0) }
        ipHeader[16...19] = withUnsafeBytes(of: destIP.bigEndian) { Data($0) }
        
        // 计算 IP Checksum
        let ipChecksum = calculateChecksum(data: ipHeader)
        ipHeader[10...11] = withUnsafeBytes(of: ipChecksum.bigEndian) { Data($0) }
        
        // --- 2. UDP Header (8 字节) ---
        var udpHeader = Data(count: 8)
        udpHeader[0...1] = withUnsafeBytes(of: srcPort.bigEndian) { Data($0) }
        udpHeader[2...3] = withUnsafeBytes(of: destPort.bigEndian) { Data($0) }
        udpHeader[4...5] = withUnsafeBytes(of: UInt16(udpLen).bigEndian) { Data($0) }
        // UDP Checksum 在 IPv4 中可选，设为 0 可跳过计算（大部分系统接受）
        udpHeader[6...7] = Data([0, 0])
        
        // --- 3. 拼接 ---
        var packet = Data(capacity: ipLen)
        packet.append(ipHeader)
        packet.append(udpHeader)
        packet.append(payload)
        
        return packet
    }
    
    /// 经典的 Internet Checksum 算法
    static func calculateChecksum(data: Data) -> UInt16 {
        var sum: UInt32 = 0
        let count = data.count
        
        data.withUnsafeBytes { (ptr: UnsafeRawBufferPointer) in
            guard let baseAddress = ptr.baseAddress else { return }
            
            // 1. 处理成对的 16-bit 单词
            let wordCount = count / 2
            let words = baseAddress.bindMemory(to: UInt16.self, capacity: wordCount)
            
            for i in 0..<wordCount {
                // 使用 bigEndian: words[i] 是正确的，但要确保不越界
                sum += UInt32(UInt16(bigEndian: words[i]))
            }
            
            // 2. 关键修复：处理奇数长度的最后一个字节
            if count % 2 != 0 {
                // 获取最后一个字节，并左移 8 位（作为 16-bit 的高位）
                let lastByte = ptr[count - 1]
                sum += UInt32(lastByte) << 8
            }
        }
        
        // 3. 将进位加回低 16 位
        while (sum >> 16) != 0 {
            sum = (sum & 0xffff) + (sum >> 16)
        }
        
        return UInt16(~sum & 0xffff)
    }

    
}