fix stun

Tun/Punchnet/Actors/SDLNATProberActor.swift

@@ -21,168 +21,153 @@ actor SDLNATProberActor {
     }
 
     // MARK: - Internal State
-
+    
-    private enum State {
+    class ProbeSession {
-        case idle
+        var cookieId: UInt32
-        case waiting
+        // 建立step -> SDLStunProbeReply的映射关系
-        case finished
+        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
+        }
     }
 
-    private var state: State = .idle
-
     // MARK: - Dependencies
 
     private let udpHole: SDLUDPHoleActor
-    private let config: SDLConfiguration
+    private let addressArray: [[SocketAddress]]
     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?
+    private var sessions: [UInt32: ProbeSession] = [:]
     
-    // 建立step -> SDLStunProbeReply的映射关系
-    private var replies: [UInt32: SDLStunProbeReply] = [:]
-    private var timeoutTask: Task<Void, Never>?
-
     // MARK: - Init
 
-    init(udpHole: SDLUDPHoleActor, config: SDLConfiguration, logger: SDLLogger) {
+    init(udpHole: SDLUDPHoleActor, addressArray: [[SocketAddress]], logger: SDLLogger) {
         self.udpHole = udpHole
-        self.config = config
+        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.start { natType in
+                await self.sendProbe(cookie: cookieId)
-                    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 {
+        guard let session = self.sessions[reply.cookie] else {
             return
         }
         
-        replies[reply.step] = reply
+        session.replies[reply.step] = reply
         
         // 提前退出的情况，没有nat映射
-        if let step1 = replies[1] {
+        if let step1 = session.replies[1] {
             let localAddress = await self.udpHole.getLocalAddress()
             if address == localAddress {
-                finish(.noNat)
+                finish(cookie: session.cookieId, .noNat)
                 return
             }
         }
         
-        if let step1 = replies[1], let step2 = replies[2] {
+        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(.symmetric)
+                finish(cookie: session.cookieId, .symmetric)
                 return
             }
         }
         
         // 收到了所有的响应, 优先判断
-        if replies[1] != nil && replies[2] != nil && replies[3] != nil && replies[4] != nil {
+        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 = replies[3] {
+            if let step3 = session.replies[3] {
-                finish(.fullCone)
+                finish(cookie: session.cookieId, .fullCone)
                 return
             }
             
             // step3: ip1:port1 <---- ip1:port2 (port改变情况)
             // 如果能收到的说明是IP地址限制锥型NAT，如果不能收到说明是端口限制锥型。
-            if let step4 = replies[4] {
+            if let step4 = session.replies[4] {
-                finish(.coneRestricted)
+                finish(cookie: session.cookieId, .coneRestricted)
                 return
             }
         }
     }
     
     /// 超时事件（由外部 Timer / Task 驱动）
-    private func handleTimeout() async {
+    private func handleTimeout(cookie: UInt32) async {
-        guard case .waiting = state else {
+        guard let session = self.sessions[cookie] else {
             return
         }
         
-        if replies[1] == nil {
+        if session.replies[1] == nil {
-            finish(.blocked)
+            finish(cookie: cookie, .blocked)
-        } else if replies[3] != nil {
+        } else if session.replies[3] != nil {
-            finish(.fullCone)
+            finish(cookie: cookie, .fullCone)
-        } else if replies[4] != nil {
+        } else if session.replies[4] != nil {
-            finish(.coneRestricted)
+            finish(cookie: cookie, .coneRestricted)
         } else {
-            finish(.portRestricted)
+            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 {
-        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()

Tun/Punchnet/SDLContext.swift

@@ -267,7 +267,7 @@ public class SDLContext {
         
         // 开始探测nat的类型
         if let udpHoleActor = self.udpHoleActor {
-            self.proberActor = SDLNATProberActor(udpHole: udpHoleActor, config: self.config, logger: self.logger)
+            self.proberActor = SDLNATProberActor(udpHole: udpHoleActor, addressArray: self.config.stunProbeSocketAddressArray, logger: self.logger)
             self.natType = await self.proberActor!.probeNatType()
         }