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(step: Step)
        case finished
    }

    private enum Step: Int {
        case step1 = 1
        case step2 = 2
        case step3 = 3
        case step4 = 4
    }

    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

    // MARK: - Init

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

    // MARK: - Public API

    /// 启动 NAT 探测（一次性）
    func start(onFinished: @escaping (NatType) -> Void) async {
        guard case .idle = state else {
            logger.log("[NAT] probe already started", level: .warning)
            return
        }
        
        self.onFinished = onFinished
        transition(to: .waiting(step: .step1))
        await self.sendProbe(step: .step1)
    }

    /// UDP 层收到 STUN 响应后调用
    func handleProbeReply(from address: SocketAddress, reply: SDLStunProbeReply) async {
        guard case .waiting(let currentStep) = state else {
            return
        }

        switch currentStep {
        case .step1:
            let localAddress = await self.udpHole.getLocalAddress()
            if address == localAddress {
                finish(.noNat)
                return
            }
            
            natAddress1 = address
            transition(to: .waiting(step: .step2))
            await self.sendProbe(step: .step2)

        case .step2:
            natAddress2 = address
            // 如果natAddress2 的IP地址与上次回来的IP是不一样的，它就是对称型NAT; 这次的包也一定能发成功并收到
            // 如果ip地址变了，这说明{dstIp, dstPort, srcIp, srcPort}, 其中有一个变了；则用新的ip地址
            if let ip1 = natAddress1?.ipAddress, let ip2 = natAddress2?.ipAddress, ip1 != ip2 {
                finish(.symmetric)
                return
            }

            transition(to: .waiting(step: .step3))
            await self.sendProbe(step: .step3)
        case .step3:
            // step3: ip1:port1 <---- ip2:port2 (ip地址和port都变的情况)
            // 如果能收到的，说明是完全锥形 说明是IP地址限制锥型NAT，如果不能收到说明是端口限制锥型。
            finish(.fullCone)
        case .step4:
            finish(.coneRestricted)
        }
    }

    /// 超时事件（由外部 Timer / Task 驱动）
    func handleTimeout() async {
        guard case .waiting(let currentStep) = state else {
            return
        }
        
        switch currentStep {
        case .step3:
            transition(to: .waiting(step: .step4))
            await sendProbe(step: .step4)
        case .step4:
            finish(.portRestricted)
        default:
            finish(.blocked)
        }
    }

    // MARK: - Internal helpers

    private func sendProbe(step: Step) async {
        let addressArray = config.stunProbeSocketAddressArray

        let remote: SocketAddress
        let attr: SDLProbeAttr

        switch step {
        case .step1:
            remote = addressArray[0][0]
            attr = .none
        case .step2:
            remote = addressArray[1][1]
            attr = .none
        case .step3:
            remote = addressArray[0][0]
            attr = .peer
        case .step4:
            remote = addressArray[0][0]
            attr = .port
        }

        var stunProbe = SDLStunProbe()
        stunProbe.cookie = self.cookieId
        stunProbe.attr = UInt32(attr.rawValue)
        
        self.cookieId &+= 1
        
        await self.udpHole.send(type: .stunProbe, data: try! stunProbe.serializedData(), remoteAddress: remote)
    }

    private func finish(_ type: NatType) {
        guard case .finished = state else {
            transition(to: .finished)
            logger.log("[NAT] finished with \(type)", level: .info)
            onFinished?(type)
            onFinished = nil
            return
        }
    }

    private func transition(to newState: State) {
        state = newState
    }
}