diff --git a/Tun/Punchnet/Actors/SDLContextActor.swift b/Tun/Punchnet/Actors/SDLContextActor.swift index f8c8d4e..d86b217 100644 --- a/Tun/Punchnet/Actors/SDLContextActor.swift +++ b/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 { diff --git a/Tun/Punchnet/DNS/DNSCloudClient.swift b/Tun/Punchnet/DNS/DNSCloudClient.swift index f068a30..aceda18 100644 --- a/Tun/Punchnet/DNS/DNSCloudClient.swift +++ b/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 - private let packetContinuation: AsyncStream.Continuation + public let packetFlow: AsyncThrowingStream + private let packetContinuation: AsyncThrowingStream.Continuation + private var isPacketContinuationFinished: Bool = false - // 处理Connection的事件 - public let eventStream: AsyncStream - private let eventContinuation: AsyncStream.Continuation - /// - 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 { return AsyncStream(bufferingPolicy: .bufferingNewest(256)) { continuation in