punchnet-macos/Tun/Super/SDLSuperService.swift

import Foundation

actor SDLSuperService {
    typealias MessageHandler = @Sendable (SDLQUICInboundMessage) async -> Void

    private let serverEndpoint: SDLConfiguration.ResolvedServerEndpoint
    private let port: UInt16

    private var onMessage: MessageHandler = { _ in }
    private var currentSession: SDLSuperSession?
    private var generation: UInt64 = 0

    init(serverEndpoint: SDLConfiguration.ResolvedServerEndpoint, port: UInt16 = 1443) {
        self.serverEndpoint = serverEndpoint
        self.port = port
    }

    func updateMessageHandler(_ onMessage: @escaping MessageHandler) {
        self.onMessage = onMessage
    }

    func run() async throws {
        let generation = self.nextGeneration()
        let session = SDLSuperSession(
            serverEndpoint: self.serverEndpoint,
            port: self.port,
            onMessage: { [weak self] message in
                await self?.handleMessage(message, generation: generation)
            }
        )

        self.currentSession = session

        do {
            try await session.run()
            self.clearCurrent(session, generation: generation)
        } catch is CancellationError {
            self.clearCurrent(session, generation: generation)
            await session.stop()
            throw CancellationError()
        } catch {
            self.clearCurrent(session, generation: generation)
            await session.stop()
            throw error
        }
    }

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

        let session = self.currentSession
        self.currentSession = nil

        await session?.stop()
    }

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

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

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

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

    private func handleMessage(_ message: SDLQUICInboundMessage, generation: UInt64) async {
        guard self.generation == generation else {
            return
        }

        await self.onMessage(message)
    }
}

final class SDLSuperSession: @unchecked Sendable {
    typealias MessageHandler = @Sendable (SDLQUICInboundMessage) async -> Void

    private let serverEndpoint: SDLConfiguration.ResolvedServerEndpoint
    private let port: UInt16
    private let onMessage: MessageHandler
    private let client: SDLSuperClient

    init(serverEndpoint: SDLConfiguration.ResolvedServerEndpoint, port: UInt16, onMessage: @escaping MessageHandler) {
        self.serverEndpoint = serverEndpoint
        self.port = port
        self.onMessage = onMessage
        self.client = SDLSuperClient(serverEndpoint: serverEndpoint, port: port)
    }

    func run() async throws {
        do {
            try await withTaskCancellationHandler {
                try await self.runLoops()
            } onCancel: {
                Task {
                    await self.client.stop()
                }
            }

            await self.stop()
        } catch {
            await self.stop()
            throw error
        }
    }

    func stop() async {
        await self.client.stop()
    }

    func send(type: SDLPacketType, data: Data) async {
        await self.client.send(type: type, data: data)
    }

    private func runLoops() async throws {
        SDLLogger.log("[SDLSuperSession] start super client: \(self.serverEndpoint.ip)")

        try await withThrowingTaskGroup(of: Void.self) { group in
            defer {
                group.cancelAll()
            }

            group.addTask {
                try await self.client.run()
            }

            group.addTask {
                try await self.readLoop()
            }

            group.addTask {
                try await self.pingLoop()
            }

            _ = try await group.next()
        }
    }

    private func readLoop() async throws {
        for try await message in self.client.messageStream {
            try Task.checkCancellation()
            await self.onMessage(message)
        }
    }

    private func pingLoop() async throws {
        while true {
            try await Task.sleep(for: .seconds(5))
            try Task.checkCancellation()
            await self.client.send(type: .ping, data: Data())
        }
    }
}