punchnet-macos/Tun/Punchnet/Actors/SDLUDPHole.swift

//
//  SDLUDPHoleActor 2.swift
//  punchnet
//
//  Created by 安礼成 on 2026/1/28.
//


//
//  SDLanServer.swift
//  Tun
//
//  Created by 安礼成 on 2024/1/31.
//

import Foundation
import NIOCore
import NIOPosix
import SwiftProtobuf

// 处理和sn-server服务器之间的通讯
final class SDLUDPHole {
    private let group = MultiThreadedEventLoopGroup(numberOfThreads: 1)
    private var channel: Channel?
    
    public let eventStream: AsyncStream<UDPHoleEvent>
    private let eventContinuation: AsyncStream<UDPHoleEvent>.Continuation
    private let logger: SDLLogger
    
    enum UDPHoleEvent {
        case ready
        case message(SocketAddress, SDLHoleInboundMessage)
    }
    
    // 启动函数
    init(logger: SDLLogger) throws {
        self.logger = logger
        (self.eventStream, self.eventContinuation) = AsyncStream.makeStream(of: UDPHoleEvent.self, bufferingPolicy: .unbounded)
    }
    
    func start() throws {
        let bootstrap = DatagramBootstrap(group: group)
            .channelOption(ChannelOptions.socketOption(.so_reuseaddr), value: 1)
            .channelInitializer { channel in
                channel.pipeline.addHandler(SDLUDPHoleHandler(eventContinuation: self.eventContinuation, logger: self.logger))
            }
        
        self.channel = try bootstrap.bind(host: "0.0.0.0", port: 0).wait()
        self.logger.log("[UDPHole] started", level: .debug)
    }
    
    func getLocalAddress() -> SocketAddress? {
        return self.channel?.localAddress
    }
    
    // MARK: 处理写入逻辑
    func send(type: SDLPacketType, data: Data, remoteAddress: SocketAddress) {
        guard let channel = self.channel else {
            return
        }
        
        var buffer = channel.allocator.buffer(capacity: data.count + 1)
        buffer.writeBytes([type.rawValue])
        buffer.writeBytes(data)
        
        let envelope = AddressedEnvelope<ByteBuffer>(remoteAddress: remoteAddress, data: buffer)
        channel.eventLoop.execute {
            channel.writeAndFlush(envelope, promise: nil)
        }
    }
    
    deinit {
        try? self.group.syncShutdownGracefully()
        self.eventContinuation.finish()
        self.channel?.close(promise: nil)
    }
    
}

extension SDLUDPHole {
    
    final class SDLUDPHoleHandler: ChannelInboundHandler {
        typealias InboundIn = AddressedEnvelope<ByteBuffer>
        
        private var eventContinuation: AsyncStream<UDPHoleEvent>.Continuation
        private var logger: SDLLogger
        
        // --MARK: ChannelInboundHandler delegate
        
        init(eventContinuation: AsyncStream<UDPHoleEvent>.Continuation, logger: SDLLogger) {
            self.eventContinuation = eventContinuation
            self.logger = logger
        }
        
        func channelActive(context: ChannelHandlerContext) {
            self.eventContinuation.yield(.ready)
        }
        
        func channelRead(context: ChannelHandlerContext, data: NIOAny) {
            let envelope = unwrapInboundIn(data)
            
            var buffer = envelope.data
            let remoteAddress = envelope.remoteAddress
            do {
                if let message = try decode(buffer: &buffer) {
                    self.eventContinuation.yield(.message(remoteAddress, message))
                } else {
                    self.logger.log("[SDLUDPHole] decode message, get null", level: .warning)
                }
            } catch let err {
                self.logger.log("[SDLUDPHole] decode message, get error: \(err)", level: .warning)
            }
        }
        
        func channelInactive(context: ChannelHandlerContext) {
            self.eventContinuation.finish()
        }
        
        func errorCaught(context: ChannelHandlerContext, error: any Error) {
            context.close(promise: nil)
        }
        
        // --MARK: 编解码器
        private func decode(buffer: inout ByteBuffer) throws -> SDLHoleInboundMessage? {
            guard let type = buffer.readInteger(as: UInt8.self),
                  let packetType = SDLPacketType(rawValue: type),
                  let bytes = buffer.readBytes(length: buffer.readableBytes) else {
                return nil
            }
            
            switch packetType {
            case .data:
                let dataPacket = try SDLData(serializedBytes: bytes)
                return .data(dataPacket)
            case .register:
                let registerPacket = try SDLRegister(serializedBytes: bytes)
                return .register(registerPacket)
            case .registerAck:
                let registerAck = try SDLRegisterAck(serializedBytes: bytes)
                return .registerAck(registerAck)
            case .stunProbeReply:
                let stunProbeReply = try SDLStunProbeReply(serializedBytes: bytes)
                return .stunProbeReply(stunProbeReply)
            case .registerSuperAck:
                guard let bytes = buffer.readBytes(length: buffer.readableBytes),
                      let registerSuperAck = try? SDLRegisterSuperAck(serializedBytes: bytes) else {
                    return nil
                }
                return .registerSuperAck(registerSuperAck)
            case .registerSuperNak:
                guard let bytes = buffer.readBytes(length: buffer.readableBytes),
                      let registerSuperNak = try? SDLRegisterSuperNak(serializedBytes: bytes) else {
                    return nil
                }
                return .registerSuperNak(registerSuperNak)
                
            case .peerInfo:
                guard let bytes = buffer.readBytes(length: buffer.readableBytes),
                      let peerInfo = try? SDLPeerInfo(serializedBytes: bytes) else {
                    return nil
                }
                
                return .peerInfo(peerInfo)
            case .event:
                guard let eventVal = buffer.readInteger(as: UInt8.self),
                      let event = SDLEventType(rawValue: eventVal),
                      let bytes = buffer.readBytes(length: buffer.readableBytes) else {
                    return nil
                }
                
                switch event {
                case .natChanged:
                    guard let natChangedEvent = try? SDLNatChangedEvent(serializedBytes: bytes) else {
                        return nil
                    }
                    return .event(.natChanged(natChangedEvent))
                case .sendRegister:
                    guard let sendRegisterEvent = try? SDLSendRegisterEvent(serializedBytes: bytes) else {
                        return nil
                    }
                    return .event(.sendRegister(sendRegisterEvent))
                case .networkShutdown:
                    guard let networkShutdownEvent = try? SDLNetworkShutdownEvent(serializedBytes: bytes) else {
                        return nil
                    }
                    return .event(.networkShutdown(networkShutdownEvent))
                }
                
            default:
                return nil
            }
        }
        
    }
    
}