punchnet-macos/Tun/Punchnet/Context/SDLPacketReaderService.swift
2026-05-08 10:28:44 +08:00

99 lines
2.9 KiB
Swift

import Foundation
import NetworkExtension
actor SDLPacketReaderService {
enum Event {
case packet(IPPacket)
}
typealias EventHandler = @Sendable (Event) async -> Void
private typealias PacketReadResult = (packets: [Data], protocols: [NSNumber])?
private final class CancellationToken: @unchecked Sendable {
private let lock = NSLock()
private var cancelled: Bool = false
var isCancelled: Bool {
lock.lock()
defer {
lock.unlock()
}
return cancelled
}
func cancel() {
lock.lock()
cancelled = true
lock.unlock()
}
}
private let provider: NEPacketTunnelProvider
private let onEvent: EventHandler
private var readTask: Task<Void, Never>?
private var readToken: CancellationToken?
init(provider: NEPacketTunnelProvider, onEvent: @escaping EventHandler) {
self.provider = provider
self.onEvent = onEvent
}
func start() {
guard self.readTask == nil else {
return
}
let provider = self.provider
let onEvent = self.onEvent
let readToken = CancellationToken()
self.readToken = readToken
self.readTask = Task(priority: .high) {
while !Task.isCancelled && !readToken.isCancelled {
guard let batch = await Self.readPackets(from: provider) else {
break
}
if Task.isCancelled || readToken.isCancelled {
break
}
for (data, number) in zip(batch.packets, batch.protocols) where number.int32Value == 2 {
if Task.isCancelled || readToken.isCancelled {
break
}
if let packet = IPPacket(data) {
await onEvent(.packet(packet))
}
}
}
SDLLogger.log("[SDLPacketReaderService] readTask finished")
}
}
func stop() {
let readTask = self.readTask
self.readTask = nil
let readToken = self.readToken
self.readToken = nil
readToken?.cancel()
readTask?.cancel()
}
private static func readPackets(from provider: NEPacketTunnelProvider) async -> PacketReadResult {
let readContinuation = OnceContinuation<PacketReadResult, Never>()
return await withTaskCancellationHandler {
await withCheckedContinuation { continuation in
readContinuation.set(continuation)
provider.packetFlow.readPackets { packets, protocols in
readContinuation.resume(returning: (packets: packets, protocols: protocols))
}
}
} onCancel: {
readContinuation.resume(returning: nil)
}
}
}