From c7f216e810530bfb10e3ca53dd25c9a779c45228 Mon Sep 17 00:00:00 2001
From: anlicheng <244108715@qq.com>
Date: Thu, 28 May 2026 16:03:04 +0800
Subject: [PATCH] fix dns

---
 Tun/DNS/DNSCloudService.swift | 17 ++---------------
 Tun/DNS/DNSLocalClient.swift  | 36 +++++++++++++++++++----------------
 Tun/DNS/DNSLocalService.swift | 21 ++++----------------
 3 files changed, 26 insertions(+), 48 deletions(-)

diff --git a/Tun/DNS/DNSCloudService.swift b/Tun/DNS/DNSCloudService.swift
index b0266ad..ac9a937 100644
--- a/Tun/DNS/DNSCloudService.swift
+++ b/Tun/DNS/DNSCloudService.swift
@@ -5,7 +5,6 @@ actor DNSCloudService {
     private var onEvent: DNSEventHandler = { _ in }
 
     private var currentClient: DNSCloudClient?
-    private var generation: UInt64 = 0
 
     init(serverIP: String) {
         self.serverIP = serverIP
@@ -16,12 +15,11 @@ actor DNSCloudService {
     }
 
     func run() async throws {
-        let generation = self.nextGeneration()
         let client = DNSCloudClient(serverIP: self.serverIP, port: 15353)
         self.currentClient = client
 
         defer {
-            self.clearCurrent(client, generation: generation)
+            self.clearCurrent(client)
         }
 
         do {
@@ -57,8 +55,6 @@ actor DNSCloudService {
     }
 
     func stop() async {
-        self.generation &+= 1
-
         let client = self.currentClient
         self.currentClient = nil
 
@@ -69,16 +65,7 @@ actor DNSCloudService {
         await self.currentClient?.forward(ipPacketData: ipPacketData)
     }
 
-    private func nextGeneration() -> UInt64 {
-        self.generation &+= 1
-        return self.generation
-    }
-
-    private func clearCurrent(_ client: DNSCloudClient, generation: UInt64) {
-        guard self.generation == generation else {
-            return
-        }
-
+    private func clearCurrent(_ client: DNSCloudClient) {
         if self.currentClient === client {
             self.currentClient = nil
         }
diff --git a/Tun/DNS/DNSLocalClient.swift b/Tun/DNS/DNSLocalClient.swift
index e9f420f..98c615d 100644
--- a/Tun/DNS/DNSLocalClient.swift
+++ b/Tun/DNS/DNSLocalClient.swift
@@ -18,6 +18,7 @@ actor DNSLocalClient {
         case failed(Error)
         case cancelled
         case sendFailed(Error)
+        case invalidData
     }
     
     private let queue = DispatchQueue(label: "com.sdl.DNSCloudClient.queue")
@@ -82,10 +83,9 @@ actor DNSLocalClient {
                 
                 group.addTask {
                     try await self.readySignal.wait()
-                    
-                    let stream = Self.makeReceiveStream(for: self.connection)
-                    for await data in stream {
+                    while true {
                         try Task.checkCancellation()
+                        let data = try await self.readOnce()
                         await self.handleResponse(data: data)
                     }
                 }
@@ -241,23 +241,27 @@ actor DNSLocalClient {
         return rewrittenPayload
     }
     
-    private static func makeReceiveStream(for connection: NWConnection) -> AsyncStream<Data> {
-        return AsyncStream(bufferingPolicy: .bufferingNewest(256)) { continuation in
-            func receiveNext() {
-                connection.receiveMessage { content, _, _, error in
-                    if let data = content, !data.isEmpty {
-                        continuation.yield(data)
-                    }
-                    
-                    if error == nil && connection.state == .ready {
-                        receiveNext()
+    private func readOnce() async throws -> Data {
+        guard self.connection.state == .ready else {
+            throw DNSLocalError.cancelled
+        }
+        
+        let readContinuation = OnceContinuation<Data, Error>()
+        return try await withTaskCancellationHandler {
+            try await withCheckedThrowingContinuation { cont in
+                readContinuation.set(cont)
+                self.connection.receiveMessage { content, _, _, error in
+                    if let error {
+                        readContinuation.resume(throwing: error)
+                    } else if let data = content, !data.isEmpty {
+                        readContinuation.resume(returning: data)
                     } else {
-                        continuation.finish()
+                        readContinuation.resume(throwing: DNSLocalError.invalidData)
                     }
                 }
             }
-            
-            receiveNext()
+        } onCancel: {
+            readContinuation.resume(throwing: CancellationError())
         }
     }
     
diff --git a/Tun/DNS/DNSLocalService.swift b/Tun/DNS/DNSLocalService.swift
index 8719622..2706766 100644
--- a/Tun/DNS/DNSLocalService.swift
+++ b/Tun/DNS/DNSLocalService.swift
@@ -5,7 +5,6 @@ actor DNSLocalService {
     private var onEvent: DNSEventHandler = { _ in }
 
     private var currentClient: DNSLocalClient?
-    private var generation: UInt64 = 0
 
     init(publicDnsServers: [String]) {
         self.publicDnsServers = publicDnsServers
@@ -16,7 +15,6 @@ actor DNSLocalService {
     }
 
     func run() async throws {
-        let generation = self.nextGeneration()
         let dnsServer = self.publicDnsServers.randomElement() ?? "223.5.5.5"
         let client = DNSLocalClient(host: dnsServer)
         self.currentClient = client
@@ -25,22 +23,20 @@ actor DNSLocalService {
 
         do {
             try await self.run(client: client)
-            self.clearCurrent(client, generation: generation)
+            self.clearCurrent(client)
             await client.stop()
         } catch is CancellationError {
-            self.clearCurrent(client, generation: generation)
+            self.clearCurrent(client)
             await client.stop()
             throw CancellationError()
         } catch {
-            self.clearCurrent(client, generation: generation)
+            self.clearCurrent(client)
             await client.stop()
             throw error
         }
     }
 
     func stop() async {
-        self.generation &+= 1
-
         let client = self.currentClient
         self.currentClient = nil
 
@@ -74,16 +70,7 @@ actor DNSLocalService {
         }
     }
 
-    private func nextGeneration() -> UInt64 {
-        self.generation &+= 1
-        return self.generation
-    }
-
-    private func clearCurrent(_ client: DNSLocalClient, generation: UInt64) {
-        guard self.generation == generation else {
-            return
-        }
-
+    private func clearCurrent(_ client: DNSLocalClient) {
         if self.currentClient === client {
             self.currentClient = nil
         }