feat: add a new config to support for closing overtime connections fr…

…om user client automatically (#279)
monkeyWie · Jun 24, 2023 · e2f2d5e · e2f2d5e
1 parent b8bc8ad
commit e2f2d5e
Show file tree

Hide file tree

Showing 5 changed files with 88 additions and 0 deletions.
diff --git a/src/main/java/com/github/monkeywie/proxyee/config/IdleStateCheck.java b/src/main/java/com/github/monkeywie/proxyee/config/IdleStateCheck.java
@@ -0,0 +1,48 @@
+package com.github.monkeywie.proxyee.config;
+
+/**
+ * @Author ffei
+ * @Date 2023/6/24 10:27
+ */
+public class IdleStateCheck {
+    private long readerIdleTime;
+    private long writerIdleTime;
+    private long allIdleTime;
+
+    public IdleStateCheck() {
+        this.readerIdleTime = 10;
+        this.writerIdleTime = 15;
+        this.allIdleTime = 20;
+    }
+
+    public IdleStateCheck(long readerIdleTime, long writerIdleTime, long allIdleTime) {
+        this.readerIdleTime = readerIdleTime;
+        this.writerIdleTime = writerIdleTime;
+        this.allIdleTime = allIdleTime;
+    }
+
+    public long getReaderIdleTime() {
+        return readerIdleTime;
+    }
+
+    public long getWriterIdleTime() {
+        return writerIdleTime;
+    }
+
+    public long getAllIdleTime() {
+        return allIdleTime;
+    }
+
+    public void setReaderIdleTime(long readerIdleTime) {
+        this.readerIdleTime = readerIdleTime;
+    }
+
+    public void setWriterIdleTime(long writerIdleTime) {
+        this.writerIdleTime = writerIdleTime;
+    }
+
+    public void setAllIdleTime(long allIdleTime) {
+        this.allIdleTime = allIdleTime;
+    }
+}
+
diff --git a/src/main/java/com/github/monkeywie/proxyee/handler/HttpProxyServerHandler.java b/src/main/java/com/github/monkeywie/proxyee/handler/HttpProxyServerHandler.java
@@ -24,6 +24,7 @@
 import io.netty.handler.proxy.ProxyHandler;
 import io.netty.handler.ssl.SslContext;
 import io.netty.handler.ssl.SslContextBuilder;
+import io.netty.handler.timeout.IdleStateEvent;
 import io.netty.resolver.NoopAddressResolverGroup;
 import io.netty.util.ReferenceCountUtil;
 
@@ -273,6 +274,13 @@ public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws E
         exceptionHandle.beforeCatch(ctx.channel(), cause);
     }
 
+    @Override
+    public void userEventTriggered(ChannelHandlerContext ctx, Object evt) throws Exception {
+        if (evt instanceof IdleStateEvent) {
+            ctx.channel().close();
+        }
+    }
+
     private boolean authenticate(ChannelHandlerContext ctx, HttpRequest request) {
         if (serverConfig.getAuthenticationProvider() != null) {
             HttpProxyAuthenticationProvider authProvider = serverConfig.getAuthenticationProvider();

diff --git a/src/main/java/com/github/monkeywie/proxyee/handler/TunnelProxyInitializer.java b/src/main/java/com/github/monkeywie/proxyee/handler/TunnelProxyInitializer.java
@@ -6,6 +6,7 @@
 import io.netty.channel.ChannelInboundHandlerAdapter;
 import io.netty.channel.ChannelInitializer;
 import io.netty.handler.proxy.ProxyHandler;
+import io.netty.util.ReferenceCountUtil;
 
 /**
  * http代理隧道，转发原始报文
@@ -29,6 +30,10 @@ protected void initChannel(Channel ch) throws Exception {
         ch.pipeline().addLast(new ChannelInboundHandlerAdapter() {
             @Override
             public void channelRead(ChannelHandlerContext ctx0, Object msg0) throws Exception {
+                if (!clientChannel.isOpen()) {
+                    ReferenceCountUtil.release(msg0);
+                    return;
+                }
                 clientChannel.writeAndFlush(msg0);
             }
 

diff --git a/src/main/java/com/github/monkeywie/proxyee/server/HttpProxyServer.java b/src/main/java/com/github/monkeywie/proxyee/server/HttpProxyServer.java
@@ -1,5 +1,6 @@
 package com.github.monkeywie.proxyee.server;
 
+import com.github.monkeywie.proxyee.config.IdleStateCheck;
 import com.github.monkeywie.proxyee.crt.CertPool;
 import com.github.monkeywie.proxyee.crt.CertUtil;
 import com.github.monkeywie.proxyee.exception.HttpProxyExceptionHandle;
@@ -19,6 +20,7 @@
 import io.netty.handler.logging.LoggingHandler;
 import io.netty.handler.ssl.SslContextBuilder;
 import io.netty.handler.ssl.util.InsecureTrustManagerFactory;
+import io.netty.handler.timeout.IdleStateHandler;
 import io.netty.util.internal.logging.InternalLogger;
 import io.netty.util.internal.logging.InternalLoggerFactory;
 
@@ -28,6 +30,7 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
 import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 
 public class HttpProxyServer {
 
@@ -184,6 +187,13 @@ protected void initChannel(Channel ch) throws Exception {
                                 serverConfig.getMaxInitialLineLength(),
                                 serverConfig.getMaxHeaderSize(),
                                 serverConfig.getMaxChunkSize()));
+                        if (serverConfig.getIdleStateCheck() != null) {
+                            IdleStateCheck idleStateCheck = serverConfig.getIdleStateCheck();
+                            ch.pipeline().addLast("idleStateCheck",
+                                    new IdleStateHandler(idleStateCheck.getReaderIdleTime(), idleStateCheck.getWriterIdleTime(),
+                                            idleStateCheck.getAllIdleTime(), TimeUnit.MILLISECONDS)
+                            );
+                        }
                         ch.pipeline().addLast("serverHandle",
                                 new HttpProxyServerHandler(serverConfig, proxyInterceptInitializer, proxyConfig,
                                         httpProxyExceptionHandle));

diff --git a/src/main/java/com/github/monkeywie/proxyee/server/HttpProxyServerConfig.java b/src/main/java/com/github/monkeywie/proxyee/server/HttpProxyServerConfig.java
@@ -3,6 +3,7 @@
 import com.github.monkeywie.proxyee.server.accept.HttpProxyAcceptHandler;
 import com.github.monkeywie.proxyee.server.accept.HttpProxyMitmMatcher;
 import com.github.monkeywie.proxyee.server.auth.HttpProxyAuthenticationProvider;
+import com.github.monkeywie.proxyee.config.IdleStateCheck;
 import io.netty.channel.EventLoopGroup;
 import io.netty.handler.codec.http.HttpObjectDecoder;
 import io.netty.handler.ssl.SslContext;
@@ -35,6 +36,7 @@ public class HttpProxyServerConfig {
     private int maxInitialLineLength = HttpObjectDecoder.DEFAULT_MAX_INITIAL_LINE_LENGTH;
     private int maxHeaderSize = HttpObjectDecoder.DEFAULT_MAX_HEADER_SIZE;
     private int maxChunkSize = HttpObjectDecoder.DEFAULT_MAX_CHUNK_SIZE;
+    private IdleStateCheck idleStateCheck;
 
     public HttpProxyServerConfig() {
         this(DefaultAddressResolverGroup.INSTANCE);
@@ -62,6 +64,7 @@ private HttpProxyServerConfig(Builder builder) {
         this.maxInitialLineLength = builder.maxInitialLineLength;
         this.maxHeaderSize = builder.maxHeaderSize;
         this.maxChunkSize = builder.maxChunkSize;
+        this.idleStateCheck = builder.idleStateCheck;
     }
 
     public SslContext getClientSslCtx() {
@@ -220,6 +223,14 @@ public void setMaxChunkSize(int maxChunkSize) {
         this.maxChunkSize = maxChunkSize;
     }
 
+    public IdleStateCheck getIdleStateCheck() {
+        return idleStateCheck;
+    }
+
+    public void setIdleStateCheck(IdleStateCheck idleStateCheck) {
+        this.idleStateCheck = idleStateCheck;
+    }
+
     public static class Builder {
         private SslContext clientSslCtx;
         private String issuer;
@@ -239,6 +250,7 @@ public static class Builder {
         private int maxInitialLineLength = HttpObjectDecoder.DEFAULT_MAX_INITIAL_LINE_LENGTH;
         private int maxHeaderSize = HttpObjectDecoder.DEFAULT_MAX_HEADER_SIZE;
         private int maxChunkSize = HttpObjectDecoder.DEFAULT_MAX_CHUNK_SIZE;
+        private IdleStateCheck idleStateCheck;
 
         public Builder() {
             this(DefaultAddressResolverGroup.INSTANCE);
@@ -333,6 +345,11 @@ public Builder setMaxChunkSize(int maxChunkSize) {
             return this;
         }
 
+        public Builder setIdleStateCheck(IdleStateCheck idleStateCheck) {
+            this.idleStateCheck = idleStateCheck;
+            return this;
+        }
+
         public HttpProxyServerConfig build() {
             HttpProxyServerConfig config = new HttpProxyServerConfig(this);
             return config;