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

    private enum State {
        case idle
        case waiting
        case finished
    }

    private var state: State = .idle

    // MARK: - Dependencies

    private let udpHole: SDLUDPHoleActor
    private let config: SDLConfiguration
    private let logger: SDLLogger

    // MARK: - Probe Data

    private var natAddress1: SocketAddress?
    private var natAddress2: SocketAddress?

    // MARK: - Completion

    private var onFinished: ((NatType) -> Void)?
    
    private var cookieId: UInt32 = 1
    
    private var currentCookieId: UInt32?
    
    // 建立step -> SDLStunProbeReply的映射关系
    private var replies: [UInt32: SDLStunProbeReply] = [:]
    private var timeoutTask: Task<Void, Never>?

    // MARK: - Init

    init(udpHole: SDLUDPHoleActor, config: SDLConfiguration, logger: SDLLogger) {
        self.udpHole = udpHole
        self.config = config
        self.logger = logger
    }

    // MARK: - Public API
    
    func probeNatType() async -> NatType {
        return await withCheckedContinuation { continuation in
            Task {
                await self.start { natType in
                    continuation.resume(returning: natType)
                }
            }
        }
    }
    
    /// 启动 NAT 探测（一次性）
    private func start(onFinished: @escaping (NatType) -> Void) async {
        guard case .idle = state else {
            return
        }
        
        self.onFinished = onFinished
        transition(to: .waiting)
        
        let cookieId = self.cookieId
        self.cookieId &+= 1
        self.currentCookieId = cookieId
        
        await self.sendProbe(cookie: cookieId)
        
        self.timeoutTask = Task {
            try? await Task.sleep(nanoseconds: 5_000_000_000)
            await self.handleTimeout()
        }
    }

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

    // MARK: - Internal helpers

    private func sendProbe(cookie: UInt32) async {
        let addressArray = config.stunProbeSocketAddressArray
        
        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 finish(_ type: NatType) {
        guard case .waiting = state else {
            return
        }
        
        transition(to: .finished)
        onFinished?(type)
        onFinished = nil
        
        // 取消定时器
        self.timeoutTask?.cancel()
        self.timeoutTask = nil
    }

    private func transition(to newState: State) {
        state = newState
    }
    
    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()
    }
    
}