使用Netty处理WebSocket请求

使用Netty处理WebSocket请求

前言

前面我根据网上的一些资料，整理了一个使用Netty处理WebSocket的案例代码出来SpringBoot整合Netty处理WebSocket（支持url参数），后来觉得那份代码不是最优解，因此，我决定重新写了以下代码，更好的处理URL和请求头。

依赖

<dependency>
  <groupId>io.netty</groupId>
  <artifactId>netty-all</artifactId>
  <version>4.1.51.Final</version>
</dependency>

服务端代码

@Slf4j
public class WebSocketServer implements Closeable {

  private EventLoopGroup bossGroup;
  private EventLoopGroup workGroup;
  private ChannelFuture future;

  public WebSocketServer(int port, String websocketPath) throws InterruptedException {
    log.info("Netty WebSocket server starting...");
    try {
      this.bossGroup = new NioEventLoopGroup(2);
      this.workGroup = new NioEventLoopGroup(2);

      ServerBootstrap bootstrap = new ServerBootstrap();
      bootstrap.group(bossGroup, workGroup)
          .option(ChannelOption.SO_BACKLOG, 1 << 10)
          .channel(NioServerSocketChannel.class)
          .localAddress(port)
          .childHandler(new ChannelInitializer<SocketChannel>() {
            @Override
            protected void initChannel(SocketChannel channel) {
              channel.pipeline()
                  .addLast(new HttpServerCodec())
                  .addLast(new ChunkedWriteHandler())
                  .addLast(new HttpObjectAggregator(1 << 13))
                  .addLast(new WebSocketServerProtocolHandler(websocketPath, true))
                  .addLast(new WebSocketMessageHandler(websocketPath));
            }
          });

      future = bootstrap.bind().sync();
    } finally {
      if (future != null && future.isSuccess())
        log.info("Netty WebSocket server started.");
      else
        close();
    }
  }

  @Override
  public void close() {
    log.info("Netty WebSocket server closing...");
    try {
      if (future != null)
        future = future.channel().closeFuture();
      if (bossGroup != null)
        bossGroup.shutdownGracefully().sync();
      if (workGroup != null)
        workGroup.shutdownGracefully().sync();
      if (future != null)
        future.sync();
    } catch (InterruptedException e) {
      log.error(e.getMessage(), e);
    }
    log.info("Netty WebSocket server closed.");
  }
}

连接、消息处理

@Slf4j
@RequiredArgsConstructor
public class WebSocketMessageHandler extends SimpleChannelInboundHandler<WebSocketFrame> {

  private final String websocketPath;
  private ChannelHandlerContext ctx;

  @Override
  public void channelActive(ChannelHandlerContext ctx) throws Exception {
    super.channelActive(ctx);
    this.ctx = ctx;
    log.info("客户端连接：{}", ctx.channel().id().asLongText());
  }

  @Override
  public void channelInactive(ChannelHandlerContext ctx) throws Exception {
    super.channelInactive(ctx);
    log.info("客户端断开：{}", ctx.channel().id().asLongText());
  }

  @Override
  public void userEventTriggered(ChannelHandlerContext ctx, Object evt) throws Exception {
    super.userEventTriggered(ctx, evt);
    if (evt instanceof WebSocketServerProtocolHandler.HandshakeComplete) {
      WebSocketServerProtocolHandler.HandshakeComplete complete = (WebSocketServerProtocolHandler.HandshakeComplete) evt;
      String uri = complete.requestUri(); // 获取请求路径，接着根据自己的业务逻辑判断以及获取自己想要的参数
      String regex = "^" + websocketPath + "/(\\w+)$";
      if (!uri.matches(regex)) {
        log.info("请求路径不合法：{}", uri);
        ctx.close();
        return;
      }
      log.info("用户名：{}", uri.replaceAll(regex, "$1"));
    }
  }

  @Override
  protected void channelRead0(ChannelHandlerContext ctx, WebSocketFrame msg) {
    if (msg instanceof TextWebSocketFrame) {
      TextWebSocketFrame frame = (TextWebSocketFrame) msg;
      log.info("文本消息内容：{}", frame.text());
    } else if (msg instanceof BinaryWebSocketFrame) {
//      BinaryWebSocketFrame frame = (BinaryWebSocketFrame) msg;
      log.info("接收到二进制消息");
    }
  }

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

  public void writeAndFlush(String msg) {
    this.ctx.writeAndFlush(new TextWebSocketFrame(msg));
  }
}

程序演示

public static void main(String[] args) throws InterruptedException {
  WebSocketServer server = new WebSocketServer(8012, "/ws");
  Runtime.getRuntime().addShutdownHook(new Thread(server::close));
}