用 Netty 写一个 WebSocket 服务端

2020-08-31

这是一个用 Netty 写的 WebSocket 服务端栗子，同时结合了 Spring Boot，仅供参考。

Step 1

首先添加相关依赖，核心依赖为 Netty：

<dependency>
    <groupId>io.netty</groupId>
    <artifactId>netty-all</artifactId>
    <version>4.1.6.Final</version>
</dependency>
<dependency>
    <groupId>com.fasterxml.jackson.core</groupId>
    <artifactId>jackson-databind</artifactId>
    <version>2.9.7</version>
</dependency>
<dependency>
    <groupId>log4j</groupId>
    <artifactId>log4j</artifactId>
    <version>1.2.17</version>
</dependency>

Step 2

编写 Netty WebSocket Server 主类：（实际使用时把注释去掉）

import io.netty.bootstrap.ServerBootstrap;
import io.netty.channel.ChannelFuture;
import io.netty.channel.ChannelInitializer;
import io.netty.channel.ChannelOption;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.SocketChannel;
import io.netty.channel.socket.nio.NioServerSocketChannel;
import io.netty.handler.codec.http.HttpObjectAggregator;
import io.netty.handler.codec.http.HttpServerCodec;
import io.netty.handler.codec.http.websocketx.WebSocketServerProtocolHandler;
import io.netty.handler.codec.http.websocketx.extensions.compression.WebSocketServerCompressionHandler;
import io.netty.handler.ssl.SslHandler;
import io.netty.handler.stream.ChunkedWriteHandler;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.scheduling.annotation.Async;

import javax.net.ssl.SSLContext;
import javax.net.ssl.SSLEngine;

@Component
public class WebSocketServer {

    private static final Logger LOGGER = LoggerFactory.getLogger(WebSocketServer.class);

    @Value("${websocket.server.port}")
    private int port;

    @Async
    public void start() {
        // 定义一对线程组
        // 主线程组，用于接受客户端的连接
        NioEventLoopGroup bossGroup = new NioEventLoopGroup();
        // 从线程组，负责 IO 交互工作
        NioEventLoopGroup group = new NioEventLoopGroup();

        try {
            // 创建 netty 服务器和辅助工具类，并进行相关配置
            ServerBootstrap bootstrap = new ServerBootstrap();
            bootstrap.option(ChannelOption.SO_BACKLOG, 1024)
                    // 绑定两个线程组
                    .group(group, bossGroup)
                    .childOption(ChannelOption.TCP_NODELAY, true)
                    .childOption(ChannelOption.SO_KEEPALIVE, true)
                    // 指定 NIO 的模式
                    .channel(NioServerSocketChannel.class)
                    // 绑定端口
                    .localAddress(this.port)
                    // 子处理器，用于处理 WorkerGroup
                    .childHandler(new ChannelInitializer<SocketChannel>() {
                        @Override
                        protected void initChannel(SocketChannel ch) throws Exception {
                            // 如果要使用 WSS 协议，则需添加 SSL/TLS 的配置
                            SSLContext sslContext = SslUtil.createSSLContext(
                                    "JKS", "N:\\temp\\server.keystore", "123456");
                            SSLEngine sslEngine = sslContext.createSSLEngine();
                            sslEngine.setUseClientMode(false);
                            sslEngine.setNeedClientAuth(false);
                            ch.pipeline().addLast("ssl", new SslHandler(sslEngine))
                                    // WebSocket 基于 HTTP 协议，所以要有 HTTP 的编、解码器，服务端用的是 HttpServerCodec
                                    .addLast(new HttpServerCodec())
                                    // 对写大数据流的支持
                                    .addLast(new ChunkedWriteHandler())
                                    /**
                                     * 我们通常接收到的是一个 HTTP 片段，如果想完整接收一次请求的所有数据，
                                     * 我们需要绑定 HttpObjectAggregator，对 HttpMessage 进行聚合，聚
                                     * 合成 FullHttpRequest 或 FullHttpResponse，然后我们就可以收到一
                                     * 个 FullHttpRequest，它表示一个完整的请求信息。在 Netty 中的编程，
                                     * 几乎都会使用到此 Handler。
                                     */
                                    .addLast(new HttpObjectAggregator(8192))

                                    // == 以上配置跟 HTTP 协议相关，以下配置则跟 httpWebsocket 相关 ==

                                    .addLast(new WebSocketServerCompressionHandler())
                                    /**
                                     * WebSocket 服务器处理的协议，用于指定给客户端连接访问的路由 : /road
                                     * 本 Handler 会帮你处理一些繁重的复杂的事，比如会帮你处理握手动作：
                                     * handshaking（close, ping, pong） ping + pong = 心跳
                                     * 对于 WebSocket 来讲，都是以 frames 进行传输的，不同的数据类型
                                     * 对应的 frames 不同。
                                     */
                                    .addLast(new WebSocketServerProtocolHandler(
                                            "/road", null, true, 65536 * 10))
                                    // 添加自定义的 Handler
                                    .addLast(new TextWebSocketHandler())
                                    .addLast(new BinaryWebSocketHandler())
                            ;
                        }
                    });
            // 启动 server，并且设置端口号（前面已绑定，因此此处可不填），同时设置启动方式为同步
            ChannelFuture future = bootstrap.bind().sync();
            // 监听关闭的 channel，也设置为同步
            future.channel().closeFuture().sync();

        } catch (InterruptedException e) {
            LOGGER.error(e.getMessage());
        } finally {
            try {
                group.shutdownGracefully().sync();
                bossGroup.shutdownGracefully().sync();
            } catch (InterruptedException e) {
                LOGGER.error(e.getMessage());
            }
        }
    }
}

Step 3

添加 SSL 协议的工具类：

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.net.ssl.KeyManagerFactory;
import javax.net.ssl.SSLContext;
import java.io.FileInputStream;
import java.io.IOException;
import java.security.*;
import java.security.cert.CertificateException;

public class SslUtil {

    private static final Logger LOGGER = LoggerFactory.getLogger(SslUtil.class);

    public static SSLContext createSSLContext(String type, String path, String password) {
        try {
            KeyStore keyStore = KeyStore.getInstance(type);
            FileInputStream fileInputStream = new FileInputStream(path);
            keyStore.load(fileInputStream, password.toCharArray());
            KeyManagerFactory factory = KeyManagerFactory.getInstance(
                KeyManagerFactory.getDefaultAlgorithm());
            factory.init(keyStore, password.toCharArray());
            SSLContext context = SSLContext.getInstance("TLS");
            context.init(factory.getKeyManagers(), null, null);
            return context;

        } catch (KeyStoreException
                | CertificateException
                | NoSuchAlgorithmException
                | IOException
                | UnrecoverableKeyException
                | KeyManagementException e) {
            LOGGER.error(e.getMessage());
        }
        return null;
    }
}

添加自定义的 Handler：

import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.SimpleChannelInboundHandler;
import io.netty.handler.codec.DecoderException;
import io.netty.handler.codec.http.websocketx.TextWebSocketFrame;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class TextWebSocketHandler extends SimpleChannelInboundHandler<TextWebSocketFrame> {

    private static final Logger LOGGER = LoggerFactory.getLogger(TextWebSocketHandler.class);

    @Override
    public void channelActive(ChannelHandlerContext ctx) throws Exception {
        AppConstant.WEB_CLIENT_ONLINE.add(ctx.channel());
    }

    @Override
    public void channelInactive(ChannelHandlerContext ctx) throws Exception {
        AppConstant.WEB_CLIENT_ONLINE.remove(ctx.channel());
    }

    @Override
    public void channelReadComplete(ChannelHandlerContext ctx) throws Exception {
        super.channelReadComplete(ctx);
    }

    @Override
    public void userEventTriggered(ChannelHandlerContext ctx, Object evt) throws Exception {
        super.userEventTriggered(ctx, evt);
    }

    @Override
    protected void channelRead0(ChannelHandlerContext ctx, TextWebSocketFrame msg) throws Exception {
        LOGGER.info("收到来自客户端的消息为（text message）： " + msg.text());
        // 向客户端发送数据
        AppConstant.WEB_CLIENT_ONLINE.writeAndFlush(new TextWebSocketFrame("我是服务器！"));
    }

    @Override
    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
        if (cause instanceof DecoderException)
            return;
        LOGGER.error(cause.getMessage());
    }
}

import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.SimpleChannelInboundHandler;
import io.netty.handler.codec.http.websocketx.BinaryWebSocketFrame;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class BinaryWebSocketHandler extends SimpleChannelInboundHandler<BinaryWebSocketFrame> {

    private static final Logger LOGGER = LoggerFactory.getLogger(BinaryWebSocketHandler.class);

    @Override
    public void channelActive(ChannelHandlerContext ctx) throws Exception {
        AppConstant.WEB_CLIENT_ONLINE.add(ctx.channel());
    }

    @Override
    public void channelInactive(ChannelHandlerContext ctx) throws Exception {
        AppConstant.WEB_CLIENT_ONLINE.remove(ctx.channel());
    }

    @Override
    public void channelReadComplete(ChannelHandlerContext ctx) throws Exception {
        super.channelReadComplete(ctx);
    }

    @Override
    public void userEventTriggered(ChannelHandlerContext ctx, Object evt) throws Exception {
        super.userEventTriggered(ctx, evt);
    }

    @Override
    protected void channelRead0(ChannelHandlerContext ctx, BinaryWebSocketFrame msg) throws Exception {
        LOGGER.info("Receive binary message.");
        LOGGER.info("收到来自客户端的消息为（binary message）： " + msg.content());
    }

    @Override
    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
        LOGGER.error(cause.getMessage());
    }
}

其中：

1	public static final ChannelGroup WEB_CLIENT_ONLINE = new DefaultChannelGroup(GlobalEventExecutor.INSTANCE);

Step 4

启动 Server：

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationListener;
import org.springframework.context.event.ContextRefreshedEvent;
import org.springframework.stereotype.Component;

@Component
public class StartupEvent implements ApplicationListener<ContextRefreshedEvent> {

    private static final Logger LOGGER = LoggerFactory.getLogger(StartupEvent.class);

    @Override
    public void onApplicationEvent(ContextRefreshedEvent event) {
        LOGGER.info("Application started, refresh Context.");
        ApplicationContext context = event.getApplicationContext();
        context.getBean(WebSocketServer.class).start();
    }
}