fix Local DNS

Tun/Punchnet/DNS/DNSLocalClient.swift

@@ -22,6 +22,7 @@ actor DNSLocalClient {
     
     private var state: State = .idle
     private var connections: [NWConnection] = []
+    private var receiveTasks: [ObjectIdentifier: Task<Void, Never>] = [:]
     private let dnsServers = ["223.5.5.5", "119.29.29.29"]
     
     let packetFlow: AsyncStream<Data>
@@ -35,7 +36,7 @@ actor DNSLocalClient {
     private var didFinishPacketFlow = false
     
     init() {
-        let (stream, continuation) = AsyncStream.makeStream(of: Data.self, bufferingPolicy: .unbounded)
+        let (stream, continuation) = AsyncStream.makeStream(of: Data.self, bufferingPolicy: .bufferingNewest(256))
         self.packetFlow = stream
         self.packetContinuation = continuation
     }
@@ -104,6 +105,8 @@ actor DNSLocalClient {
         }
         
         self.state = .stopped
+        self.receiveTasks.values.forEach { $0.cancel() }
+        self.receiveTasks.removeAll()
         self.connections.forEach { $0.cancel() }
         self.connections.removeAll()
         
@@ -122,11 +125,13 @@ actor DNSLocalClient {
         
         switch state {
         case .ready:
-            self.receiveLoop(for: conn)
+            self.startReceiveTask(for: conn)
         case .failed(let error):
             SDLLogger.log("[DNSLocalClient] failed with error: \(error.localizedDescription)", for: .debug)
             self.stop()
         case .cancelled:
+            let key = ObjectIdentifier(conn)
+            self.receiveTasks.removeValue(forKey: key)?.cancel()
             self.connections.removeAll { $0 === conn }
             if self.connections.isEmpty {
                 self.stop()
@@ -136,28 +141,31 @@ actor DNSLocalClient {
         }
     }
     
-    private func receiveLoop(for conn: NWConnection) {
+    private func startReceiveTask(for conn: NWConnection) {
-        conn.receiveMessage { [weak self] content, _, _, error in
+        let key = ObjectIdentifier(conn)
-            Task {
+        guard self.receiveTasks[key] == nil else {
-                await self?.handleReceive(content: content, error: error, for: conn)
-            }
-        }
-    }
-    
-    private func handleReceive(content: Data?, error: NWError?, for conn: NWConnection) {
-        guard case .running = self.state else {
             return
         }
         
-        if let data = content {
+        let stream = Self.makeReceiveStream(for: conn)
-            self.handleResponse(data: data)
+        self.receiveTasks[key] = Task { [weak self] in
-        }
+            guard let self else {
+                return
+            }
             
-        if error == nil && conn.state == .ready {
+            for await data in stream {
-            self.receiveLoop(for: conn)
+                await self.handleResponse(data: data)
+            }
+            
+            await self.didFinishReceiving(for: conn)
         }
     }
     
+    private func didFinishReceiving(for conn: NWConnection) {
+        let key = ObjectIdentifier(conn)
+        self.receiveTasks.removeValue(forKey: key)
+    }
+    
     private func handleResponse(data: Data) {
         guard case .running = self.state,
               let rewrittenTransactionID = Self.readTransactionID(from: data),
@@ -235,6 +243,26 @@ actor DNSLocalClient {
         rewrittenPayload[1] = UInt8(transactionID & 0xFF)
         return rewrittenPayload
     }
+    
+    private static func makeReceiveStream(for conn: NWConnection) -> AsyncStream<Data> {
+        return AsyncStream(bufferingPolicy: .bufferingNewest(256)) { continuation in
+            func receiveNext() {
+                conn.receiveMessage { content, _, _, error in
+                    if let data = content, !data.isEmpty {
+                        continuation.yield(data)
+                    }
+                    
+                    if error == nil && conn.state == .ready {
+                        receiveNext()
+                    } else {
+                        continuation.finish()
+                    }
+                }
+            }
+            
+            receiveNext()
+        }
+    }
 }
 
 extension DNSLocalClient {