feat: add support for udp

Author: SenseBin

src/main/java/sense/test/socks5/constant/ChannelAttrKey.groovy

@@ -8,4 +8,5 @@ import io.netty.util.AttributeKey
  */
 interface ChannelAttrKey {
     AttributeKey<Bootstrap> BOOTSTRAP = new AttributeKey('bootstrap')
+    AttributeKey<Bootstrap> BOOTSTRAP_UDP = new AttributeKey('bootstrap_udp')
 }

src/main/java/sense/test/socks5/handler/CmdHandler.groovy

@@ -42,6 +42,9 @@ class CmdHandler extends ChannelDuplexHandler {
             case SocksCmdType.CONNECT:
                 connect(ctx, request)
                 break
+            case SocksCmdType.UDP:
+                udp(ctx, request)
+                break
             default:
                 ctx.close()
         }
@@ -81,4 +84,39 @@ class CmdHandler extends ChannelDuplexHandler {
             }
         })
     }
+
+    protected void udp(final ChannelHandlerContext ctx, final SocksCmdRequest request) {
+        InetSocketAddress clientAddress = ctx.channel().remoteAddress() as InetSocketAddress
+        logger.debug('以udp模式进行代理，[{}:{} -> {}:{}]',
+                clientAddress.hostString, clientAddress.port, request.host(), request.port())
+
+        ctx.channel().attr(ChannelAttrKey.BOOTSTRAP_UDP).get().connect(request.host(), request.port()).addListener(new ChannelFutureListener() {
+            @Override
+            void operationComplete(ChannelFuture future) throws Exception {
+                Channel channel = null
+                try {
+                    channel = future.sync().channel()
+                    InetSocketAddress socketAddress = channel.remoteAddress() as InetSocketAddress
+                    boolean ipv4 = socketAddress.address instanceof Inet4Address
+                    SocksAddressType addressType = ipv4 ? SocksAddressType.IPv4 : SocksAddressType.IPv6
+                    SocksCmdResponse cmdResponse = new SocksCmdResponse(SocksCmdStatus.SUCCESS, addressType, socketAddress.address.getHostAddress(), socketAddress.port)
+
+                    ctx.pipeline().addFirst(new ConnectHandler(channel))
+                    channel.pipeline().addLast(new ChildConnectHandler(ctx.channel()))
+                    ctx.writeAndFlush(cmdResponse)
+                    ctx.pipeline().remove(CmdHandler.this)
+                } catch (Exception e) {
+                    logger.warn('代理udp请求失败', e)
+                    try {
+                        channel?.close()
+                    } catch (ignore) {
+                    }
+                    try {
+                        ctx.writeAndFlush(new SocksCmdResponse(SocksCmdStatus.FAILURE, SocksAddressType.UNKNOWN)).channel().close()
+                    } catch (ignore) {
+                    }
+                }
+            }
+        })
+    }
 }

src/main/java/sense/test/socks5/strap/ServerStrap.groovy

@@ -7,6 +7,7 @@ import io.netty.channel.ChannelOption
 import io.netty.channel.EventLoopGroup
 import io.netty.channel.nio.NioEventLoopGroup
 import io.netty.channel.socket.SocketChannel
+import io.netty.channel.socket.nio.NioDatagramChannel
 import io.netty.channel.socket.nio.NioServerSocketChannel
 import io.netty.channel.socket.nio.NioSocketChannel
 import io.netty.handler.codec.socks.SocksCmdRequestDecoder
@@ -39,16 +40,23 @@ class ServerStrap {
             return new Thread(r, "childGroup-" + atomicLong.getAndIncrement())
         }
     })
-    private Bootstrap childConnectBootstrap = new Bootstrap().with {
-        EventLoopGroup group = new NioEventLoopGroup(0, new ThreadFactory() {
-            private AtomicLong atomicLong = new AtomicLong()
+    private EventLoopGroup childConnectGroup = new NioEventLoopGroup(0, new ThreadFactory() {
+        private AtomicLong atomicLong = new AtomicLong()
 
+        @Override
+        Thread newThread(Runnable r) {
+            return new Thread(r, "childConnectGroup-" + atomicLong.getAndIncrement())
+        }
+    })
+    private Bootstrap childConnectBootstrap = new Bootstrap().with {
+        it.group(childConnectGroup).channel(NioSocketChannel.class).handler(new ChannelInitializer<SocketChannel>() {
             @Override
-            Thread newThread(Runnable r) {
-                return new Thread(r, "childConnectGroup-" + atomicLong.getAndIncrement())
+            protected void initChannel(SocketChannel socketChannel) throws Exception {
             }
-        })
-        it.group(group).channel(NioSocketChannel.class).handler(new ChannelInitializer<SocketChannel>() {
+        }).option(ChannelOption.CONNECT_TIMEOUT_MILLIS, 30000)
+    }
+    private Bootstrap childUdpConnectBootstrap = new Bootstrap().with {
+        it.group(childConnectGroup).channel(NioDatagramChannel.class).handler(new ChannelInitializer<SocketChannel>() {
             @Override
             protected void initChannel(SocketChannel socketChannel) throws Exception {
             }
@@ -62,14 +70,15 @@ class ServerStrap {
         serverBootstrapTemplate.childHandler(new ChannelInitializer<SocketChannel>() {
             @Override
             protected void initChannel(SocketChannel ch) throws Exception {
-                ch.pipeline()
+                def channel = ch.pipeline()
                         .addLast(new SocksMessageEncoder())
                         .addLast(new SocksInitRequestDecoder())
                         .addLast(new InitHandler())
                         .addLast(new SocksCmdRequestDecoder())
                         .addLast(new CmdHandler())
                         .channel()
-                        .attr(ChannelAttrKey.BOOTSTRAP).set(childConnectBootstrap)
+                channel.attr(ChannelAttrKey.BOOTSTRAP).set(childConnectBootstrap)
+                channel.attr(ChannelAttrKey.BOOTSTRAP_UDP).set(childUdpConnectBootstrap)
             }
         })
         serverBootstrapTemplate.bind(port).sync().channel().closeFuture().await()