fix dnsClient

This commit is contained in:
anlicheng 2026-05-05 12:41:16 +08:00
parent e755203013
commit dde40ceb68
2 changed files with 107 additions and 80 deletions

View File

@ -341,49 +341,23 @@ actor SDLContextActor {
SDLLogger.log("[SDLContext] dnsLocalClient started")
self.dnsLocalClient = dnsLocalClient
try await withThrowingTaskGroup { group in
defer {
group.cancelAll()
self.dnsLocalClient = nil
//dnsLocalClient.stop()
}
group.addTask {
try await withTaskCancellationHandler {
//
for await packet in dnsLocalClient.packetFlow {
for try await packet in dnsLocalClient.packetFlow {
try Task.checkCancellation()
// Ip
let nePacket = NEPacket(data: packet, protocolFamily: 2)
self.provider.packetFlow.writePacketObjects([nePacket])
}
throw SDLContextError.dnsLocalClientClosed
}
group.addTask {
for await event in dnsLocalClient.eventStream {
try Task.checkCancellation()
switch event {
case .failed(let error):
SDLLogger.log("[SDLContext] dnsLocalClient failed: \(error)")
throw error
case .cancelled:
SDLLogger.log("[SDLContext] dnsLocalClient cancelled")
throw SDLContextError.dnsLocalClientCancelled
case .sendFailed(let error):
SDLLogger.log("[SDLContext] dnsLocalClient sendFailed: \(error)")
throw error
}
}
}
do {
try await group.next()
} onCancel: {
Task {
await dnsLocalClient.stop()
self.dnsLocalClient = nil
} catch let err {
await dnsLocalClient.stop()
self.dnsLocalClient = nil
throw err
}
}
}

View File

@ -16,11 +16,12 @@ actor DNSLocalClient {
private enum State {
case idle
case starting
case running
case stopped
}
enum Event {
enum DNSLocalError: Error {
case failed(Error)
case cancelled
case sendFailed(Error)
@ -35,16 +36,10 @@ actor DNSLocalClient {
private var cleanupTask: Task<Void, Never>?
private let timeoutInterval: TimeInterval = 3.0
public let packetFlow: AsyncStream<Data>
nonisolated let packetFlow: AsyncThrowingStream<Data, Error>
@ObservationIgnored
private let packetContinuation: AsyncStream<Data>.Continuation
//
public let eventStream: AsyncStream<Event>
@ObservationIgnored
private let eventContinuation: AsyncStream<Event>.Continuation
private let packetContinuation: AsyncThrowingStream<Data, Error>.Continuation
private var isPacketContinuationFinished: Bool = false
private var pendingRequests: [UInt16: PendingRequest] = [:]
private var nextTransactionID: UInt16 = 1
@ -52,16 +47,21 @@ actor DNSLocalClient {
init(host: String) {
self.dnsServerEndpoint = .hostPort(host: NWEndpoint.Host(host), port: 53)
let (stream, continuation) = AsyncStream.makeStream(of: Data.self, bufferingPolicy: .bufferingNewest(256))
let (stream, continuation) = AsyncThrowingStream.makeStream(of: Data.self, bufferingPolicy: .bufferingNewest(256))
self.packetFlow = stream
self.packetContinuation = continuation
let eventPair = AsyncStream.makeStream(of: Event.self)
self.eventStream = eventPair.stream
self.eventContinuation = eventPair.continuation
self.packetContinuation.onTermination = { termination in
SDLLogger.log("[DNSLocalClient] packetFlow terminated: \(termination)")
}
}
func start() {
guard self.state == .idle else {
return
}
self.state = .starting
let parameters = NWParameters.udp
parameters.prohibitedInterfaceTypes = [.other]
// 2. pathSelectionOptions
@ -73,44 +73,47 @@ actor DNSLocalClient {
await self?.handleConnectionStateUpdate(state, for: connection)
}
}
connection.start(queue: .global())
self.connection = connection
self.cleanupTask = Task { [weak self] in
let cleanupTask = Task { [weak self] in
while !Task.isCancelled {
try? await Task.sleep(nanoseconds: 3 * 1_000_000_000)
await self?.performCleanup()
}
}
self.connection = connection
self.cleanupTask = cleanupTask
connection.start(queue: .global())
}
func query(tracker: DNSTracker, dnsPayload: Data) {
guard let connection = self.connection, connection.state == .ready, dnsPayload.count >= 2 else {
let transactionID: UInt16
let connection: NWConnection
let rewrittenPayload: Data
guard self.state != .stopped,
let connection = self.connection, connection.state == .ready, dnsPayload.count >= 2 else {
return
}
guard let transactionID = self.allocateTransactionID() else {
guard let allocatedTransactionID = self.allocateTransactionID() else {
SDLLogger.log("[DNSLocalClient] no available transaction id", for: .debug)
return
}
transactionID = allocatedTransactionID
self.pendingRequests[transactionID] = PendingRequest(tracker: tracker)
let rewrittenPayload = Self.rewriteTransactionID(in: dnsPayload, to: transactionID)
rewrittenPayload = Self.rewriteTransactionID(in: dnsPayload, to: transactionID)
connection.send(content: rewrittenPayload, completion: .contentProcessed { error in
connection.send(content: rewrittenPayload, completion: .contentProcessed { [weak self] error in
if let error {
self.eventContinuation.yield(.sendFailed(error))
Task {
await self.removePendingRequest(forKey: transactionID)
await self?.handleSendFailure(transactionID: transactionID, error: error)
}
}
})
}
private func removePendingRequest(forKey id: UInt16) {
self.pendingRequests.removeValue(forKey: id)
}
func stop() {
guard self.state != .stopped else {
return
@ -118,31 +121,35 @@ actor DNSLocalClient {
self.state = .stopped
self.receiveTask?.cancel()
let receiveTask = self.receiveTask
self.receiveTask = nil
self.connection?.cancel()
let connection = self.connection
self.connection = nil
self.cleanupTask?.cancel()
let cleanupTask = self.cleanupTask
self.cleanupTask = nil
self.pendingRequests.removeAll()
self.nextTransactionID = 1
self.packetContinuation.finish()
receiveTask?.cancel()
connection?.cancel()
cleanupTask?.cancel()
self.finishPacketContinuationIfNeed(throwing: nil)
}
private func handleConnectionStateUpdate(_ state: NWConnection.State, for conn: NWConnection) {
switch state {
case .ready:
if self.markConnectionReady(conn) {
self.startReceiveTask(for: conn)
self.state = .running
}
case .failed(let error):
SDLLogger.log("[DNSLocalClient] failed with error: \(error.localizedDescription)", for: .debug)
self.eventContinuation.yield(.failed(error))
self.finishPacketContinuationIfNeed(throwing: .failed(error))
case .cancelled:
self.eventContinuation.yield(.cancelled)
self.finishPacketContinuationIfNeed(throwing: .cancelled)
default:
()
}
@ -151,7 +158,7 @@ actor DNSLocalClient {
private func startReceiveTask(for conn: NWConnection) {
let stream = Self.makeReceiveStream(for: conn)
self.receiveTask = Task { [weak self] in
let task = Task { [weak self] in
for await data in stream {
guard let self else {
break
@ -159,6 +166,30 @@ actor DNSLocalClient {
await self.handleResponse(data: data)
}
}
let shouldKeepTask = self.state != .stopped && self.isCurrentConnection(conn)
if shouldKeepTask {
self.receiveTask?.cancel()
self.receiveTask = task
}
if !shouldKeepTask {
task.cancel()
}
}
private func finishPacketContinuationIfNeed(throwing error: DNSLocalError?) {
guard !self.isPacketContinuationFinished else {
return
}
self.isPacketContinuationFinished = true
if let error {
self.packetContinuation.finish(throwing: error)
} else {
self.packetContinuation.finish()
}
}
private func handleResponse(data: Data) {
@ -179,6 +210,28 @@ actor DNSLocalClient {
self.packetContinuation.yield(packet)
}
private func markConnectionReady(_ conn: NWConnection) -> Bool {
guard self.state != .stopped, self.isCurrentConnection(conn) else {
return false
}
self.state = .running
return true
}
private func isCurrentConnection(_ conn: NWConnection) -> Bool {
guard let currentConnection = self.connection else {
return false
}
return currentConnection === conn
}
private func handleSendFailure(transactionID: UInt16, error: NWError) {
self.pendingRequests.removeValue(forKey: transactionID)
self.finishPacketContinuationIfNeed(throwing: .sendFailed(error))
}
private func performCleanup() {
let now = Date()
self.pendingRequests = self.pendingRequests.filter { _, request in