feat(client): move TrafficShapingHandler to `ClientInitializationCo…

…ntext`
Zmax0 · Aug 14, 2024 · a6a9b8e · a6a9b8e
1 parent fb03c60
commit a6a9b8e
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 30 deletions.
diff --git a/urban-spork-client/src/com/urbanspork/client/Client.java b/urban-spork-client/src/com/urbanspork/client/Client.java
@@ -41,24 +41,23 @@ public static void main(String[] args) {
     public static void launch(ClientConfig config, CompletableFuture<Instance> promise) {
         EventLoopGroup bossGroup = new NioEventLoopGroup();
         EventLoopGroup workerGroup = new NioEventLoopGroup();
-        GlobalChannelTrafficShapingHandler trafficShapingHandler = new GlobalChannelTrafficShapingHandler(workerGroup);
-        ServerConfig current = config.getCurrent();
-        current.setTrafficShapingHandler(trafficShapingHandler);
+        GlobalChannelTrafficShapingHandler traffic = new GlobalChannelTrafficShapingHandler(workerGroup);
+        ClientInitializationContext context = new ClientInitializationContext(config, traffic);
         try {
             new ServerBootstrap().group(bossGroup, workerGroup)
                 .channel(NioServerSocketChannel.class)
                 .childOption(ChannelOption.SO_KEEPALIVE, true) // socks5 require
                 .childOption(ChannelOption.TCP_NODELAY, false)
                 .childOption(ChannelOption.SO_LINGER, 1)
-                .childHandler(new ClientInitializer(current))
+                .childHandler(new ClientInitializer(context))
                 .bind(InetAddress.getLoopbackAddress(), config.getPort()).sync().addListener((ChannelFutureListener) future -> {
                     ServerSocketChannel tcp = (ServerSocketChannel) future.channel();
                     InetSocketAddress tcpLocalAddress = tcp.localAddress();
                     int localPort = tcpLocalAddress.getPort();
                     config.setPort(localPort);
-                    DatagramChannel udp = launchUdp(bossGroup, workerGroup, config);
+                    DatagramChannel udp = launchUdp(bossGroup, workerGroup, context);
                     logger.info("Launch client => tcp{} udp{} ", tcpLocalAddress, udp.localAddress());
-                    Instance client = new Instance(tcp, udp, trafficShapingHandler.trafficCounter());
+                    Instance client = new Instance(tcp, udp, traffic.trafficCounter());
                     promise.complete(client);
                 });
             Instance client = promise.get();
@@ -72,13 +71,14 @@ public static void launch(ClientConfig config, CompletableFuture<Instance> promi
             logger.error("Launch client failed", e);
             promise.completeExceptionally(e);
         } finally {
+            traffic.trafficCounter().stop();
             workerGroup.shutdownGracefully();
             bossGroup.shutdownGracefully();
         }
     }
 
-    private static DatagramChannel launchUdp(EventLoopGroup bossGroup, EventLoopGroup workerGroup, ClientConfig config) throws InterruptedException {
-        ServerConfig current = config.getCurrent();
+    private static DatagramChannel launchUdp(EventLoopGroup bossGroup, EventLoopGroup workerGroup, ClientInitializationContext context) throws InterruptedException {
+        ServerConfig current = context.config().getCurrent();
         ChannelHandler udpTransportHandler;
         if (Protocol.vmess == current.getProtocol()) {
             udpTransportHandler = new com.urbanspork.client.vmess.ClientUdpOverTcpHandler(current, workerGroup);
@@ -92,20 +92,19 @@ private static DatagramChannel launchUdp(EventLoopGroup bossGroup, EventLoopGrou
                 @Override
                 protected void initChannel(Channel ch) {
                     ch.pipeline().addLast(
-                        current.getTrafficShapingHandler(),
+                        context.traffic(),
                         new DatagramPacketEncoder(),
                         new DatagramPacketDecoder(),
                         udpTransportHandler
                     );
                 }
             })
-            .bind(InetAddress.getLoopbackAddress(), config.getPort()).sync().channel();
+            .bind(InetAddress.getLoopbackAddress(), context.config().getPort()).sync().channel();
     }
 
     public record Instance(ServerSocketChannel tcp, DatagramChannel udp, TrafficCounter traffic) implements Closeable {
         @Override
         public void close() {
-            traffic.stop();
             tcp.close().awaitUninterruptibly();
             udp.close().awaitUninterruptibly();
         }

diff --git a/urban-spork-client/src/com/urbanspork/client/ClientInitializationContext.java b/urban-spork-client/src/com/urbanspork/client/ClientInitializationContext.java
@@ -0,0 +1,6 @@
+package com.urbanspork.client;
+
+import com.urbanspork.common.config.ClientConfig;
+import io.netty.handler.traffic.GlobalChannelTrafficShapingHandler;
+
+public record ClientInitializationContext(ClientConfig config, GlobalChannelTrafficShapingHandler traffic) {}
diff --git a/urban-spork-client/src/com/urbanspork/client/ClientInitializer.java b/urban-spork-client/src/com/urbanspork/client/ClientInitializer.java
@@ -1,21 +1,20 @@
 package com.urbanspork.client;
 
 import com.urbanspork.common.channel.AttributeKeys;
-import com.urbanspork.common.config.ServerConfig;
 import io.netty.channel.ChannelInitializer;
 import io.netty.channel.socket.nio.NioSocketChannel;
 
 class ClientInitializer extends ChannelInitializer<NioSocketChannel> {
 
-    private final ServerConfig config;
+    private final ClientInitializationContext context;
 
-    ClientInitializer(ServerConfig config) {
-        this.config = config;
+    ClientInitializer(ClientInitializationContext context) {
+        this.context = context;
     }
 
     @Override
     protected void initChannel(NioSocketChannel channel) {
-        channel.attr(AttributeKeys.SERVER_CONFIG).set(config);
-        channel.pipeline().addLast(config.getTrafficShapingHandler(), new ClientProxyUnificationHandler());
+        channel.attr(AttributeKeys.SERVER_CONFIG).set(context.config().getCurrent());
+        channel.pipeline().addLast(context.traffic(), new ClientProxyUnificationHandler());
     }
 }
diff --git a/urban-spork-common/src/com/urbanspork/common/config/ServerConfig.java b/urban-spork-common/src/com/urbanspork/common/config/ServerConfig.java
@@ -1,11 +1,9 @@
 package com.urbanspork.common.config;
 
-import com.fasterxml.jackson.annotation.JsonIgnore;
 import com.urbanspork.common.codec.CipherKind;
 import com.urbanspork.common.protocol.Protocol;
 import com.urbanspork.common.transport.Transport;
 import com.urbanspork.common.transport.udp.PacketEncoding;
-import io.netty.handler.traffic.AbstractTrafficShapingHandler;
 
 import java.util.List;
 
@@ -33,9 +31,6 @@ public class ServerConfig {
 
     private WebSocketSetting ws;
 
-    @JsonIgnore
-    private AbstractTrafficShapingHandler trafficShapingHandler;
-
     public String getHost() {
         return host;
     }
@@ -124,14 +119,6 @@ public void setWs(WebSocketSetting ws) {
         this.ws = ws;
     }
 
-    public AbstractTrafficShapingHandler getTrafficShapingHandler() {
-        return trafficShapingHandler;
-    }
-
-    public void setTrafficShapingHandler(AbstractTrafficShapingHandler trafficShapingHandler) {
-        this.trafficShapingHandler = trafficShapingHandler;
-    }
-
     @Override
     public String toString() {
         StringBuilder builder = new StringBuilder(listItemText());