netty实现websocket握手前鉴权和针对广播的优化

1、netty实现websocket握手前鉴权和针对广播的优化

1.1、背景

现在用netty实现了websocket，决定在建立连接前添加鉴权。首先想到的就是在
WebSocketServerProtocolHandler后面添加一个HandshakeHandler鉴权，这有一个问题。就是执行到HandshakeHandler时，服务端已经向客户端发送了101，同意握手了。这个时候如果鉴权不通过，相当于多响应了一次服务端。

1.2、解决方案

那最好的执行时机是在握手前就执行鉴权，鉴权通过则放过，实现代码如下：

@Slf4j
public class WsHandler extends SimpleChannelInboundHandler<WebSocketFrame> {
    private static final TextWebSocketFrame HEARTBEAT_SEQUENCE =
            new TextWebSocketFrame(Unpooled.unreleasableBuffer(Unpooled.copiedBuffer("pong", CharsetUtil.UTF_8)));
            
    @Override
    public void handlerAdded(ChannelHandlerContext ctx) throws Exception {
        log.info("------------执行初始化操作---------------");
        // 在这里执行初始化操作
        // 例如设置状态或注册定时器
        ChannelPipeline cp = ctx.pipeline();
        if (cp.get(HandshakeHandler.class) == null) {

            WebSocketServerProtocolHandler webSocketServerProtocolHandler = cp.get(WebSocketServerProtocolHandler.class);
            Field filed = webSocketServerProtocolHandler.getClass().getDeclaredField("serverConfig");
            ReflectionUtils.makeAccessible(filed);
            WebSocketServerProtocolConfig serverConfig = (WebSocketServerProtocolConfig)ReflectionUtils.getField(filed, webSocketServerProtocolHandler);

            cp.addBefore("io.netty.handler.codec.http.websocketx.WebSocketServerProtocolHandshakeHandler",  HandshakeHandler.class.getName(),
                    new HandshakeHandler(serverConfig));
        }
    }
    
    ......其他执行......
}

@Slf4j
public class HandshakeHandler extends ChannelInboundHandlerAdapter {
    private final WebSocketServerProtocolConfig serverConfig;

    public HandshakeHandler(WebSocketServerProtocolConfig serverConfig) {
        this.serverConfig = checkNotNull(serverConfig, "serverConfig");
    }

    @Override
    public void handlerAdded(ChannelHandlerContext ctx) {

    }

    @Override
    public void channelRead(final ChannelHandlerContext ctx, Object msg) throws Exception {
        final FullHttpRequest req = (FullHttpRequest) msg;
        if (!isWebSocketPath(req)) {
            ctx.fireChannelRead(msg);
            return;
        }

        if (!GET.equals(req.method())) {
            sendHttpResponse(ctx, req, new DefaultFullHttpResponse(HTTP_1_1, FORBIDDEN, ctx.alloc().buffer(0)));
            //避免内存泄漏
            ReferenceCountUtil.release(msg);
            return;
        }

        // 实现鉴权
                    String requestUri = req.uri();
            UrlQuery urlQuery = UrlQuery.of(requestUri, CharsetUtil.CHARSET_UTF_8, true);
            Object tokenObj = urlQuery.get(TOKEN);
        if (鉴权不通过){
            sendNoAuthHttpResponse(ctx);
            ReferenceCountUtil.release(msg);
            return;
        }

        ctx.pipeline().remove(this);
        ctx.fireChannelRead(msg);
    }


    private boolean isWebSocketPath(FullHttpRequest req) {
        String websocketPath = serverConfig.websocketPath();
        String uri = req.uri();
        boolean checkStartUri = uri.startsWith(websocketPath);
        boolean checkNextUri = "/".equals(websocketPath) || checkNextUri(uri, websocketPath);
        return serverConfig.checkStartsWith() ? (checkStartUri && checkNextUri) : uri.equals(websocketPath);
    }

    private boolean checkNextUri(String uri, String websocketPath) {
        int len = websocketPath.length();
        if (uri.length() > len) {
            char nextUri = uri.charAt(len);
            return nextUri == '/' || nextUri == '?';
        }
        return true;
    }

    private static void sendNoAuthHttpResponse(ChannelHandlerContext ctx) {
        HttpResponse response = new DefaultHttpResponse(HttpVersion.HTTP_1_1, HttpResponseStatus.UNAUTHORIZED);
        response.headers().set(HttpHeaderNames.WWW_AUTHENTICATE, "Basic realm=\"webconsole\"");
        response.headers().set(HttpHeaderNames.CONTENT_TYPE, "text/plain");
        response.headers().set(HttpHeaderNames.CONTENT_LENGTH, 0);
        ctx.writeAndFlush(response).addListener(ChannelFutureListener.CLOSE);
    }

    private static void sendHttpResponse(ChannelHandlerContext ctx, HttpRequest req, HttpResponse res) {
        ChannelFuture f = ctx.channel().writeAndFlush(res);
        if (!isKeepAlive(req) || res.status().code() != 200) {
            f.addListener(ChannelFutureListener.CLOSE);
        }
    }

}

咱们来看看WsHandler#handlerAdded，

这个代码就是在向pipeline添加handler时，在WebSocketServerProtocolHandshakeHandler之前添加一个HandshakeHandler。

这个WebSocketServerProtocolHandler什么时候来的了？我们明明没有注册进去呀？大家可以看看WebSocketServerProtocolHandler这个handler，这个handler添加pipeline时会添加。这个类就是用来处理握手协议的，咱们在这个handler前添加鉴权的handler，完美无缝衔接。

接着咱们来看看这个类HandshakeHandler，这个类很简单就是鉴权不通过响应401给前端，然后拒绝协议升级。如果鉴权通过则将该handler从pipeline剔除，因为鉴权只存在握手前，握手前鉴权通过，就可以收发websocket协议了，没必要鉴权这些数据了。

2、针对广播的优化

我们有一个SessionManager记录了所有存活的channel，在执行广播的时候将所有的channel进行分片，然后异步线程针对分片的每个channel循环发送消息。

我们发送消息首先想到的是new TextWebSocketFrame(content)，然后writeAndFlush发送。创建的TextWebSocketFrame发送之后会执行release方法将refCnt置0，然后内存就标记可以被回收了。

我们循环到第二个channel发送该对象就会抛异常，因为对象已经被释放了。那要想解决这个问题，增加引用次数。

现在我们在发送第一个channel发送成功，后续所有channel发送都是成功，但是报文都是“”。那这是为啥了？

因为TextWebSocketFrame对象内部把持一个byteBuf，发送一次消息后readIndex就到终止读的位置了。第二次就读不出任何数据了，就读到""发送出去了。那怎么解决了？

我们可以在创建TextWebSocketFrame对象时，执行payload.content().markReaderIndex()标识标记当前读的索引位置。然后再发送writeAndFlush后，payload.content().resetReaderIndex()重置读的位置。可是现在发现循环发送消息时会断开连接，因为writeAndFlush看起来是立刻发送消息，以为是同步发送，其实他是异步发送消息。刚消息发送到一半，你重置了读索引，就导致读取数据不对，不符合websocket协议，就断开链接了。

解决的最终办法就是②和③，我发送一个副本，这样副本的读写索引都会重新生成一个，但是其持有的内存不会重新分配。紧接着假设该channel不存活了，我就手动执行ReferenceCountUtil.release(payload)释放状态。如果writeAndFlush内部发送失败也会主动执行ReferenceCountUtil#release释放内存的，所以问题没有。

3、心跳优化

心跳可以直接在websocket协议级别，收到ping后返回pong。什么叫协议级别了？就是websocket建立连接后，会通过websocket的协议发送报文。当opcode是09表示ping帧，10是pong帧。大家可以点击下
这个是我基于netty实现的my-netty，然后基于my-netty实现了websocket，socks5，http，dns，tcp协议

我们可以利用ws协议级别发送pingpong，这样就不需要明文设置ping和pong字符串发送，极大减少内存，增加发送速率。以下是netty对协议级别的支持。

if (webSocketFrame instanceof PingWebSocketFrame) {
            ctx.writeAndFlush(new PongWebSocketFrame());
        }

但是浏览器通过js很难去实现，所以也有需要明文发送pingpong的场景，下面是对这种场景的支持。
由于心跳数据是固定的，所以创建一个静态变量。

 private static final TextWebSocketFrame HEARTBEAT_SEQUENCE =
            new TextWebSocketFrame(Unpooled.unreleasableBuffer(Unpooled.copiedBuffer("pong", CharsetUtil.UTF_8).asReadOnly()));

Unpooled.unreleasableBuffer表示这个哪怕writeAndFlush后执行release时，也不会释放内存，也就是refCnt不会减一成0，导致内存被释放。注意这个方法由于不会释放内存，所以建议只在固定不变的内容，且内容较少的场景下使用，避免内存泄漏的风险。
接着发送心跳报文，这里会创建一个副本。为什么创建一个副本不直接发送了？是因为直接发送虽然引用次数不会减少，不会释放内存，但是readIndex会读到writeIndex。导致第二次发送消息时只能读取到""。而又是多线程的场景，通过markReaderIndex和resetReaderIndex会导致并发问题。所以创建一个副本就是最好的，虽然创建了一个副本，但是共享的还是同一片内存。引用次数和读写索引位置也是一样的。

if (isPing()){
                ctx.writeAndFlush(HEARTBEAT_SEQUENCE.duplicate());
                return;
            }

长时间没收到心跳则断开链接。

@Override
public void userEventTriggered(ChannelHandlerContext ctx, Object evt) throws Exception {
if (evt instanceof WebSocketServerProtocolHandler.HandshakeComplete){
            WebSocketServerProtocolHandler.HandshakeComplete handshakeComplete = (WebSocketServerProtocolHandler.HandshakeComplete) evt;
            HttpHeaders handshakeHeaders = handshakeComplete.requestHeaders();
            Channel channel = ctx.channel();
            String channelId = channel.id().asLongText();
            String userId = channel.attr(NettyConstants.USER_ATTR).get();

        
            log.info("握手成功，存储回话信息 userId: {} , channelId: {} ",userId,channelId);
        } 
    else if (evt instanceof IdleStateEvent) {
        if (IdleState.READER_IDLE == ((IdleStateEvent) evt).state()) {
            // 心跳没有回应，则主动关闭连接。
            String channelId = ctx.channel().id().asLongText();
            ctx.close();
        }
    } else {
        super.userEventTriggered(ctx, evt);
    }
}