ai-security-xdr/haobang-security-dm/syslog-serve/src/main/java/com/netty/SyslogServerBoth.java

package com.netty;

import io.netty.bootstrap.Bootstrap;
import io.netty.bootstrap.ServerBootstrap;
import io.netty.channel.ChannelFuture;
import io.netty.channel.ChannelOption;
import io.netty.channel.EventLoopGroup;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.nio.NioDatagramChannel;
import io.netty.channel.socket.nio.NioServerSocketChannel;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

/**
 * Syslog 服务器主类，支持 TCP 和 UDP 协议
 */
public class SyslogServerBoth {
    private static final Logger logger = LoggerFactory.getLogger(SyslogServer.class);

    private final int tcpPort;
    private final int udpPort;

    // 共享的EventLoopGroup以优化资源使用
    private EventLoopGroup bossGroup;
    private EventLoopGroup workerGroup;
    private EventLoopGroup udpGroup;

    private ChannelFuture tcpChannelFuture;
    private ChannelFuture udpChannelFuture;

    private final ExecutorService serverExecutor = Executors.newFixedThreadPool(2);

    public SyslogServerBoth(int tcpPort, int udpPort) {
        this.tcpPort = tcpPort;
        this.udpPort = udpPort;
    }

    public void start() throws InterruptedException {
        logger.info("Starting Syslog server with TCP port {} and UDP port {}", tcpPort, udpPort);

        // 创建共享的EventLoopGroup
        bossGroup = new NioEventLoopGroup(1);
        workerGroup = new NioEventLoopGroup();
        udpGroup = new NioEventLoopGroup();

        final CountDownLatch latch = new CountDownLatch(2);

        try {
            // 启动UDP服务器（使用独立的线程）
            serverExecutor.submit(() -> {
                try {
                    startUdpServer();
                    latch.countDown();
                } catch (Exception e) {
                    logger.error("UDP server failed to start", e);
                    latch.countDown();
                    throw new RuntimeException(e);
                }
            });

            // 启动TCP服务器（使用独立的线程）
            serverExecutor.submit(() -> {
                try {
                    startTcpServer();
                    latch.countDown();
                } catch (Exception e) {
                    logger.error("TCP server failed to start", e);
                    latch.countDown();
                    throw new RuntimeException(e);
                }
            });

            // 等待两个服务器都启动完成
            latch.await();
            logger.info("Both TCP and UDP Syslog servers are running");

            // 注册关闭钩子
            Runtime.getRuntime().addShutdownHook(new Thread(this::shutdown));

            // 保持主线程运行
            Thread.currentThread().join();

        } catch (Exception e) {
            logger.error("Failed to start servers", e);
            shutdown();
        }
    }

    private void startTcpServer() throws InterruptedException {
        ServerBootstrap b = new ServerBootstrap();
        b.group(bossGroup, workerGroup)
                .channel(NioServerSocketChannel.class)
                .childHandler(new SyslogTcpChannelInitializer())
                .option(ChannelOption.SO_BACKLOG, 128)
                .childOption(ChannelOption.SO_KEEPALIVE, true);

        tcpChannelFuture = b.bind(tcpPort).sync();
        logger.info("TCP Syslog server started on port {}", tcpPort);

        // 添加关闭监听器
        tcpChannelFuture.channel().closeFuture().addListener(future -> {
            if (future.isSuccess()) {
                logger.info("TCP server channel closed");
            }
        });
    }

    private void startUdpServer() throws InterruptedException {
        Bootstrap b = new Bootstrap();
        b.group(udpGroup)
                .channel(NioDatagramChannel.class)
                .handler(new SyslogUdpChannelInitializer())
                .option(ChannelOption.SO_BROADCAST, true);

        udpChannelFuture = b.bind(udpPort).sync();
        logger.info("UDP Syslog server started on port {}", udpPort);

        // 添加关闭监听器
        udpChannelFuture.channel().closeFuture().addListener(future -> {
            if (future.isSuccess()) {
                logger.info("UDP server channel closed");
            }
        });
    }

    /**
     * 优雅关闭服务器
     */
    public void shutdown() {
        logger.info("Shutting down Syslog servers...");

        try {
            // 关闭TCP服务器
            if (tcpChannelFuture != null) {
                tcpChannelFuture.channel().close().sync();
            }

            // 关闭UDP服务器
            if (udpChannelFuture != null) {
                udpChannelFuture.channel().close().sync();
            }

            // 关闭EventLoopGroup
            if (bossGroup != null) {
                bossGroup.shutdownGracefully().sync();
            }
            if (workerGroup != null) {
                workerGroup.shutdownGracefully().sync();
            }
            if (udpGroup != null) {
                udpGroup.shutdownGracefully().sync();
            }

            // 关闭线程池
            serverExecutor.shutdown();

            logger.info("Syslog servers shutdown complete");
        } catch (Exception e) {
            logger.error("Error during shutdown", e);
        }
    }

    /**
     * 检查服务器是否在运行
     */
    public boolean isRunning() {
        return (tcpChannelFuture != null && tcpChannelFuture.channel().isActive()) ||
                (udpChannelFuture != null && udpChannelFuture.channel().isActive());
    }

    public static void main(String[] args) throws Exception {
        int tcpPort = 514;
        int udpPort = 514;

        // 解析命令行参数
        if (args.length >= 1) {
            tcpPort = Integer.parseInt(args[0]);
        }
        if (args.length >= 2) {
            udpPort = Integer.parseInt(args[1]);
        }

        // 如果两个端口相同，UDP端口自动+1避免冲突
        if (tcpPort == udpPort) {
            udpPort = tcpPort + 1;
            logger.warn("TCP and UDP ports cannot be the same. UDP port changed to {}", udpPort);
        }

        SyslogServerBoth server = new SyslogServerBoth(tcpPort, udpPort);

        // 注册全局异常处理器
        Thread.setDefaultUncaughtExceptionHandler((t, e) -> {
            logger.error("Uncaught exception in thread {}", t.getName(), e);
        });

        server.start();
    }
}