fix quicClient

Tun/Punchnet/Actors/SDLContextActor.swift

@@ -267,7 +267,6 @@ actor SDLContextActor {
             quicClient.stop()
         }
         
-        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)")
@@ -278,7 +277,7 @@ actor SDLContextActor {
             }
             
             group.addTask {
-                for try await message in quicClient.messageStream() {
+                for await message in quicClient.messageStream {
                     await self.handleQUICMessage(message: message)
                 }
             }

Tun/Punchnet/Actors/SDLQuicClient.swift

@@ -18,6 +18,7 @@ enum SDLQUICError: Error {
     case timeout
     case decodeError(String)
     case packetTooLarge
+    case dataStreamClosed
 }
 
 enum SDLQUICEvent: Error {
@@ -27,26 +28,28 @@ enum SDLQUICEvent: Error {
 }
 
 final class SDLQUICClient {
-    private let allocator = ByteBufferAllocator()
+    private let frameParser: SDLQUICFrameParser
-    // 最大缓冲区区为2M
-    private let maxBufferSize: Int
     
-    private static let maxPacketSize: Int = 64 * 1024
+    // 数据流
-
+    public var messageStream: AsyncStream<SDLQUICInboundMessage>
-    private let readyLatch = AsyncOneShot<Void>()
+    private let messageCont: AsyncStream<SDLQUICInboundMessage>.Continuation
+    private var readTask: Task<Void, Never>?
 
     // 事件流
     public var eventStream: AsyncStream<SDLQUICEvent>
     private let eventCont: AsyncStream<SDLQUICEvent>.Continuation
     
+    private var isFinished: Bool = false
+
     private let connection: NWConnection
     private let queue = DispatchQueue(label: "com.sdl.QUICClient.queue") // 专用队列保证线程安全
 
     init(host: String, port: UInt16, maxBufferSize: Int = 2 * 1024 * 1024) {
         let options = NWProtocolQUIC.Options(alpn: ["punchnet/1.0"])
+        self.frameParser = SDLQUICFrameParser(maxBufferSize: maxBufferSize)
         
-        self.maxBufferSize = maxBufferSize
         (self.eventStream, self.eventCont) = AsyncStream.makeStream(of: SDLQUICEvent.self)
+        (self.messageStream, self.messageCont) = AsyncStream.makeStream(of: SDLQUICInboundMessage.self)
 
         // 这里设置证书的校验逻辑
         sec_protocol_options_set_verify_block(
@@ -64,31 +67,38 @@ final class SDLQUICClient {
     
     func start() {
         connection.stateUpdateHandler = { [weak self] state in
-            Task {
+            SDLLogger.log("[SDLQUICClient] new state: \(state)", for: .debug)
-                await self?.handleConnectionStateUpdate(state)
+            switch state {
+            case .ready:
+                self?.startReadTask()
+            case .failed(let error):
+                self?.eventCont.yield(.failed(error))
+            case .cancelled:
+                self?.eventCont.yield(.cancelled)
+            default:
+                ()
             }
         }
         connection.start(queue: self.queue)
     }
     
-    private func handleConnectionStateUpdate(_ state: NWConnection.State) async {
+    private func startReadTask() {
-        SDLLogger.log("[SDLQUICClient] new state: \(state)", for: .debug)
+        self.readTask?.cancel()
-        switch state {
+        self.readTask = Task {
-        case .ready:
+            do {
-            await self.readyLatch.succeed(())
+                while !Task.isCancelled {
-        case .failed(let error):
+                    let data = try await self.readOnce()
-            await self.readyLatch.fail(error)
+                    let frames = try self.frameParser.parseFrames(data: data)
-            self.eventCont.yield(.failed(error))
+                    for frame in frames {
-        case .cancelled:
+                        if let message = SDLQUICCodec.decode(frame: frame) {
-            await self.readyLatch.fail(SDLQUICError.connectionCancelled)
+                            self.messageCont.yield(message)
-            self.eventCont.yield(.cancelled)
+                        }
-        default:
+                    }
-            ()
+                }
+            } catch {
+                self.messageCont.finish()
             }
         }
-
-    func waitReady(timeout: Duration = .seconds(5)) async throws {
-        try await self.readyLatch.wait(timeout: timeout, timeoutError: SDLQUICError.timeout)
     }
     
     func send(type: SDLPacketType, data: Data) {
@@ -106,63 +116,51 @@ final class SDLQUICClient {
         })
     }
     
-    func stop() {
+    private func readOnce() async throws -> Data {
-        self.connection.cancel()
+        return try await withCheckedThrowingContinuation { cont in
-        Task {
+            self.connection.receive(minimumIncompleteLength: 1, maximumLength: 64 * 1024) { data, _, isComplete, error in
-            await self.readyLatch.fail(SDLQUICError.connectionCancelled)
-        }
-        self.eventCont.finish()
-    }
-    
-}
-
-// --MARK: Reader
-extension SDLQUICClient {
-    
-    func messageStream() -> AsyncThrowingStream<SDLQUICInboundMessage, Error> {
-        return AsyncThrowingStream { continuation in
-            var buffer = allocator.buffer(capacity: self.maxBufferSize)
-            let threshold = self.maxBufferSize / 10 * 6
-            
-            func readOnce() {
-                self.connection.receive(minimumIncompleteLength: 1, maximumLength: Self.maxPacketSize) { data, _, isComplete, error in
                 if let error {
-                        continuation.finish(throwing: error)
+                    cont.resume(throwing: error)
-                        return
-                    }
-                    
-                    do {
-                        if let data, !data.isEmpty {
-                            buffer.writeBytes(data)
-                            let frames = try Self.parseFrames(buffer: &buffer)
-                            if buffer.readerIndex > threshold {
-                                buffer.discardReadBytes()
-                            }
-                            
-                            for frame in frames {
-                                if let message = Self.decode(frame: frame) {
-                                    continuation.yield(message)
-                                }
-                            }
-                        }
-                    } catch let err {
-                        continuation.finish(throwing: err)
                     return
                 }
                 
                 if isComplete {
-                        continuation.finish()
+                    cont.resume(throwing: SDLQUICError.dataStreamClosed)
                 } else {
-                        readOnce()
+                    cont.resume(returning: data ?? Data())
                 }
             }
         }
-            readOnce()
+    }
-        }
+
+    func stop() {
+        self.readTask?.cancel()
+        self.connection.cancel()
+        self.eventCont.finish()
+        self.messageCont.finish()
+    }
+    
+}
+
+// --MARK: 数据累加器
+extension SDLQUICClient {
+    
+    class SDLQUICFrameParser {
+        private let allocator = ByteBufferAllocator()
+        // 最大缓冲区区为2M
+        private let maxPacketSize: Int = 64 * 1024
+        private let maxBufferSize: Int
+        private var buffer: ByteBuffer
+        
+        init(maxBufferSize: Int) {
+            self.buffer = allocator.buffer(capacity: maxBufferSize)
+            self.maxBufferSize = maxBufferSize
         }
         
         // 尝试解析数据
-    private static func parseFrames(buffer: inout ByteBuffer) throws -> [ByteBuffer] {
+        public func parseFrames(data: Data) throws -> [ByteBuffer] {
+            self.buffer.writeBytes(data)
+            
             guard buffer.readableBytes >= 2 else {
                 return []
             }
@@ -173,7 +171,7 @@ extension SDLQUICClient {
                     break
                 }
                 
-            if len > Self.maxPacketSize {
+                if len > self.maxPacketSize {
                     throw SDLQUICError.packetTooLarge
                 }
                 
@@ -187,15 +185,23 @@ extension SDLQUICClient {
                 }
             }
             
+            let threshold = maxBufferSize / 10 * 6
+            if buffer.readerIndex > threshold {
+                buffer.discardReadBytes()
+            }
+            
             return frames
         }
         
+    }
+    
 }
 
 // --MARK: 编解码器
 extension SDLQUICClient {
     
-    private static func decode(frame: ByteBuffer) -> SDLQUICInboundMessage? {
+    enum SDLQUICCodec {
+        public static func decode(frame: ByteBuffer) -> SDLQUICInboundMessage? {
             var buffer = frame
             guard let type = buffer.readInteger(as: UInt8.self),
                   let packetType = SDLPacketType(rawValue: type) else {
@@ -255,6 +261,7 @@ extension SDLQUICClient {
                 return nil
             }
         }
+    }
 }
 
 // --MARK: quic验证