From 3c42aa58f3298dcfbc67a0a271ca9a2fb667ffd0 Mon Sep 17 00:00:00 2001 From: anlicheng <244108715@qq.com> Date: Tue, 28 Apr 2026 16:12:59 +0800 Subject: [PATCH] fix quicClient --- Tun/Punchnet/Actors/ArpServer.swift | 4 +- Tun/Punchnet/Actors/SDLContextActor.swift | 155 ++++++++----- Tun/Punchnet/Actors/SDLPuncherActor.swift | 4 +- Tun/Punchnet/Actors/SDLQuicClient.swift | 265 +++++++++++++--------- Tun/Punchnet/Policy/IdentityStore.swift | 10 +- 5 files changed, 265 insertions(+), 173 deletions(-) diff --git a/Tun/Punchnet/Actors/ArpServer.swift b/Tun/Punchnet/Actors/ArpServer.swift index 89bd077..e748a7c 100644 --- a/Tun/Punchnet/Actors/ArpServer.swift +++ b/Tun/Punchnet/Actors/ArpServer.swift @@ -70,7 +70,7 @@ actor ArpServer { self.coolingDown = [:] } - func arpRequest(targetIp: UInt32, use quicClient: SDLQUICClient?) throws { + func arpRequest(targetIp: UInt32, use quicClient: SDLQUICClient?) async throws { guard let quicClient, self.coolingDown[targetIp] == nil else { return } @@ -82,7 +82,7 @@ actor ArpServer { var arpRequest = SDLArpRequest() arpRequest.targetIp = targetIp - quicClient.send(type: .arpRequest, data: try arpRequest.serializedData()) + await quicClient.send(type: .arpRequest, data: try arpRequest.serializedData()) } func handleArpResponse(arpResponse: SDLArpResponse) { diff --git a/Tun/Punchnet/Actors/SDLContextActor.swift b/Tun/Punchnet/Actors/SDLContextActor.swift index 1025688..eff735a 100644 --- a/Tun/Punchnet/Actors/SDLContextActor.swift +++ b/Tun/Punchnet/Actors/SDLContextActor.swift @@ -166,16 +166,67 @@ actor SDLContextActor { await self.supervisor.addWorker(name: "quicClient") { SDLLogger.log("[SDLContext] try start quicClient", for: .debug) - let quicClient = try await self.startQUICClient() - SDLLogger.log("[SDLContext] quicClient running!!!!") - let exit = await quicClient.run() - SDLLogger.log("[SDLContext] quicClient closed: \(exit)") - switch exit { - case .normal, .cancelled: - return - case .transportClosed, .readFailed, .writeFailed: - throw exit + + self.quicWorker?.cancel() + await self.quicClient?.stop() + + // 启动monitor + let quicClient = SDLQUICClient(host: self.config.serverHost, port: 443) + self.quicClient = quicClient + + await quicClient.start() + + try await quicClient.waitReady(timeout: .seconds(3)) + // 这里必须等待quic的协商完成 + try await Task.sleep(for: .seconds(0.3)) + SDLLogger.log("[SDLContext] start quic client: \(self.config.serverHost)") + +// self.quicWorker = Task { +// for await message in await quicClient.messageStream { +// await self.handleQUICMessage(message: message) +// } +// } + + try await withThrowingTaskGroup { group in + defer { + group.cancelAll() + } + + group.addTask { + for await message in await quicClient.messageStream { + await self.handleQUICMessage(message: message) + } + throw SDLQUICClientExit.transportClosed("messageStream finished") + } + + group.addTask { + let exit = await quicClient.run() + + switch exit { + case .normal, .cancelled: + return + + case .transportClosed, .readFailed, .writeFailed: + throw exit + } + } + + group.addTask { + for await event in await quicClient.eventStream { + switch event { + case .failed(let error): + throw error + case .cancelled: + throw SDLQUICClientExit.cancelled + case .writeFailed(let error): + throw error + } + } + throw SDLQUICClientExit.cancelled + } + try await group.next() } + } await self.supervisor.addWorker(name: "udpHole") { @@ -247,7 +298,7 @@ actor SDLContextActor { return } - await self.handleReadyTimeout() + self.handleReadyTimeout() } defer { timeoutTask.cancel() @@ -266,56 +317,36 @@ actor SDLContextActor { try await self.setNetworkSettings(config: config, dnsServer: DNSHelper.dnsServer) } - private func startQUICClient() async throws -> SDLQUICClient { - self.quicWorker?.cancel() - self.quicClient?.stop() - - // 启动monitor - let quicClient = SDLQUICClient(host: self.config.serverHost, port: 443) - quicClient.start() - - // 等待quic准备好 - try await quicClient.waitReady() - // 这里必须等待quic的协商完成 - try await Task.sleep(for: .seconds(0.2)) - SDLLogger.log("[SDLContext] start quic client: \(self.config.serverHost)") - - self.quicWorker = Task.detached { - for await message in quicClient.messageStream { - switch message { - case .welcome(let welcome): - SDLLogger.log("[SDLContext] quic welcome: \(welcome)") - // 注册 - await self.startRegisterLoop() - // 启动stun任务 - await self.startStunRequestTask(welcome: welcome) - - case .pong: - //SDLLogger.shared.log("[SDLContext] quic pong") - () - case .registerSuperAck(let registerSuperAck): - await self.handleRegisterSuperAck(registerSuperAck: registerSuperAck) - case .registerSuperNak(let registerSuperNak): - await self.handleRegisterSuperNak(nakPacket: registerSuperNak) - case .peerInfo(let peerInfo): - //SDLLogger.shared.log("[SDLContext] peer message: \(peerInfo)") - await self.puncherActor.handlePeerInfo(using: self.udpHole, udpHoleV6: self.udpHoleV6, peerInfo: peerInfo) - case .event(let event): - await self.handleEvent(event: event) - case .policyReponse(let policyResponse): - // 处理权限的请求问题 - await self.identifyStore.applyPolicyResponse(policyResponse) - case .arpResponse(let arpResponse): - //SDLLogger.shared.log("[SDLContext] get arp response: \(arpResponse)") - await self.arpServer.handleArpResponse(arpResponse: arpResponse) - } - } + private func handleQUICMessage(message: SDLQUICInboundMessage) async { + switch message { + case .welcome(let welcome): + SDLLogger.log("[SDLContext] quic welcome: \(welcome)") + // 注册 + self.startRegisterLoop() + // 启动stun任务 + await self.startStunRequestTask(welcome: welcome) + + case .pong: + //SDLLogger.shared.log("[SDLContext] quic pong") + () + case .registerSuperAck(let registerSuperAck): + await self.handleRegisterSuperAck(registerSuperAck: registerSuperAck) + case .registerSuperNak(let registerSuperNak): + self.handleRegisterSuperNak(nakPacket: registerSuperNak) + case .peerInfo(let peerInfo): + //SDLLogger.shared.log("[SDLContext] peer message: \(peerInfo)") + await self.puncherActor.handlePeerInfo(using: self.udpHole, udpHoleV6: self.udpHoleV6, peerInfo: peerInfo) + case .event(let event): + await self.handleEvent(event: event) + case .policyReponse(let policyResponse): + // 处理权限的请求问题 + await self.identifyStore.applyPolicyResponse(policyResponse) + case .arpResponse(let arpResponse): + //SDLLogger.shared.log("[SDLContext] get arp response: \(arpResponse)") + await self.arpServer.handleArpResponse(arpResponse: arpResponse) } - self.quicClient = quicClient - - return quicClient } - + private func prepareTunnelNotifier() { // 启动noticeClient // 旧的 UDP NoticeClient 已移除，改为初始化基于 App Group 的通知通道。 @@ -475,7 +506,7 @@ actor SDLContextActor { self.quicWorker?.cancel() self.quicWorker = nil - self.quicClient?.stop() + await self.quicClient?.stop() self.quicClient = nil await self.dnsClient?.stop() @@ -810,7 +841,7 @@ extension SDLContextActor { while !Task.isCancelled { switch self.superRegistrationStateMachine.makeLoopAction() { case .sendRegister: - self.doRegisterSuper() + await self.doRegisterSuper() case .stop: return } @@ -922,7 +953,7 @@ extension SDLContextActor { } } - private func doRegisterSuper() { + private func doRegisterSuper() async { // 注册 var registerSuper = SDLRegisterSuper() registerSuper.clientID = self.config.clientId @@ -936,7 +967,7 @@ extension SDLContextActor { if let registerSuperData = try? registerSuper.serializedData() { SDLLogger.log("[SDLContext] will send register super") - self.quicClient?.send(type: .registerSuper, data: registerSuperData) + await self.quicClient?.send(type: .registerSuper, data: registerSuperData) } } diff --git a/Tun/Punchnet/Actors/SDLPuncherActor.swift b/Tun/Punchnet/Actors/SDLPuncherActor.swift index fca520e..c44b05e 100644 --- a/Tun/Punchnet/Actors/SDLPuncherActor.swift +++ b/Tun/Punchnet/Actors/SDLPuncherActor.swift @@ -64,7 +64,7 @@ actor SDLPuncherActor { } } - func submitRegisterRequest(quicClient: SDLQUICClient?, request: RegisterRequest) { + func submitRegisterRequest(quicClient: SDLQUICClient?, request: RegisterRequest) async { guard let quicClient else { return } @@ -90,7 +90,7 @@ actor SDLPuncherActor { phase: .waitingPeerInfo(deadline: now.addingTimeInterval(self.peerInfoTimeout)) ) - quicClient.send(type: .queryInfo, data: queryData) + await quicClient.send(type: .queryInfo, data: queryData) } func handlePeerInfo(using udpHole: SDLUDPHole?, udpHoleV6: SDLUDPHoleV6?, peerInfo: SDLPeerInfo) async { diff --git a/Tun/Punchnet/Actors/SDLQuicClient.swift b/Tun/Punchnet/Actors/SDLQuicClient.swift index 52c5819..5bce54c 100644 --- a/Tun/Punchnet/Actors/SDLQuicClient.swift +++ b/Tun/Punchnet/Actors/SDLQuicClient.swift @@ -20,6 +20,12 @@ enum SDLQUICError: Error { case packetTooLarge } +enum SDLQUICEvent: Error { + case failed(Error) + case cancelled + case writeFailed(Error) +} + enum SDLQUICClientExit: Error, Sendable, CustomStringConvertible { case normal case cancelled @@ -43,60 +49,37 @@ enum SDLQUICClientExit: Error, Sendable, CustomStringConvertible { } } -private actor SDLQUICCloseWait { - private var exit: SDLQUICClientExit? - private var waiters: [CheckedContinuation] = [] - - func wait() async -> SDLQUICClientExit { - if let exit { - return exit - } - - return await withCheckedContinuation { continuation in - waiters.append(continuation) - } - } - - func close(_ exit: SDLQUICClientExit) { - guard self.exit == nil else { - return - } - - self.exit = exit - let waiters = self.waiters - self.waiters.removeAll() - - for waiter in waiters { - waiter.resume(returning: exit) - } - } -} - -final class SDLQUICClient { +actor SDLQUICClient { private let allocator = ByteBufferAllocator() // 单个包最大64K private let maxPacketSize: Int // 最大缓冲区区为2M private let maxBufferSize: Int + private let readyState = SDLQUICReadyState() + + // 消息流 public var messageStream: AsyncStream private let messageCont: AsyncStream.Continuation + + // 事件流 + public var eventStream: AsyncStream + private let eventCont: AsyncStream.Continuation + private var readTask: Task? private var pingTask: Task? private let connection: NWConnection private let queue = DispatchQueue(label: "com.sdl.QUICClient.queue") // 专用队列保证线程安全 - private let closeWait = SDLQUICCloseWait() - private let (readyStream, readyCont) = AsyncStream.makeStream(of: Void.self) - init(host: String, port: UInt16, maxPacketSize: Int = 64 * 1024, maxBufferSize: Int = 2 * 1024 * 1024) { let options = NWProtocolQUIC.Options(alpn: ["punchnet/1.0"]) self.maxBufferSize = maxBufferSize self.maxPacketSize = maxPacketSize (self.messageStream, self.messageCont) = AsyncStream.makeStream(of: SDLQUICInboundMessage.self) - + (self.eventStream, self.eventCont) = AsyncStream.makeStream(of: SDLQUICEvent.self) + // TODO 这里设置证书的校验逻辑 sec_protocol_options_set_verify_block( options.securityProtocolOptions, @@ -116,17 +99,22 @@ final class SDLQUICClient { SDLLogger.log("[SDLQUICClient] new state: \(state)", for: .debug) switch state { case .ready: - self.readyCont.yield() - self.readyCont.finish() - case .failed(let error): - self.readyCont.finish() Task { - await self.close(.transportClosed("failed: \(error)")) + await self.readyState.markReady() } - case .cancelled: - self.readyCont.finish() + case .failed(let error): Task { - await self.close(.cancelled) + await self.readyState.markFailed(error) + } + self.eventCont.yield(.failed(error)) + case .cancelled: + Task { + await self.readyState.markCancelled() + } + self.eventCont.yield(.cancelled) + case .setup, .preparing: + Task { + await self.readyState.markConnecting() } default: () @@ -135,32 +123,36 @@ final class SDLQUICClient { connection.start(queue: self.queue) } + func waitReady(timeout: Duration = .seconds(5)) async throws { + try await withThrowingTaskGroup(of: Void.self) { group in + group.addTask { + try await self.readyState.waitReady() + } + + group.addTask { + try await Task.sleep(for: timeout) + throw SDLQUICError.timeout + } + + try await group.next() + group.cancelAll() + } + } + func run() async -> SDLQUICClientExit { - await withTaskCancellationHandler { - await withTaskGroup(of: SDLQUICClientExit.self) { group in - group.addTask { - await self.readLoop() - } - group.addTask { - await self.heartbeatLoop() - } - group.addTask { - await self.waitClose() - } - - let exit = await group.next() ?? .normal - group.cancelAll() - self.connection.cancel() - self.messageCont.finish() - await self.close(exit) - return exit + await withTaskGroup(of: SDLQUICClientExit.self) { group in + group.addTask { + await self.readLoop() } - } onCancel: { - Task { - await self.close(.cancelled) - self.connection.cancel() - self.messageCont.finish() + + group.addTask { + await self.heartbeatLoop() } + + let exit = await group.next() ?? .normal + group.cancelAll() + + return exit } } @@ -171,26 +163,27 @@ final class SDLQUICClient { packet.append(type.rawValue) packet.append(data) - connection.send(content: packet, completion: .contentProcessed { error in + connection.send(content: packet, completion: .contentProcessed { [weak self] error in if let error { SDLLogger.log("[SDLQUICClient] send data get error: \(error)", for: .debug) - Task { - await self.close(.writeFailed("\(error)")) - } + self?.eventCont.yield(.writeFailed(error)) } }) } - func waitReady() async throws { - for await _ in readyStream { - return + private func heartbeatLoop() async -> SDLQUICClientExit { + let timerStream = SDLAsyncTimerStream() + timerStream.start(interval: .seconds(5)) + + for await _ in timerStream.stream { + if Task.isCancelled { + break + } + self.send(type: .ping, data: Data()) } - let exit = await closeWait.wait() - throw exit - } - - func waitClose() async -> SDLQUICClientExit { - await closeWait.wait() + + SDLLogger.log("[SDLQUICClient] udp pingTask cancel", for: .debug) + return .cancelled } func stop() { @@ -198,9 +191,95 @@ final class SDLQUICClient { } func close(_ exit: SDLQUICClientExit = .normal) async { - await closeWait.close(exit) + } + deinit { + self.readTask?.cancel() + self.pingTask?.cancel() + self.messageCont.finish() + } +} + +// --MARK: Ready状态机 +extension SDLQUICClient { + + actor SDLQUICReadyState { + enum State { + case idle + case connecting + case ready + case failed(Error) + case cancelled + } + + private var state: State = .idle + private var continuations: [CheckedContinuation] = [] + + func waitReady() async throws { + switch state { + case .ready: + return + + case .failed(let error): + throw error + + case .cancelled: + throw CancellationError() + + case .idle, .connecting: + try await withCheckedThrowingContinuation { continuation in + continuations.append(continuation) + } + } + } + + func markConnecting() { + switch state { + case .idle: + state = .connecting + default: + break + } + } + + func markReady() { + state = .ready + + let list = continuations + continuations.removeAll() + + for continuation in list { + continuation.resume() + } + } + + func markFailed(_ error: Error) { + state = .failed(error) + + let list = continuations + continuations.removeAll() + + for continuation in list { + continuation.resume(throwing: error) + } + } + + func markCancelled() { + state = .cancelled + + let list = continuations + continuations.removeAll() + + for continuation in list { + continuation.resume(throwing: CancellationError()) + } + } + } +} + +// --MARK: Reader +extension SDLQUICClient { private func readLoop() async -> SDLQUICClientExit { var buffer = allocator.buffer(capacity: self.maxBufferSize) let threshold = self.maxBufferSize / 10 * 6 @@ -238,21 +317,6 @@ final class SDLQUICClient { } } - private func heartbeatLoop() async -> SDLQUICClientExit { - let timerStream = SDLAsyncTimerStream() - timerStream.start(interval: .seconds(5)) - - for await _ in timerStream.stream { - if Task.isCancelled { - break - } - self.send(type: .ping, data: Data()) - } - - SDLLogger.log("[SDLQUICClient] udp pingTask cancel", for: .debug) - return .cancelled - } - // 尝试解析数据 private func parseFrames(buffer: inout ByteBuffer) throws -> [ByteBuffer] { guard buffer.readableBytes >= 2 else { @@ -294,8 +358,10 @@ final class SDLQUICClient { } } } - - // --MARK: 编解码器 +} + +// --MARK: 编解码器 +extension SDLQUICClient { private func decode(frame: ByteBuffer) -> SDLQUICInboundMessage? { var buffer = frame guard let type = buffer.readInteger(as: UInt8.self), @@ -356,14 +422,9 @@ final class SDLQUICClient { return nil } } - - deinit { - self.readTask?.cancel() - self.pingTask?.cancel() - self.messageCont.finish() - } } +// --MARK: quic验证 extension SDLQUICClient { enum QUICVerifier { diff --git a/Tun/Punchnet/Policy/IdentityStore.swift b/Tun/Punchnet/Policy/IdentityStore.swift index 17dea84..716ddb6 100644 --- a/Tun/Punchnet/Policy/IdentityStore.swift +++ b/Tun/Punchnet/Policy/IdentityStore.swift @@ -26,12 +26,12 @@ actor IdentityStore { } // 批量更新, 有外部任务驱动，因为这里依赖于当前的quicClient - func batUpdatePolicy(using quicClient: SDLQUICClient?, dstIdentityID: UInt32) { + func batUpdatePolicy(using quicClient: SDLQUICClient?, dstIdentityID: UInt32) async { guard let quicClient else { return } - self.identityMap.keys.forEach { identityId in + for identityId in self.identityMap.keys { var policyRequest = SDLPolicyRequest() policyRequest.srcIdentityID = identityId policyRequest.dstIdentityID = dstIdentityID @@ -39,13 +39,13 @@ actor IdentityStore { // 发送请求 if let queryData = try? policyRequest.serializedData() { - quicClient.send(type: .policyRequest, data: queryData) + await quicClient.send(type: .policyRequest, data: queryData) } } } // 提交权限请求 - func policyRequest(srcIdentityId: UInt32, dstIdentityId: UInt32, using quicClient: SDLQUICClient?) { + func policyRequest(srcIdentityId: UInt32, dstIdentityId: UInt32, using quicClient: SDLQUICClient?) async { guard let quicClient, !coolingDown.contains(srcIdentityId) else { return } @@ -59,7 +59,7 @@ actor IdentityStore { coolingDown.insert(srcIdentityId) // 发送请求 if let queryData = try? policyRequest.serializedData() { - quicClient.send(type: .policyRequest, data: queryData) + await quicClient.send(type: .policyRequest, data: queryData) } Task {