【工作问题】Netty与WebSocket奇遇记之“熟识”

9527威先僧
1,173 阅读7分钟

「这是我参与2022首次更文挑战的第12天，活动详情查看：2022首次更文挑战」。

序

上篇简略的分析了下WebSocket http握手的流程，分析了下我的需求点，决定在握手阶段插入用户的登录效验。没看过上篇博客的点击传送门直达：Netty与WebSocket奇遇记之“初遇”。

后端我是使用的Netty，所以今天这篇文章就是根据真实的需求，使用Netty搭建一个WebSocket服务器，同时在连接握手阶段插入身份验证。 步骤安排：从源码入手，分析Netty是如何处理WebSocket的请求的，然后根据需求编写相应的业务处理Handler，最后运行分析结果。

正文

一个WebSocket连接的建立需要经过以下几步：

  1. 握手请求阶段：前端发起http握手请求，同时申请协议升级
  2. 握手响应：后端返回握手响应，是握手成功还是失败
  3. 连接建立，开始数据传输

Netty对WebSocket的连接过程以及保活心跳和WebSocket各种生命周期事件都做了相应封装，这个类就是WebSocketServerProtocolHandler。我们先来看看该类上的注释，译文我写在下面了

This handler does all the heavy lifting for you to run a websocket server. It takes care of websocket handshaking as well as processing of control frames (Close, Ping, Pong). Text and Binary data frames are passed to the next handler in the pipeline (implemented by you) for processing. See io.netty.example.http.websocketx.html5.WebSocketServer for usage. The implementation of this handler assumes that you just want to run a websocket server and not process other types HTTP requests (like GET and POST). 这个处理程序为您运行 websocket 服务器完成了所有繁重的工作。它负责 websocket 握手以及控制帧（Close、Ping、Pong）的处理。文本和二进制数据帧被传递给管道中的下一个处理程序（由您实现）进行处理。有关用法，请参见io.netty.example.http.websocketx.html5.WebSocketServer 。此处理程序的实现假定您只想运行 websocket 服务器而不处理其他类型的 HTTP 请求（如 GET 和 POST）

看了上面的翻译，已经大致知道这个类是干啥的了，下面我们分析下代码，然后看看这个类具体做了些什么东西。 使用过netty的人都知道netty需要配置相应的handler（编解码，业务处理handler之类的），下面的代码是ChannelPipeline配置handler处理链的类。

public class WebSocketServerInitializer extends ChannelInitializer<SocketChannel> {
    public static final String WEBSOCKET_PATH = "/alert";
    @Autowired
    private WebSocketHandler webSocketHandler;
    @Autowired
    private AuthHandler authHandler;
    @Override
    protected void initChannel(SocketChannel socketChannel) throws Exception {
        ChannelPipeline pipeline = socketChannel.pipeline();
        //http解码器
        pipeline.addLast(new HttpServerCodec());
        //设置单次请求的文件的大小
        pipeline.addLast(new HttpObjectAggregator(1024*1024*10));
        //对写大数据流的支持
        pipeline.addLast(new ChunkedWriteHandler());
        pipeline.addLast(authHandler);
        //webscoket 服务器处理的协议，用于指定给客户端连接访问的路由
        pipeline.addLast(new WebSocketServerProtocolHandler(WEBSOCKET_PATH, true));
        //自定义handler(作用类似controller,客户端和服务器端之间发消息都在这个自定义handler里面)
        pipeline.addLast(webSocketHandler);
    }
}

WebSocketServerProtocolHandler源码如下，我进行了精简，只保留了相关代码

public class WebSocketServerProtocolHandler extends WebSocketProtocolHandler {
    // 省略若干。。。。。
    @Override
    public void handlerAdded(ChannelHandlerContext ctx) {
        ChannelPipeline cp = ctx.pipeline();
        if (cp.get(WebSocketServerProtocolHandshakeHandler.class) == null) {
            // Add the WebSocketHandshakeHandler before this one.
            cp.addBefore(ctx.name(), WebSocketServerProtocolHandshakeHandler.class.getName(),
                    new WebSocketServerProtocolHandshakeHandler(serverConfig));
        }
        if (serverConfig.decoderConfig().withUTF8Validator() && cp.get(Utf8FrameValidator.class) == null) {
            // Add the UFT8 checking before this one.
            cp.addBefore(ctx.name(), Utf8FrameValidator.class.getName(),
                    new Utf8FrameValidator());
        }
    }
    // 省略若干。。。。。
}

分析：该类的handlerAdded方法在WebSocketServerProtocolHandler加入到Pipeline时被调用，主要作用是在WebSocketServerProtocolHandler的前面加入了握手处理handler（WebSocketServerProtocolHandshakeHandler）、还有个不重要的啥Utf8FrameValidator。我们关心的就是WebSocketServerProtocolHandshakeHandler。它主要是处理握手相关需求的。画一个Pipeline说明图就是这样的。

websocket熟识.png

我在图中将自己定义的handler和程序自动添加的handler做了标记区分。

netty处理websocket整个连接过程的逻辑大概是这个样子的。

由于握手是使用的http，所以需要加入httpServerCodec等相关处理http请求和WebSocketServerProtocolHandshakeHandler处理握手过程的handler，然后在握手成功后移除这些handler，加入处理WebSocket协议的handler。

下面我们看看WebSocketServerProtocolHandshakeHandler是如何处理WebSocket的握手的。下面我粘贴了握手handler的channelRead方法，这也是处理握手请求和构造握手响应的方法。我在代码中关键位置进行了注释。

@Override
    public void channelRead(final ChannelHandlerContext ctx, Object msg) throws Exception {
        final FullHttpRequest req = (FullHttpRequest) msg;
        // 1-判断url是否正确
        if (!isWebSocketPath(req)) {
            ctx.fireChannelRead(msg);
            return;
        }

        try {
        // 2-握手发送的Get请求，所以只处理Get请求，不是Get请求就构造一个握手失败的响应
            if (!GET.equals(req.method())) {
                sendHttpResponse(ctx, req, new DefaultFullHttpResponse(HTTP_1_1, FORBIDDEN, ctx.alloc().buffer(0)));
                return;
            }
        // 3-根据WebSocket的版本创建不同版本WebSocketServerHandshaker
            final WebSocketServerHandshakerFactory wsFactory = new WebSocketServerHandshakerFactory(
                    getWebSocketLocation(ctx.pipeline(), req, serverConfig.websocketPath()),
                    serverConfig.subprotocols(), serverConfig.decoderConfig());
            final WebSocketServerHandshaker handshaker = wsFactory.newHandshaker(req);
            final ChannelPromise localHandshakePromise = handshakePromise;
            if (handshaker == null) {
            // 4-handshaker如果没有成功创建，说明当前不支持该版本的WebSocket，返回握手失败响应    WebSocketServerHandshakerFactory.sendUnsupportedVersionResponse(ctx.channel());
            } else {
                // Ensure we set the handshaker and replace this handler before we
                // trigger the actual handshake. Otherwise we may receive websocket bytes in this handler
                // before we had a chance to replace it.
                //
                // See https://github.com/netty/netty/issues/9471.
                WebSocketServerProtocolHandler.setHandshaker(ctx.channel(), handshaker);
                // 5-握手器生成成功后，从pipleline中移除当前的WebSocketServerProtocolHandshakeHandler
                ctx.pipeline().remove(this);
                // 6-调用握手方法进行握手效验和构造握手成功响应
                final ChannelFuture handshakeFuture = handshaker.handshake(ctx.channel(), req);
                handshakeFuture.addListener(new ChannelFutureListener() {
                    @Override
                    public void operationComplete(ChannelFuture future) {
                        if (!future.isSuccess()) {
                            localHandshakePromise.tryFailure(future.cause());
                            ctx.fireExceptionCaught(future.cause());
                        } else {
                            localHandshakePromise.trySuccess();
                            // 7-握手成功后进行fireUserEventTriggered钩子方法的回调，调用两次是为了版本兼容
                            // Kept for compatibility
                            ctx.fireUserEventTriggered(
                                    WebSocketServerProtocolHandler.ServerHandshakeStateEvent.HANDSHAKE_COMPLETE);
                            ctx.fireUserEventTriggered(
                                    new WebSocketServerProtocolHandler.HandshakeComplete(
                                            req.uri(), req.headers(), handshaker.selectedSubprotocol()));
                        }
                    }
                });
                applyHandshakeTimeout();
            }
        } finally {
            req.release();
        }
    }

分析完上面的代码，我们知道了大概的握手过程，那我们就可以在WebSocketServerProtocolHandshakeHandler之前插入一个登陆认证的AuthHandler拦截握手请求的FullRequest处理我们的登陆效验，然后将FullRequest传递给下一个handler。这样就能实现在握手阶段进行登陆效验。AuthHandler的代码如下：

import io.jsonwebtoken.lang.Collections;
import io.netty.channel.ChannelHandler;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.SimpleChannelInboundHandler;
import io.netty.handler.codec.http.FullHttpRequest;
import io.netty.handler.codec.http.QueryStringDecoder;
import io.netty.util.AttributeKey;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.Map;

/**
 * 鉴权handler
 */
@Slf4j
@ChannelHandler.Sharable
public class AuthHandler extends SimpleChannelInboundHandler<FullHttpRequest> {
    @Override
    protected void channelRead0(ChannelHandlerContext ctx, FullHttpRequest req) throws Exception {
        Map<String, List<String>> parameters = new QueryStringDecoder(req.uri()).parameters();
        if (login(ctx, parameters)){
        // 1-由于使用的SimpleChannelInboundHandler，所以FullHttpRequest消息的引用计数会自动release，所以我们需要retain后调用fireChannelRead方法传递给下一个handler
            ctx.fireChannelRead(req.retain());
            return;
        }
        ctx.channel().close();
    }

    private boolean login(ChannelHandlerContext ctx, Map<String, List<String>> parameters) {
        if (!Collections.isEmpty(parameters)) {
            //验证登陆
            if (parameters.containsKey(CommonConstant.SocketLoginConst.TOKEN)) {
                // token验证
                DecodedJWT decodedJWT = JwtTokenUtil.decodeJwt(token);
                if (decodedJWT != null) {
                    //验证成功
                    ctx.channel().attr("用户ID").set("用户ID");
                    return true;
                }
                // 验证失败
                log.error("无效token");
                return false;
            }
        }
        log.error("登陆认证数据包格式错误");
        return false;
    }

    @Override
    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
        cause.printStackTrace();
        ctx.close();
    }
}

然后我们在握手成功，连接真正创建后，需要移除掉这个AuthHandler，因为已经验证通过了。那我们在什么地方移除呢，不知道大家有没有注意到WebSocketServerProtocolHandshakeHandler中在握手成功后会调用两次ctx.fireUserEventTriggered方法，这是触发自定义事件的钩子方法，我们可以在AuthHandler之后的自定义的WebSocketHandler中的userEventTriggered方法中移除AuthHandler。下面是WebSocketHandler的精简代码，删除了一些公司的真实业务代码。

import com.alibaba.fastjson.JSON;
import io.netty.channel.ChannelHandler;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.SimpleChannelInboundHandler;
import io.netty.channel.group.ChannelGroup;
import io.netty.channel.group.DefaultChannelGroup;
import io.netty.handler.codec.http.websocketx.TextWebSocketFrame;
import io.netty.util.Attribute;
import io.netty.util.concurrent.GlobalEventExecutor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;

@Slf4j
@ChannelHandler.Sharable
public class WebSocketHandler extends SimpleChannelInboundHandler<TextWebSocketFrame> {
    @Autowired
    private CrmAlertInfoService alertInfoService;
    //所有正在连接的channel都会存在这里面
    public static ChannelGroup CHANNEL_GROUP = new DefaultChannelGroup(GlobalEventExecutor.INSTANCE);
    @Override
    protected void channelRead0(ChannelHandlerContext ctx, TextWebSocketFrame msg) throws Exception {

    }

    //客户端建立连接
    @Override
    public void handlerAdded(ChannelHandlerContext ctx) throws Exception {
        CHANNEL_GROUP.add(ctx.channel());
        log.info("{}上线了!", ctx.channel().remoteAddress());
    }

    //关闭连接
    @Override
    public void handlerRemoved(ChannelHandlerContext ctx) throws Exception {
        CHANNEL_GROUP.remove(ctx.channel());
        log.info("{}断开连接", ctx.channel().remoteAddress());
    }

    //出现异常
    @Override
    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
        cause.printStackTrace();
    }

    /**
     * 握手成功后，钩子回调函数
     * @param ctx
     * @param evt
     * @throws Exception
     */
    @Override
    public void userEventTriggered(ChannelHandlerContext ctx, Object evt) throws Exception {
                log.info("userEventTriggered移除权限认证handler");
        if (ctx.pipeline().get(AuthHandler.class) != null){
            ctx.pipeline().remove(AuthHandler.class);
            // 推送消息
            Attribute<Long> attr = ctx.channel().attr("userId");
            Long userId = attr.get();
            //登陆认证
            if (userId == null) {
                ctx.channel().close();
            }
            // 认证成功
            log.info("处理数据");
            ThreadPoolConfig.ALERT_INFO_THREAD_POOL.execute(() -> {
                // 1-连接成功后推送数据给前端
               
            });
        }
        super.userEventTriggered(ctx, evt);
    }
}

上面代码在连接创建后移除掉了token认证的handler，同时将数据推送给了前端。 好了，到此，我们的需求就完成了。

总结

整个流程下来刚开始可能会晕，但是我们要从WebSocket的整个连接流程开始梳理，再结合Netty处理WebSocket的处理源码，把WebSocket的生命周期梳理清楚，然后再结合我们自己的业务，就可以清晰的知道自己应该怎么实现了。建议大家都动手实践下，光看懂是没用的，真正难的是业务中的点点滴滴。

求波关注，公众号：java精进天路

avatar
文章
阅读
粉丝