From 5b37ac25521cdcc90d610569adc3c65fa6e3ab04 Mon Sep 17 00:00:00 2001
From: anlicheng <244108715@qq.com>
Date: Thu, 7 May 2026 08:43:48 +0800
Subject: [PATCH] fix SDLUDPHole

---
 Tun/Punchnet/Actors/SDLNATProberActor.swift | 28 ++++++++++++++++++++-
 Tun/Punchnet/Context/SDLContextActor.swift  | 24 +++++++++++++++---
 Tun/Punchnet/UDPHole/SDLUDPHole.swift       |  3 ++-
 3 files changed, 50 insertions(+), 5 deletions(-)

diff --git a/Tun/Punchnet/Actors/SDLNATProberActor.swift b/Tun/Punchnet/Actors/SDLNATProberActor.swift
index 9ba707c..bdc68f0 100644
--- a/Tun/Punchnet/Actors/SDLNATProberActor.swift
+++ b/Tun/Punchnet/Actors/SDLNATProberActor.swift
@@ -27,6 +27,7 @@ actor SDLNATProberActor {
         // 建立step -> SDLStunProbeReply的映射关系
         var replies: [UInt32: SDLStunProbeReply]
         var timeoutTask: Task<Void, Never>?
+        var sendTask: Task<Void, Never>?
         var continuation: CheckedContinuation<NatType, Never>
         
         private var isFinished: Bool = false
@@ -46,6 +47,7 @@ actor SDLNATProberActor {
             self.continuation.resume(returning: type)
             // 取消定时器
             self.timeoutTask?.cancel()
+            self.sendTask?.cancel()
             self.isFinished = true
         }
     }
@@ -68,6 +70,10 @@ actor SDLNATProberActor {
     // MARK: - Public API
     
     func probeNatType(using udpHole: SDLUDPHole) async -> NatType {
+        if Task.isCancelled {
+            return .blocked
+        }
+
         let cookieId = self.cookieId
         self.cookieId &+= 1
         
@@ -83,7 +89,7 @@ actor SDLNATProberActor {
                 continuation: continuation
             )
             self.sessions[cookieId] = session
-            Task {
+            session.sendTask = Task {
                 await self.sendProbe(using: udpHole, cookie: cookieId)
             }
         }
@@ -131,6 +137,14 @@ actor SDLNATProberActor {
             }
         }
     }
+
+    func cancelAll() {
+        let sessions = self.sessions
+        self.sessions.removeAll()
+        sessions.values.forEach { session in
+            session.finished(with: .blocked)
+        }
+    }
     
     /// 超时事件（由外部 Timer / Task 驱动）
     private func handleTimeout(cookie: UInt32) async {
@@ -158,9 +172,21 @@ actor SDLNATProberActor {
     // MARK: - Internal helpers
 
     private func sendProbe(using udpHole: SDLUDPHole, cookie: UInt32) async {
+        guard !Task.isCancelled else {
+            return
+        }
         await udpHole.send(type: .stunProbe, data: makeProbePacket(cookieId: cookie, step: 1, attr: .none), remoteAddress: addressArray[0][0])
+        guard !Task.isCancelled else {
+            return
+        }
         await udpHole.send(type: .stunProbe, data: makeProbePacket(cookieId: cookie, step: 2, attr: .none), remoteAddress: addressArray[1][1])
+        guard !Task.isCancelled else {
+            return
+        }
         await udpHole.send(type: .stunProbe, data: makeProbePacket(cookieId: cookie, step: 3, attr: .peer), remoteAddress: addressArray[0][0])
+        guard !Task.isCancelled else {
+            return
+        }
         await udpHole.send(type: .stunProbe, data: makeProbePacket(cookieId: cookie, step: 4, attr: .port), remoteAddress: addressArray[0][0])
     }
     
diff --git a/Tun/Punchnet/Context/SDLContextActor.swift b/Tun/Punchnet/Context/SDLContextActor.swift
index 690443c..6e7b339 100644
--- a/Tun/Punchnet/Context/SDLContextActor.swift
+++ b/Tun/Punchnet/Context/SDLContextActor.swift
@@ -182,17 +182,32 @@ actor SDLContextActor {
         let udpHole = self.udpHole
         self.udpHole = nil
         self.udpHoleLocalAddress = nil
-        self.udpHoleMonitorTask?.cancel()
+
+        let udpHoleMonitorTask = self.udpHoleMonitorTask
         self.udpHoleMonitorTask = nil
-        self.natProbeTask?.cancel()
+        let natProbeTask = self.natProbeTask
         self.natProbeTask = nil
+
+        udpHoleMonitorTask?.cancel()
+        natProbeTask?.cancel()
+        await self.proberActor.cancelAll()
         await udpHole?.stop()
+        if let natProbeTask {
+            await natProbeTask.value
+        }
+        if let udpHoleMonitorTask {
+            await udpHoleMonitorTask.value
+        }
 
         let udpHoleV6 = self.udpHoleV6
         self.udpHoleV6 = nil
-        self.udpHoleV6MonitorTask?.cancel()
+        let udpHoleV6MonitorTask = self.udpHoleV6MonitorTask
         self.udpHoleV6MonitorTask = nil
+        udpHoleV6MonitorTask?.cancel()
         udpHoleV6?.stop()
+        if let udpHoleV6MonitorTask {
+            await udpHoleV6MonitorTask.value
+        }
 
         let dnsClient = self.dnsClient
         self.dnsClient = nil
@@ -680,6 +695,9 @@ extension SDLContextActor {
         let proberActor = self.proberActor
         self.natProbeTask = Task { [weak self] in
             SDLLogger.log("[SDLContext] start probeNatType")
+            if Task.isCancelled {
+                return
+            }
             let natType = await proberActor.probeNatType(using: udpHole)
             if Task.isCancelled {
                 return
diff --git a/Tun/Punchnet/UDPHole/SDLUDPHole.swift b/Tun/Punchnet/UDPHole/SDLUDPHole.swift
index 453d1eb..1c480a4 100644
--- a/Tun/Punchnet/UDPHole/SDLUDPHole.swift
+++ b/Tun/Punchnet/UDPHole/SDLUDPHole.swift
@@ -172,8 +172,9 @@ private final class SDLUDPHoleHandler: ChannelInboundHandler {
         
     func stop() {
         self.finishMessageContinuationIfNeed(throwing: nil)
-        try? self.channel?.close().wait()
+        let channel = self.channel
         self.channel = nil
+        try? channel?.close().wait()
         try? self.group.syncShutdownGracefully()
         
         let counterActor = self.counterActor