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)));
return;
}
// 实现鉴权
if (鉴权不通过){
sendNoAuthHttpResponse(ctx);
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);
ctx.channel().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)));
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 IdleStateEvent) {
if (IdleState.READER_IDLE == ((IdleStateEvent) evt).state()) {
// 心跳没有回应,则主动关闭连接。
String channelId = ctx.channel().id().asLongText();
ctx.close();
}
} else {
super.userEventTriggered(ctx, evt);
}
}