307 lines
9.9 KiB
Swift
307 lines
9.9 KiB
Swift
import Foundation
|
||
import Network
|
||
|
||
actor DNSLocalClient {
|
||
|
||
struct DNSTracker {
|
||
let transactionID: UInt16
|
||
let clientIP: UInt32
|
||
let clientPort: UInt16
|
||
let createdAt: Date
|
||
}
|
||
|
||
private struct PendingRequest {
|
||
let tracker: DNSTracker
|
||
}
|
||
|
||
private enum State {
|
||
case idle
|
||
case running
|
||
case stopped
|
||
}
|
||
|
||
enum Event {
|
||
case failed(Error)
|
||
case cancelled
|
||
case sendFailed(Error)
|
||
}
|
||
|
||
private var state: State = .idle
|
||
|
||
private let dnsServerEndpoint: NWEndpoint
|
||
private var connection: NWConnection?
|
||
private var receiveTask: Task<Void, Never>?
|
||
|
||
private var cleanupTask: Task<Void, Never>?
|
||
private let timeoutInterval: TimeInterval = 3.0
|
||
|
||
|
||
|
||
public let packetFlow: AsyncStream<Data>
|
||
@ObservationIgnored
|
||
private let packetContinuation: AsyncStream<Data>.Continuation
|
||
|
||
// 事件处理
|
||
public let eventStream: AsyncStream<Event>
|
||
@ObservationIgnored
|
||
private let eventContinuation: AsyncStream<Event>.Continuation
|
||
|
||
private var pendingRequests: [UInt16: PendingRequest] = [:]
|
||
private var nextTransactionID: UInt16 = 1
|
||
|
||
init(host: String) {
|
||
self.dnsServerEndpoint = .hostPort(host: NWEndpoint.Host(host), port: 53)
|
||
|
||
let (stream, continuation) = AsyncStream.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
|
||
}
|
||
|
||
func start() {
|
||
let parameters = NWParameters.udp
|
||
parameters.prohibitedInterfaceTypes = [.other]
|
||
// 2. 增强健壮性：启用多路径切换（替代 pathSelectionOptions 的意图）
|
||
parameters.multipathServiceType = .handover
|
||
|
||
let connection = NWConnection(to: self.dnsServerEndpoint, using: parameters)
|
||
connection.stateUpdateHandler = { [weak self] state in
|
||
Task {
|
||
await self?.handleConnectionStateUpdate(state, for: connection)
|
||
}
|
||
}
|
||
connection.start(queue: .global())
|
||
self.connection = connection
|
||
|
||
self.cleanupTask = Task { [weak self] in
|
||
while !Task.isCancelled {
|
||
try? await Task.sleep(nanoseconds: 3 * 1_000_000_000)
|
||
await self?.performCleanup()
|
||
}
|
||
}
|
||
}
|
||
|
||
func query(tracker: DNSTracker, dnsPayload: Data) {
|
||
guard let connection = self.connection, connection.state == .ready, dnsPayload.count >= 2 else {
|
||
return
|
||
}
|
||
|
||
guard let transactionID = self.allocateTransactionID() else {
|
||
SDLLogger.log("[DNSLocalClient] no available transaction id", for: .debug)
|
||
return
|
||
}
|
||
|
||
self.pendingRequests[transactionID] = PendingRequest(tracker: tracker)
|
||
let rewrittenPayload = Self.rewriteTransactionID(in: dnsPayload, to: transactionID)
|
||
|
||
connection.send(content: rewrittenPayload, completion: .contentProcessed { error in
|
||
if let error {
|
||
self.eventContinuation.yield(.sendFailed(error))
|
||
Task {
|
||
await self.removePendingRequest(forKey: transactionID)
|
||
}
|
||
}
|
||
})
|
||
}
|
||
|
||
private func removePendingRequest(forKey id: UInt16) {
|
||
self.pendingRequests.removeValue(forKey: id)
|
||
}
|
||
|
||
func stop() {
|
||
guard self.state != .stopped else {
|
||
return
|
||
}
|
||
|
||
self.state = .stopped
|
||
|
||
self.receiveTask?.cancel()
|
||
self.receiveTask = nil
|
||
|
||
self.connection?.cancel()
|
||
self.connection = nil
|
||
|
||
self.cleanupTask?.cancel()
|
||
self.cleanupTask = nil
|
||
|
||
self.pendingRequests.removeAll()
|
||
self.nextTransactionID = 1
|
||
|
||
self.packetContinuation.finish()
|
||
}
|
||
|
||
private func handleConnectionStateUpdate(_ state: NWConnection.State, for conn: NWConnection) {
|
||
switch state {
|
||
case .ready:
|
||
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))
|
||
case .cancelled:
|
||
self.eventContinuation.yield(.cancelled)
|
||
default:
|
||
()
|
||
}
|
||
}
|
||
|
||
private func startReceiveTask(for conn: NWConnection) {
|
||
let stream = Self.makeReceiveStream(for: conn)
|
||
|
||
self.receiveTask = Task { [weak self] in
|
||
for await data in stream {
|
||
guard let self else {
|
||
break
|
||
}
|
||
await self.handleResponse(data: data)
|
||
}
|
||
}
|
||
}
|
||
|
||
private func handleResponse(data: Data) {
|
||
guard let rewrittenTransactionID = Self.readTransactionID(from: data),
|
||
let pendingRequest = self.pendingRequests.removeValue(forKey: rewrittenTransactionID) else {
|
||
return
|
||
}
|
||
|
||
let restoredPayload = Self.rewriteTransactionID(in: data, to: pendingRequest.tracker.transactionID)
|
||
|
||
let packet = Self.createDNSResponse(
|
||
payload: restoredPayload,
|
||
srcIP: DNSHelper.dnsDestIpAddr,
|
||
srcPort: 53,
|
||
destIP: pendingRequest.tracker.clientIP,
|
||
destPort: pendingRequest.tracker.clientPort
|
||
)
|
||
self.packetContinuation.yield(packet)
|
||
}
|
||
|
||
private func performCleanup() {
|
||
let now = Date()
|
||
self.pendingRequests = self.pendingRequests.filter { _, request in
|
||
now.timeIntervalSince(request.tracker.createdAt) < self.timeoutInterval
|
||
}
|
||
}
|
||
|
||
private func allocateTransactionID() -> UInt16? {
|
||
var candidate = self.nextTransactionID == 0 ? 1 : self.nextTransactionID
|
||
let start = candidate
|
||
|
||
repeat {
|
||
if self.pendingRequests[candidate] == nil {
|
||
self.nextTransactionID = Self.nextTransactionID(after: candidate)
|
||
return candidate
|
||
}
|
||
|
||
candidate = Self.nextTransactionID(after: candidate)
|
||
} while candidate != start
|
||
|
||
return nil
|
||
}
|
||
|
||
private static func nextTransactionID(after id: UInt16) -> UInt16 {
|
||
return id == UInt16.max ? 1 : id &+ 1
|
||
}
|
||
|
||
private static func readTransactionID(from payload: Data) -> UInt16? {
|
||
guard payload.count >= 2 else {
|
||
return nil
|
||
}
|
||
|
||
return UInt16(payload[0]) << 8 | UInt16(payload[1])
|
||
}
|
||
|
||
private static func rewriteTransactionID(in payload: Data, to transactionID: UInt16) -> Data {
|
||
guard payload.count >= 2 else {
|
||
return payload
|
||
}
|
||
|
||
var rewrittenPayload = payload
|
||
rewrittenPayload[0] = UInt8((transactionID >> 8) & 0xFF)
|
||
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 {
|
||
static func createDNSResponse(payload: Data, srcIP: UInt32, srcPort: UInt16, destIP: UInt32, destPort: UInt16) -> Data {
|
||
let udpLen = 8 + payload.count
|
||
let ipLen = 20 + udpLen
|
||
|
||
var ipHeader = Data(count: 20)
|
||
ipHeader[0] = 0x45
|
||
ipHeader[2...3] = withUnsafeBytes(of: UInt16(ipLen).bigEndian) { Data($0) }
|
||
ipHeader[8] = 64
|
||
ipHeader[9] = 17
|
||
|
||
ipHeader[12...15] = withUnsafeBytes(of: srcIP.bigEndian) { Data($0) }
|
||
ipHeader[16...19] = withUnsafeBytes(of: destIP.bigEndian) { Data($0) }
|
||
|
||
let ipChecksum = calculateChecksum(data: ipHeader)
|
||
ipHeader[10...11] = withUnsafeBytes(of: ipChecksum.bigEndian) { Data($0) }
|
||
|
||
var udpHeader = Data(count: 8)
|
||
udpHeader[0...1] = withUnsafeBytes(of: srcPort.bigEndian) { Data($0) }
|
||
udpHeader[2...3] = withUnsafeBytes(of: destPort.bigEndian) { Data($0) }
|
||
udpHeader[4...5] = withUnsafeBytes(of: UInt16(udpLen).bigEndian) { Data($0) }
|
||
udpHeader[6...7] = Data([0, 0])
|
||
|
||
var packet = Data(capacity: ipLen)
|
||
packet.append(ipHeader)
|
||
packet.append(udpHeader)
|
||
packet.append(payload)
|
||
|
||
return packet
|
||
}
|
||
|
||
static func calculateChecksum(data: Data) -> UInt16 {
|
||
var sum: UInt32 = 0
|
||
let count = data.count
|
||
|
||
data.withUnsafeBytes { (ptr: UnsafeRawBufferPointer) in
|
||
guard let baseAddress = ptr.baseAddress else { return }
|
||
|
||
let wordCount = count / 2
|
||
let words = baseAddress.bindMemory(to: UInt16.self, capacity: wordCount)
|
||
|
||
for i in 0..<wordCount {
|
||
sum += UInt32(UInt16(bigEndian: words[i]))
|
||
}
|
||
|
||
if count % 2 != 0 {
|
||
let lastByte = ptr[count - 1]
|
||
sum += UInt32(lastByte) << 8
|
||
}
|
||
}
|
||
|
||
while (sum >> 16) != 0 {
|
||
sum = (sum & 0xffff) + (sum >> 16)
|
||
}
|
||
|
||
return UInt16(~sum & 0xffff)
|
||
}
|
||
}
|