punchnet-macos/Tun/UDPHole/SDLUDPHoleService.swift

import Foundation
import NIOCore

enum SDLUDPHoleKind: Equatable {
    case v4
    case v6

    func convertAddressType() -> Session.AddressType {
        switch self {
        case .v4:
            return .v4
        case .v6:
            return .v6
        }
    }
}

actor SDLUDPHoleService {
    enum Event {
        case ready(SocketAddress)
        case natType(SDLNATProberActor.NatType)
        case packet(SocketAddress, SDLHoleControlMessage, source: SDLUDPHoleKind)
        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

    init(proberActor: SDLNATProberActor) {
        self.proberActor = proberActor
    }

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

    func run(includeV6: Bool = false) async throws {
        let generation = self.nextGeneration()
        let session = SDLUDPHoleSession(
            proberActor: self.proberActor,
            includeV6: includeV6,
            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)
        } 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()
        await self.proberActor.cancelAll()
    }

    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 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)
    }
}

actor SDLUDPHoleSession {
    private let proberActor: SDLNATProberActor
    private let includeV6: Bool
    private let onEvent: SDLUDPHoleService.EventHandler
    private let onData: SDLUDPHoleService.DataHandler

    private var udpHole: SDLUDPHole?
    private var udpHoleV6: SDLUDPHoleV6?
    private var localAddress: SocketAddress?

    init(
        proberActor: SDLNATProberActor,
        includeV6: Bool,
        onEvent: @escaping SDLUDPHoleService.EventHandler,
        onData: @escaping SDLUDPHoleService.DataHandler
    ) {
        self.proberActor = proberActor
        self.includeV6 = includeV6
        self.onEvent = onEvent
        self.onData = onData
    }

    func run() async throws {
        do {
            try await withThrowingTaskGroup(of: Void.self) { group in
                defer {
                    group.cancelAll()
                }

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

                if self.includeV6 {
                    group.addTask {
                        try await self.runV6()
                    }
                }

                try await group.waitForAll()
            }
            await self.stop()
        } catch {
            await self.stop()
            throw error
        }
    }

    func stop() async {
        let udpHole = self.udpHole
        self.udpHole = nil
        self.localAddress = nil

        let udpHoleV6 = self.udpHoleV6
        self.udpHoleV6 = nil

        await self.proberActor.cancelAll()
        await udpHole?.stop()
        udpHoleV6?.stop()
    }

    func send(type: SDLPacketType, data: Data, remoteAddress: SocketAddress) async {
        switch remoteAddress {
        case .v4:
            guard let udpHole else {
                SDLLogger.log("[SDLUDPHoleSession] udpHole is nil for remoteAddress: \(remoteAddress)", for: .debug)
                return
            }

            await udpHole.send(type: type, data: data, remoteAddress: remoteAddress)
        case .v6:
            guard let udpHoleV6 else {
                SDLLogger.log("[SDLUDPHoleSession] udpHoleV6 is nil for remoteAddress: \(remoteAddress)", for: .debug)
                return
            }

            udpHoleV6.send(type: type, data: data, remoteAddress: remoteAddress)
        default:
            SDLLogger.log("[SDLUDPHoleSession] unsupported socket family: \(remoteAddress)", for: .debug)
        }
    }

    private func runV4() async throws {
        let udpHole = try SDLUDPHole()
        let localAddress = try await udpHole.start()
        self.udpHole = udpHole
        self.localAddress = localAddress

        SDLLogger.log("[SDLUDPHoleSession] udpHole started, on address: \(localAddress)")
        await self.onEvent(.ready(localAddress))

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

                    group.addTask {
                        try await self.readV4Loop(udpHole: udpHole)
                    }

                    group.addTask {
                        await self.probeNatType(udpHole: udpHole)
                    }

                    try await group.waitForAll()
                }
            } onCancel: {
                Task {
                    await udpHole.stop()
                }
            }
        } catch {
            await udpHole.stop()
            if self.udpHole === udpHole {
                self.udpHole = nil
                self.localAddress = nil
            }
            throw error
        }
    }

    private func readV4Loop(udpHole: SDLUDPHole) async throws {
        for try await datagram in await udpHole.messageStream() {
            try Task.checkCancellation()
            try await self.handleV4Message(remoteAddress: datagram.remoteAddress, message: datagram.message)
        }
    }

    private func probeNatType(udpHole: SDLUDPHole) async {
        if Task.isCancelled {
            return
        }

        let natType = await self.proberActor.probeNatType(using: udpHole)
        if Task.isCancelled {
            return
        }

        await self.onEvent(.natType(natType))
    }

    private func handleV4Message(remoteAddress: SocketAddress, message: SDLHoleMessage) async throws {
        switch message {
        case .control(let control):
            switch control {
            case .stunProbeReply(let probeReply):
                await self.proberActor.handleProbeReply(localAddress: self.localAddress, reply: probeReply)
            default:
                await self.onEvent(.packet(remoteAddress, control, source: .v4))
            }
        case .data(let data):
            await self.onData(data)
        }
    }

    private func runV6() async throws {
        let udpHoleV6 = try SDLUDPHoleV6()
        let localAddress = try udpHoleV6.start()
        self.udpHoleV6 = udpHoleV6

        if let localAddress {
            SDLLogger.log("[SDLUDPHoleSession] udpHoleV6 started, on address: \(localAddress)")
        } else {
            SDLLogger.log("[SDLUDPHoleSession] udpHoleV6 started, no local address")
        }

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

                let onEvent = self.onEvent
                let onData = self.onData
                group.addTask {
                    for await (remoteAddress, message) in udpHoleV6.messageStream {
                        try Task.checkCancellation()
                        switch message {
                        case .control(let control):
                            await onEvent(.packet(remoteAddress, control, source: .v6))
                        case .data(let data):
                            await onData(data)
                        }
                    }
                }

                group.addTask {
                    for await event in udpHoleV6.eventStream {
                        try Task.checkCancellation()
                        switch event {
                        case .ready:
                            SDLLogger.log("[SDLUDPHoleSession] udpHoleV6 ready")
                        case .closed, .errorCaught:
                            throw SDLContextError.udpHoleClosed
                        }
                    }
                }

                _ = try await group.next()
            }
        } catch {
            udpHoleV6.stop()
            if self.udpHoleV6 === udpHoleV6 {
                self.udpHoleV6 = nil
            }
            throw error
        }
    }
}