fix(ws): have netty close websocket channel

Zmax0 · Jul 5, 2024 · c6268fd · c6268fd
1 parent 4dd6cc2
commit c6268fd
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 4 deletions.
diff --git a/urban-spork-client/src/com/urbanspork/client/ClientSocksConnectHandler.java b/urban-spork-client/src/com/urbanspork/client/ClientSocksConnectHandler.java
@@ -16,13 +16,15 @@
 import io.netty.channel.ChannelFutureListener;
 import io.netty.channel.ChannelHandler;
 import io.netty.channel.ChannelHandlerContext;
+import io.netty.channel.ChannelInboundHandlerAdapter;
 import io.netty.channel.ChannelInitializer;
 import io.netty.channel.ChannelOption;
 import io.netty.channel.SimpleChannelInboundHandler;
 import io.netty.handler.codec.MessageToMessageCodec;
 import io.netty.handler.codec.http.HttpClientCodec;
 import io.netty.handler.codec.http.HttpObjectAggregator;
 import io.netty.handler.codec.http.websocketx.BinaryWebSocketFrame;
+import io.netty.handler.codec.http.websocketx.CloseWebSocketFrame;
 import io.netty.handler.codec.http.websocketx.WebSocketClientProtocolHandler;
 import io.netty.handler.codec.socks.SocksCmdType;
 import io.netty.handler.codec.socksx.v5.DefaultSocks5CommandResponse;
@@ -101,6 +103,7 @@ private static void enableWebSocket(Channel inbound, Channel outbound, ServerCon
             ClientSocksInitializer.buildWebSocketHandler(config),
             new WebSocketCodec(inbound, config, request)
         );
+        inbound.closeFuture().addListener(future -> outbound.writeAndFlush(new CloseWebSocketFrame()));
     }
 
     static class WebSocketCodec extends MessageToMessageCodec<BinaryWebSocketFrame, ByteBuf> {
@@ -127,9 +130,14 @@ protected void decode(ChannelHandlerContext ctx, BinaryWebSocketFrame msg, List<
         @Override
         public void userEventTriggered(ChannelHandlerContext ctx, Object evt) {
             if (evt == WebSocketClientProtocolHandler.ClientHandshakeStateEvent.HANDSHAKE_COMPLETE) {
+                inbound.pipeline().addLast(new ChannelInboundHandlerAdapter() {
+                    @Override
+                    public void channelRead(ChannelHandlerContext ctx2, Object msg) {
+                        ctx.channel().writeAndFlush(msg);
+                    }
+                }); // L → R
                 Socks5CommandRequest bndRequest = Socks5.toCommandRequest(Socks5CommandType.CONNECT, (InetSocketAddress) inbound.localAddress());
-                inbound.writeAndFlush(new DefaultSocks5CommandResponse(Socks5CommandStatus.SUCCESS, bndRequest.dstAddrType(), bndRequest.dstAddr(), bndRequest.dstPort()))
-                    .addListener((ChannelFutureListener) channelFuture -> inbound.pipeline().addLast(new DefaultChannelInboundHandler(ctx.channel()))); // L → R
+                inbound.writeAndFlush(new DefaultSocks5CommandResponse(Socks5CommandStatus.SUCCESS, bndRequest.dstAddrType(), bndRequest.dstAddr(), bndRequest.dstPort()));
             }
             if (evt == WebSocketClientProtocolHandler.ClientHandshakeStateEvent.HANDSHAKE_TIMEOUT) {
                 logger.error("Connect proxy websocket server {}:{} time out", config.getHost(), config.getPort());

diff --git a/urban-spork-client/src/com/urbanspork/client/ClientSocksMessageHandler.java b/urban-spork-client/src/com/urbanspork/client/ClientSocksMessageHandler.java
@@ -8,9 +8,9 @@
 import io.netty.handler.codec.socksx.v5.*;
 
 @ChannelHandler.Sharable
-public class ClientSocksMessageHandler extends SimpleChannelInboundHandler<Socks5Message> {
+class ClientSocksMessageHandler extends SimpleChannelInboundHandler<Socks5Message> {
 
-    public static final ClientSocksMessageHandler INSTANCE = new ClientSocksMessageHandler();
+    static final ClientSocksMessageHandler INSTANCE = new ClientSocksMessageHandler();
 
     private ClientSocksMessageHandler() {}