punchnet-macos/Tun/Punchnet/Actors/SDLNATProberActor.swift

//
//  SDLNATProberActor.swift
//  punchnet
//
//  Created by 安礼成 on 2026/1/28.
//
import Foundation
import NIOCore

actor SDLNATProberActor {

    // MARK: - NAT Type

    enum NatType: UInt8, Encodable {
        case blocked = 0
        case noNat = 1
        case fullCone = 2
        case portRestricted = 3
        case coneRestricted = 4
        case symmetric = 5
    }

    // MARK: - Internal State
    
    class ProbeSession {
        var cookieId: UInt32
        // 建立step -> SDLStunProbeReply的映射关系
        var replies: [UInt32: SDLStunProbeReply]
        var timeoutTask: Task<Void, Never>?
        var continuation: CheckedContinuation<NatType, Never>
        
        private var isFinished: Bool = false
        
        init(cookieId: UInt32, timeoutTask: Task<Void, Never>? = nil, continuation: CheckedContinuation<NatType, Never>) {
            self.cookieId = cookieId
            self.replies = [:]
            self.timeoutTask = timeoutTask
            self.continuation = continuation
        }
        
        func finished(with type: NatType) {
            guard !isFinished else {
                return
            }
            
            self.continuation.resume(returning: type)
            // 取消定时器
            self.timeoutTask?.cancel()
            self.isFinished = true
        }
    }

    // MARK: - Dependencies

    private let udpHole: SDLUDPHoleActor
    private let addressArray: [[SocketAddress]]
    private let logger: SDLLogger

    // MARK: - Completion
    private var cookieId: UInt32 = 1
    
    private var sessions: [UInt32: ProbeSession] = [:]
    
    // MARK: - Init

    init(udpHole: SDLUDPHoleActor, addressArray: [[SocketAddress]], logger: SDLLogger) {
        self.udpHole = udpHole
        self.addressArray = addressArray
        self.logger = logger
    }

    // MARK: - Public API
    
    func probeNatType() async -> NatType {
        let cookieId = self.cookieId
        self.cookieId &+= 1
        
        return await withCheckedContinuation { continuation in
            let timeoutTask = Task {
                try? await Task.sleep(nanoseconds: 5_000_000_000)
                await self.handleTimeout(cookie: cookieId)
            }
            
            let session = ProbeSession(
                cookieId: cookieId,
                timeoutTask: timeoutTask,
                continuation: continuation
            )
            self.sessions[cookieId] = session
            Task {
                await self.sendProbe(cookie: cookieId)
            }
        }
    }
    
    /// UDP 层收到 STUN 响应后调用
    func handleProbeReply(from address: SocketAddress, reply: SDLStunProbeReply) async {
        guard let session = self.sessions[reply.cookie] else {
            return
        }
        
        session.replies[reply.step] = reply
        
        // 提前退出的情况，没有nat映射
        if let step1 = session.replies[1] {
            let localAddress = await self.udpHole.getLocalAddress()
            if address == localAddress {
                finish(cookie: session.cookieId, .noNat)
                return
            }
        }
        
        if let step1 = session.replies[1], let step2 = session.replies[2] {
            // 如果natAddress2 的IP地址与上次回来的IP是不一样的，它就是对称型NAT; 这次的包也一定能发成功并收到
            // 如果ip地址变了，这说明{dstIp, dstPort, srcIp, srcPort}, 其中有一个变了；则用新的ip地址
            if let addr1 = step1.socketAddress(), let addr2 = step2.socketAddress(), addr1 != addr2 {
                finish(cookie: session.cookieId, .symmetric)
                return
            }
        }
        
        // 收到了所有的响应, 优先判断
        if session.replies[1] != nil && session.replies[2] != nil && session.replies[3] != nil && session.replies[4] != nil {
            // step3: ip2:port2 <---- ip1:port1 (ip地址和port都变的情况)
            // 如果能收到的，说明是完全锥形 说明是IP地址限制锥型NAT，如果不能收到说明是端口限制锥型。
            if let step3 = session.replies[3] {
                finish(cookie: session.cookieId, .fullCone)
                return
            }
            
            // step3: ip1:port1 <---- ip1:port2 (port改变情况)
            // 如果能收到的说明是IP地址限制锥型NAT，如果不能收到说明是端口限制锥型。
            if let step4 = session.replies[4] {
                finish(cookie: session.cookieId, .coneRestricted)
                return
            }
        }
    }
    
    /// 超时事件（由外部 Timer / Task 驱动）
    private func handleTimeout(cookie: UInt32) async {
        guard let session = self.sessions[cookie] else {
            return
        }
        
        if session.replies[1] == nil {
            finish(cookie: cookie, .blocked)
        } else if session.replies[3] != nil {
            finish(cookie: cookie, .fullCone)
        } else if session.replies[4] != nil {
            finish(cookie: cookie, .coneRestricted)
        } else {
            finish(cookie: cookie, .portRestricted)
        }
    }
    
    private func finish(cookie: UInt32, _ type: NatType) {
        if let session = self.sessions.removeValue(forKey: cookie) {
            session.finished(with: type)
        }
    }

    // MARK: - Internal helpers

    private func sendProbe(cookie: UInt32) async {
        await self.udpHole.send(type: .stunProbe, data: makeProbePacket(cookieId: cookie, step: 1, attr: .none), remoteAddress: addressArray[0][0])
        await self.udpHole.send(type: .stunProbe, data: makeProbePacket(cookieId: cookie, step: 2, attr: .none), remoteAddress: addressArray[1][1])
        await self.udpHole.send(type: .stunProbe, data: makeProbePacket(cookieId: cookie, step: 3, attr: .peer), remoteAddress: addressArray[0][0])
        await self.udpHole.send(type: .stunProbe, data: makeProbePacket(cookieId: cookie, step: 4, attr: .port), remoteAddress: addressArray[0][0])
    }
    
    private func makeProbePacket(cookieId: UInt32, step: UInt32, attr: SDLProbeAttr) -> Data {
        var stunProbe = SDLStunProbe()
        stunProbe.cookie = cookieId
        stunProbe.step = step
        stunProbe.attr = UInt32(attr.rawValue)
        
        return try! stunProbe.serializedData()
    }
    
}