fix quicClient

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) {

Tun/Punchnet/Actors/SDLContextActor.swift

@@ -166,18 +166,69 @@ 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!!!!")
+            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()
-            SDLLogger.log("[SDLContext] quicClient closed: \(exit)")
+                    
                     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") {
             let udpHole = try await self.startUDPHole()
             SDLLogger.log("[SDLContext] udp running!!!!")
@@ -247,7 +298,7 @@ actor SDLContextActor {
                 return
             }
 
-            await self.handleReadyTimeout()
+            self.handleReadyTimeout()
         }
         defer {
             timeoutTask.cancel()
@@ -266,27 +317,12 @@ actor SDLContextActor {
         try await self.setNetworkSettings(config: config, dnsServer: DNSHelper.dnsServer)
     }
     
-    private func startQUICClient() async throws -> SDLQUICClient {
+    private func handleQUICMessage(message: SDLQUICInboundMessage) async {
-        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()
+            self.startRegisterLoop()
             // 启动stun任务
             await self.startStunRequestTask(welcome: welcome)
             
@@ -296,7 +332,7 @@ actor SDLContextActor {
         case .registerSuperAck(let registerSuperAck):
             await self.handleRegisterSuperAck(registerSuperAck: registerSuperAck)
         case .registerSuperNak(let registerSuperNak):
-                    await self.handleRegisterSuperNak(nakPacket: 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)
@@ -310,11 +346,6 @@ actor SDLContextActor {
             await self.arpServer.handleArpResponse(arpResponse: arpResponse)
         }
     }
-        }
-        self.quicClient = quicClient
-        
-        return quicClient
-    }
     
     private func prepareTunnelNotifier() {
         // 启动noticeClient
@@ -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)
         }
     }
     

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 {

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,59 +49,36 @@ enum SDLQUICClientExit: Error, Sendable, CustomStringConvertible {
     }
 }
 
-private actor SDLQUICCloseWait {
+actor SDLQUICClient {
-    private var exit: SDLQUICClientExit?
-    private var waiters: [CheckedContinuation<SDLQUICClientExit, Never>] = []
-    
-    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 {
     private let allocator = ByteBufferAllocator()
     // 单个包最大64K
     private let maxPacketSize: Int
     // 最大缓冲区区为2M
     private let maxBufferSize: Int
     
+    private let readyState = SDLQUICReadyState()
+    
+    // 消息流
     public var messageStream: AsyncStream<SDLQUICInboundMessage>
     private let messageCont: AsyncStream<SDLQUICInboundMessage>.Continuation
+    
+    // 事件流
+    public var eventStream: AsyncStream<SDLQUICEvent>
+    private let eventCont: AsyncStream<SDLQUICEvent>.Continuation
+    
     private var readTask: Task<Void, Never>?
     private var pingTask: Task<Void, Never>?
 
     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(
@@ -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:
+            case .failed(let error):
-                self.readyCont.finish()
                 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,33 +123,37 @@ 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
         }
-        } onCancel: {
-            Task {
-                await self.close(.cancelled)
-                self.connection.cancel()
-                self.messageCont.finish()
-            }
-        }
     }
     
     func send(type: SDLPacketType, data: Data) {
@@ -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 {
+                self?.eventCont.yield(.writeFailed(error))
-                    await self.close(.writeFailed("\(error)"))
-                }
             }
         })
     }
     
-    func waitReady() async throws {
+    private func heartbeatLoop() async -> SDLQUICClientExit {
-        for await _ in readyStream {
+        let timerStream = SDLAsyncTimerStream()
-            return
+        timerStream.start(interval: .seconds(5))
+        
+        for await _ in timerStream.stream {
+            if Task.isCancelled {
+                break
             }
-        let exit = await closeWait.wait()
+            self.send(type: .ping, data: Data())
-        throw exit
         }
         
-    func waitClose() async -> SDLQUICClientExit {
+        SDLLogger.log("[SDLQUICClient] udp pingTask cancel", for: .debug)
-        await closeWait.wait()
+        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<Void, Error>] = []
+
+        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: 编解码器
+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 {

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 {