punchnet-macos/Tun/UDPHole/SDLUDPHoleService.swift

import Foundation
import NIOCore

actor SDLUDPHoleService {
    enum Event {
        case ready(SocketAddress)
        case natType(SDLNATProberActor.NatType)
        case packet(SocketAddress, SDLHoleControlMessage)
        case closed(Error)
    }

    typealias EventHandler = @Sendable (Event) async -> Void
    typealias DataHandler = @Sendable (SDLData) async -> Void

    private let proberActor: SDLNATProberActor

    private var onEvent: EventHandler = { _ in }
    private var onData: DataHandler = { _ in }
    private var currentSession: SDLUDPHoleSession?
    private var generation: UInt64 = 0
    private var isRunning = false
    private var isStopping = false
    private var needsImmediateRestart = false
    private let retryDelay: Duration

    init(proberActor: SDLNATProberActor, retryDelay: Duration = .seconds(5)) {
        self.proberActor = proberActor
        self.retryDelay = retryDelay
    }

    func updateHandlers(onEvent: @escaping EventHandler, onData: @escaping DataHandler) {
        self.onEvent = onEvent
        self.onData = onData
    }

    func run() async throws {
        guard !self.isRunning else {
            return
        }

        self.isRunning = true
        self.isStopping = false

        defer {
            self.isRunning = false
            self.currentSession = nil
        }

        while !Task.isCancelled, !self.isStopping {
            let generation = self.nextGeneration()
            let session = SDLUDPHoleSession(
                proberActor: self.proberActor,
                onEvent: { [weak self] event in
                    await self?.handleEvent(event, generation: generation)
                },
                onData: self.onData
            )

            self.currentSession = session

            do {
                try await session.run()
                self.clearCurrent(session, generation: generation)
                await session.stop()

                guard !self.isStopping else {
                    break
                }

                SDLLogger.log("[SDLUDPHoleService] session ended, will restart", category: .udpHole)
            } catch is CancellationError {
                self.clearCurrent(session, generation: generation)
                await session.stop()
                throw CancellationError()
            } catch {
                self.clearCurrent(session, generation: generation)
                await session.stop()

                guard !self.isStopping else {
                    break
                }

                SDLLogger.log("[SDLUDPHoleService] session failed: \(error.localizedDescription), will restart", category: .udpHole)
            }

            if self.consumeImmediateRestartRequest() {
                SDLLogger.log("[SDLUDPHoleService] session invalidated after wakeup, will restart immediately", category: .udpHole)
                continue
            }

            try await Task.sleep(for: self.retryDelay)
        }
    }

    func stop() async {
        self.isStopping = true
        self.needsImmediateRestart = false
        await self.invalidateCurrentSession()
    }

    func recoverAfterWake() async {
        guard !self.isStopping else {
            return
        }

        self.needsImmediateRestart = self.currentSession != nil
        await self.invalidateCurrentSession()
    }

    func send(type: SDLPacketType, data: Data, remoteAddress: SocketAddress) async {
        await self.currentSession?.send(type: type, data: data, remoteAddress: remoteAddress)
    }

    private func nextGeneration() -> UInt64 {
        self.generation &+= 1
        return self.generation
    }

    private func invalidateCurrentSession() async {
        self.generation &+= 1

        let session = self.currentSession
        self.currentSession = nil

        await session?.stop()
        await self.proberActor.cancelAll()
    }

    private func consumeImmediateRestartRequest() -> Bool {
        let needsImmediateRestart = self.needsImmediateRestart
        self.needsImmediateRestart = false
        return needsImmediateRestart
    }

    private func clearCurrent(_ session: SDLUDPHoleSession, generation: UInt64) {
        guard self.generation == generation else {
            return
        }

        if self.currentSession === session {
            self.currentSession = nil
        }
    }

    private func handleEvent(_ event: Event, generation: UInt64) async {
        guard self.generation == generation else {
            return
        }

        await self.onEvent(event)
    }
}