punchnet-macos/Tun/Context/SDLSuperControlPlane.swift

//
//  SDLSuperControlPlane.swift
//  Tun
//
//  Created by Codex on 2026/5/28.
//

import Foundation
import NIOCore

enum SuperDecision {
    case updateIPv6Assist(SDLV6Info?)
    case completeRegistration(cipher: CCDataCipher)
    case failTunnel(Error)
    case publishTunnelEvent(code: Int?, message: String)
    case sendSuper(type: SDLPacketType, data: Data)
    case sendPacket(type: SDLPacketType, data: Data, remoteAddress: SocketAddress)
    case resolvePeerInfo(SDLPeerInfo)
    case removeSession(dstMac: Data)
    case requestExposedService
    case shutdown(message: String)
    case applyPolicyResponse(SDLPolicyResponse)
    case applyExposedServiceResponse(SDLExposedServiceResponse)
    case handleARPResponse(SDLArpResponse)
}

actor SDLSuperControlPlane {
    typealias DecisionHandler = @Sendable (SuperDecision) async -> Void

    private let config: SDLConfiguration
    private let rsaCipher: RSACipher
    private var sessionToken: Data?
    private var onDecision: DecisionHandler = { _ in }

    init(config: SDLConfiguration, rsaCipher: RSACipher) {
        self.config = config
        self.rsaCipher = rsaCipher
    }

    func updateDecisionHandler(_ onDecision: @escaping DecisionHandler) {
        self.onDecision = onDecision
    }

    func reset() {
        self.sessionToken = nil
    }

    func handle(_ message: SDLSuperMessage) async {
        switch message {
        case .welcome(let welcome):
            await self.handleWelcome(welcome)
        case .pong:
            ()
        case .registerSuperAck(let registerSuperAck):
            await self.handleRegisterSuperAck(registerSuperAck)
        case .registerSuperNak(let registerSuperNak):
            await self.handleRegisterSuperNak(registerSuperNak)
        case .peerInfo(let peerInfo):
            SDLLogger.log("[SDLContext] peer message: \(peerInfo)")
            await self.onDecision(.resolvePeerInfo(peerInfo))
        case .event(let event):
            await self.handleEvent(event)
        case .policyReponse(let policyResponse):
            await self.onDecision(.applyPolicyResponse(policyResponse))
        case .exposedServiceResponse(let response):
            await self.onDecision(.applyExposedServiceResponse(response))
        case .arpResponse(let arpResponse):
            SDLLogger.log("[SDLContext] get arp response: \(arpResponse)")
            await self.onDecision(.handleARPResponse(arpResponse))
        }
    }

    func sendStunRequest(natType: SDLNATProberActor.NatType, v6Info: SDLV6Info?) async {
        guard let sessionToken else {
            return
        }

        var stunRequest = SDLStunRequest()
        stunRequest.clientID = self.config.clientId
        stunRequest.networkID = self.config.networkAddress.networkId
        stunRequest.ip = self.config.networkAddress.ip
        stunRequest.mac = self.config.networkAddress.mac
        stunRequest.natType = UInt32(natType.rawValue)
        stunRequest.sessionToken = sessionToken

        if let v6Info {
            stunRequest.v6Info = v6Info
        }

        if let stunData = try? stunRequest.serializedData() {
            await self.onDecision(.sendPacket(type: .stunRequest, data: stunData, remoteAddress: self.config.stunSocketAddress))
        }
    }

    private func handleWelcome(_ welcome: SDLWelcome) async {
        SDLLogger.log("[SDLContext] quic welcome: \(welcome)")

        if welcome.hasIpv6Assist {
            await self.onDecision(.updateIPv6Assist(welcome.ipv6Assist))
        } else {
            await self.onDecision(.updateIPv6Assist(nil))
        }

        await self.doRegisterSuper()
        SDLLogger.log("[SDLContext] quic doRegisterSuper")
    }

    private func handleRegisterSuperAck(_ registerSuperAck: SDLRegisterSuperAck) async {
        guard let key = try? self.rsaCipher.decode(data: Data(registerSuperAck.key)) else {
            SDLLogger.log("[SDLContext] registerSuperAck invalid key")
            await self.onDecision(.failTunnel(SDLError.invalidKey))
            return
        }

        let algorithm = registerSuperAck.algorithm.lowercased()
        let regionId = registerSuperAck.regionID
        self.sessionToken = registerSuperAck.sessionToken

        let cipher: CCDataCipher
        switch algorithm {
        case "aes":
            cipher = CCAESChiper(key: key)
        case "chacha20":
            cipher = CCChaCha20Cipher(regionId: regionId, keyData: key)
        default:
            SDLLogger.log("[SDLContext] registerSuperAck invalid algorithm \(algorithm)")
            await self.onDecision(.failTunnel(SDLError.unsupportedAlgorithm(algorithm: algorithm)))
            return
        }

        SDLLogger.log("[SDLContext] registerSuperAck, use algorithm \(algorithm), key len: \(key.count)")
        await self.onDecision(.completeRegistration(cipher: cipher))
    }

    private func handleRegisterSuperNak(_ nakPacket: SDLRegisterSuperNak) async {
        let errorMessage = nakPacket.errorMessage
        guard let errorCode = SDLNAKErrorCode(rawValue: UInt8(nakPacket.errorCode)) else {
            return
        }

        switch errorCode {
        case .invalidToken, .nodeDisabled:
            await self.onDecision(.publishTunnelEvent(code: Int(errorCode.rawValue), message: errorMessage))
            await self.onDecision(.failTunnel(NSError(domain: "com.jihe.punchnet.tun", code: -1)))
        case .noIpAddress, .networkFault, .internalFault:
            await self.onDecision(.publishTunnelEvent(code: Int(errorCode.rawValue), message: errorMessage))
        }

        SDLLogger.log("[SDLContext] Get a SuperNak message exit")
    }

    private func handleEvent(_ event: SDLEvent) async {
        switch event.event {
        case .natChanged(let natChangedEvent):
            let dstMac = natChangedEvent.mac
            SDLLogger.log("[SDLContext] natChangedEvent, dstMac: \(dstMac)")
            await self.onDecision(.removeSession(dstMac: dstMac))
        case .sendRegister(let sendRegisterEvent):
            await self.handleSendRegisterEvent(sendRegisterEvent)
        case .exposedServiceChanged:
            SDLLogger.log("[SDLContext] exposedServiceChanged event")
            await self.onDecision(.requestExposedService)
        case .shutdown(let shutdownEvent):
            await self.onDecision(.shutdown(message: shutdownEvent.message))
        case .none:
            ()
        }
    }

    private func handleSendRegisterEvent(_ event: SDLEvent.SendRegister) async {
        var register = SDLRegister()
        register.networkID = self.config.networkAddress.networkId
        register.srcMac = self.config.networkAddress.mac
        register.dstMac = event.dstMac

        guard let registerData = try? register.serializedData() else {
            return
        }

        SDLLogger.log("[SDLContext] sendRegisterEvent, ip: \(event)")

        if event.natIp > 0 && event.natPort > 0 {
            let address = SDLUtil.int32ToIp(event.natIp)
            if let remoteAddress = try? SocketAddress(ipAddress: address, port: Int(event.natPort)) {
                await self.onDecision(.sendPacket(type: .register, data: registerData, remoteAddress: remoteAddress))
            }
        }

        if event.hasV6Info, let remoteAddress = try? await event.v6Info.socketAddress() {
            await self.onDecision(.sendPacket(type: .register, data: registerData, remoteAddress: remoteAddress))
        }
    }

    private func doRegisterSuper() async {
        var registerSuper = SDLRegisterSuper()
        registerSuper.clientID = self.config.clientId
        registerSuper.networkID = self.config.networkAddress.networkId
        registerSuper.mac = self.config.networkAddress.mac
        registerSuper.ip = self.config.networkAddress.ip
        registerSuper.maskLen = UInt32(self.config.networkAddress.maskLen)
        registerSuper.hostname = self.config.hostname
        registerSuper.pubKey = self.rsaCipher.pubKey
        registerSuper.accessToken = self.config.accessToken

        if let registerSuperData = try? registerSuper.serializedData() {
            SDLLogger.log("[SDLContext] will send register super")
            await self.onDecision(.sendSuper(type: .registerSuper, data: registerSuperData))
        }
    }
}