fix dnsClient

Tun/Punchnet/Actors/SDLContextActor.swift

@@ -312,53 +312,25 @@ actor SDLContextActor {
     private func startDnsClient() async throws {
         // 启动dns服务
         let dnsClient = DNSCloudClient(host: self.config.serverHost, port: 15353)
-        dnsClient.start()
-        SDLLogger.log("[SDLContext] dnsClient started")
         self.dnsClient = dnsClient
+        dnsClient.start()
         
-        try await withThrowingTaskGroup { group in
-            defer {
-                group.cancelAll()
-            }
-            
-            group.addTask {
-                for await packet in dnsClient.packetFlow {
-                    try Task.checkCancellation()
-                    
-                    let nePacket = NEPacket(data: packet, protocolFamily: 2)
-                    self.provider.packetFlow.writePacketObjects([nePacket])
-                }
-                throw SDLContextError.dnsClientClosed
-            }
-            
-            group.addTask {
-                for await event in dnsClient.eventStream {
-                    try Task.checkCancellation()
-                    switch event {
-                    case .failed(let error):
-                        SDLLogger.log("[SDLContext] dnsClient failed with error: \(error)")
-                        throw error
-                    case .cancelled:
-                        SDLLogger.log("[SDLContext] dnsClient cancelled")
-                        throw SDLContextError.dnsClientCancelled
-                    case .sendFailed(let error):
-                        SDLLogger.log("[SDLContext] dnsClient sendFailed with error: \(error)")
-                        throw error
-                    }
-                }
-            }
-            
-            do {
-                try await group.next()
-                dnsClient.stop()
-                self.dnsClient = nil
-            } catch let err {
-                dnsClient.stop()
-                self.dnsClient = nil
-                throw err
-            }
+        defer {
+            self.dnsClient = nil
+            dnsClient.stop()
         }
         
+        try await withTaskCancellationHandler {
+            for try await packet in dnsClient.packetFlow {
+                try Task.checkCancellation()
+                
+                let nePacket = NEPacket(data: packet, protocolFamily: 2)
+                self.provider.packetFlow.writePacketObjects([nePacket])
+            }
+            throw SDLContextError.dnsClientClosed
+        } onCancel: {
+            dnsClient.stop()
+        }
     }
     
     private func startDnsLocalClient() async throws {

Tun/Punchnet/DNS/DNSCloudClient.swift

@@ -9,7 +9,7 @@ import Network
 
 final class DNSCloudClient {
     
-    enum Event {
+    enum DNSCloudError: Error {
         case failed(Error)
         case cancelled
         case sendFailed(Error)
@@ -28,25 +28,18 @@ final class DNSCloudClient {
     private let dnsServerAddress: NWEndpoint
     
     // 用于对外输出收到的 DNS 响应包
-    public let packetFlow: AsyncStream<Data>
-    private let packetContinuation: AsyncStream<Data>.Continuation
-    
-    // 处理Connection的事件
-    public let eventStream: AsyncStream<Event>
-    private let eventContinuation: AsyncStream<Event>.Continuation
+    public let packetFlow: AsyncThrowingStream<Data, Error>
+    private let packetContinuation: AsyncThrowingStream<Data, Error>.Continuation
+    private var isPacketContinuationFinished: Bool = false
     
     /// - Parameter host: 你的 sn-server 地址 (如 "8.8.8.8")
     /// - Parameter port: 端口 (如 53)
     init(host: String, port: UInt16) {
         self.dnsServerAddress = .hostPort(host: NWEndpoint.Host(host), port: NWEndpoint.Port(integerLiteral: port))
         
-        let packetPair = AsyncStream.makeStream(of: Data.self, bufferingPolicy: .bufferingNewest(256))
+        let packetPair = AsyncThrowingStream.makeStream(of: Data.self)
         self.packetFlow = packetPair.stream
         self.packetContinuation = packetPair.continuation
-        
-        let eventPair = AsyncStream.makeStream(of: Event.self)
-        self.eventStream = eventPair.stream
-        self.eventContinuation = eventPair.continuation
     }
     
     func start() {
@@ -76,7 +69,7 @@ final class DNSCloudClient {
         
         connection.send(content: ipPacketData, completion: .contentProcessed { error in
             if let error = error {
-                self.eventContinuation.yield(.sendFailed(error))
+                self.finishPacketContinuationIfNeed(throwing: .sendFailed(error))
             }
         })
     }
@@ -94,8 +87,7 @@ final class DNSCloudClient {
         self.connection?.cancel()
         self.connection = nil
         
-        self.packetContinuation.finish()
-        self.eventContinuation.finish()
+        self.finishPacketContinuationIfNeed(throwing: nil)
     }
     
     private func handleConnectionStateUpdate(_ state: NWConnection.State, for connection: NWConnection) {
@@ -105,9 +97,9 @@ final class DNSCloudClient {
             self.startReceiveTask(for: connection)
             self.state = .running
         case .failed(let error):
-            self.eventContinuation.yield(.failed(error))
+            self.finishPacketContinuationIfNeed(throwing: .failed(error))
         case .cancelled:
-            self.eventContinuation.yield(.cancelled)
+            self.finishPacketContinuationIfNeed(throwing: .cancelled)
         default:
             break
         }
@@ -129,6 +121,19 @@ final class DNSCloudClient {
         }
     }
     
+    private func finishPacketContinuationIfNeed(throwing error: DNSCloudError?) {
+        guard !self.isPacketContinuationFinished else {
+            return
+        }
+        
+        self.isPacketContinuationFinished = true
+        if let error {
+            self.packetContinuation.finish(throwing: error)
+        } else {
+            self.packetContinuation.finish()
+        }
+    }
+    
     /// 接收数据的递归循环
     private static func makeReceiveStream(for connection: NWConnection) -> AsyncStream<Data> {
         return AsyncStream(bufferingPolicy: .bufferingNewest(256)) { continuation in