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 let onEvent: EventHandler
    private let onData: DataHandler

    private var udpHole: SDLUDPHole?
    private var udpHoleMonitorTask: Task<Void, Never>?
    private var natProbeTask: Task<Void, Never>?
    private var localAddress: SocketAddress?

    private var udpHoleV6: SDLUDPHoleV6?
    private var udpHoleV6MonitorTask: Task<Void, Never>?

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

    func start(includeV6: Bool = false) {
        self.startV4()
        if includeV6 {
            self.startV6()
        }
    }

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

        let udpHoleMonitorTask = self.udpHoleMonitorTask
        self.udpHoleMonitorTask = nil

        let natProbeTask = self.natProbeTask
        self.natProbeTask = nil

        udpHoleMonitorTask?.cancel()
        natProbeTask?.cancel()
        await self.proberActor.cancelAll()
        await udpHole?.stop()

        if let natProbeTask {
            await natProbeTask.value
        }
        if let udpHoleMonitorTask {
            await udpHoleMonitorTask.value
        }

        let udpHoleV6 = self.udpHoleV6
        self.udpHoleV6 = nil
        let udpHoleV6MonitorTask = self.udpHoleV6MonitorTask
        self.udpHoleV6MonitorTask = nil
        udpHoleV6MonitorTask?.cancel()
        udpHoleV6?.stop()
        if let udpHoleV6MonitorTask {
            await udpHoleV6MonitorTask.value
        }
    }

    func send(type: SDLPacketType, data: Data, remoteAddress: SocketAddress) async {
        switch remoteAddress {
        case .v4:
            guard let udpHole else {
                SDLLogger.log("[SDLUDPHoleService] 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("[SDLUDPHoleService] udpHoleV6 is nil for remoteAddress: \(remoteAddress)", for: .debug)
                return
            }
            udpHoleV6.send(type: type, data: data, remoteAddress: remoteAddress)
        default:
            SDLLogger.log("[SDLUDPHoleService] unsupported socket family: \(remoteAddress)", for: .debug)
        }
    }

    private func startV4() {
        guard self.udpHoleMonitorTask == nil else {
            return
        }

        self.udpHoleMonitorTask = startMonitorTask(name: "udpHoleServiceV4Monitor") { [weak self] in
            guard let self else {
                throw CancellationError()
            }
            try await self.runV4()
        }
    }

    private func runV4() async throws {
        let udpHole = try SDLUDPHole()
        let localAddress = try await udpHole.start()
        self.udpHole = udpHole
        self.localAddress = localAddress
        SDLLogger.log("[SDLUDPHoleService] udpHole started, on address: \(localAddress)")

        await self.onEvent(.ready(localAddress))
        self.startNatProbe(using: udpHole)

        defer {
            if self.udpHole === udpHole {
                self.udpHole = nil
                self.localAddress = nil
            }
        }

        do {
            try await withTaskCancellationHandler {
                for try await (remoteAddress, message) in await udpHole.messageStream() {
                    try Task.checkCancellation()
                    try await self.handleV4Message(remoteAddress: remoteAddress, message: message)
                }
            } onCancel: {
                Task {
                    await udpHole.stop()
                }
            }
        } catch {
            await udpHole.stop()
            throw error
        }
    }

    private func startNatProbe(using udpHole: SDLUDPHole) {
        self.natProbeTask?.cancel()
        let proberActor = self.proberActor
        let onEvent = self.onEvent
        self.natProbeTask = Task {
            if Task.isCancelled {
                return
            }
            let natType = await proberActor.probeNatType(using: udpHole)
            if Task.isCancelled {
                return
            }
            await 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 startV6() {
        guard self.udpHoleV6MonitorTask == nil else {
            return
        }

        self.udpHoleV6MonitorTask = startMonitorTask(name: "udpHoleServiceV6Monitor") { [weak self] in
            guard let self else {
                throw CancellationError()
            }
            try await self.runV6()
        }
    }

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

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

        defer {
            if self.udpHoleV6 === udpHoleV6 {
                udpHoleV6.stop()
                self.udpHoleV6 = nil
            }
        }

        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("[SDLUDPHoleService] udpHoleV6 ready")
                    case .closed, .errorCaught:
                        throw SDLContextError.udpHoleClosed
                    }
                }
            }

            _ = try await group.next()
        }
    }
}