From 537d56d1f106d7d91873382bd408a1b32beabdc4 Mon Sep 17 00:00:00 2001
From: hengyunabc <hengyunabc@gmail.com>
Date: Tue, 23 Feb 2021 00:30:47 +0800
Subject: [PATCH] websocket support idle PingWebSocketFrame. #1573

---
 .../com/taobao/arthas/common/ArthasConstants.java   |  2 ++
 .../term/impl/http/LocalTtyServerInitializer.java   |  2 ++
 .../shell/term/impl/http/TtyServerInitializer.java  |  2 ++
 .../term/impl/http/TtyWebSocketFrameHandler.java    |  5 +++++
 .../term/impl/httptelnet/ProtocolDetectHandler.java |  2 ++
 .../alibaba/arthas/tunnel/client/TunnelClient.java  |  2 ++
 .../client/TunnelClientSocketClientHandler.java     | 13 ++++++++++++-
 7 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/common/src/main/java/com/taobao/arthas/common/ArthasConstants.java b/common/src/main/java/com/taobao/arthas/common/ArthasConstants.java
index 94947905d..cfee4b65f 100644
--- a/common/src/main/java/com/taobao/arthas/common/ArthasConstants.java
+++ b/common/src/main/java/com/taobao/arthas/common/ArthasConstants.java
@@ -24,4 +24,6 @@ public class ArthasConstants {
     public static final String SPRING_APPLICATION_NAME = "spring.application.name";
 
     public static final int TELNET_PORT = 3658;
+
+    public static final int WEBSOCKET_IDLE_SECONDS = 60;
 }
diff --git a/core/src/main/java/com/taobao/arthas/core/shell/term/impl/http/LocalTtyServerInitializer.java b/core/src/main/java/com/taobao/arthas/core/shell/term/impl/http/LocalTtyServerInitializer.java
index 0ff80cc5f..bf51b8302 100644
--- a/core/src/main/java/com/taobao/arthas/core/shell/term/impl/http/LocalTtyServerInitializer.java
+++ b/core/src/main/java/com/taobao/arthas/core/shell/term/impl/http/LocalTtyServerInitializer.java
@@ -8,6 +8,7 @@ import io.netty.handler.codec.http.HttpObjectAggregator;
 import io.netty.handler.codec.http.HttpServerCodec;
 import io.netty.handler.codec.http.websocketx.WebSocketServerProtocolHandler;
 import io.netty.handler.stream.ChunkedWriteHandler;
+import io.netty.handler.timeout.IdleStateHandler;
 import io.netty.util.concurrent.EventExecutorGroup;
 import io.termd.core.function.Consumer;
 import io.termd.core.tty.TtyConnection;
@@ -43,6 +44,7 @@ public class LocalTtyServerInitializer extends ChannelInitializer<LocalChannel>
         pipeline.addLast(new HttpObjectAggregator(ArthasConstants.MAX_HTTP_CONTENT_LENGTH));
         pipeline.addLast(workerGroup, "HttpRequestHandler", new HttpRequestHandler("/ws"));
         pipeline.addLast(new WebSocketServerProtocolHandler("/ws"));
+        pipeline.addLast(new IdleStateHandler(0, 0, ArthasConstants.WEBSOCKET_IDLE_SECONDS));
         pipeline.addLast(new TtyWebSocketFrameHandler(group, handler));
     }
 }
diff --git a/core/src/main/java/com/taobao/arthas/core/shell/term/impl/http/TtyServerInitializer.java b/core/src/main/java/com/taobao/arthas/core/shell/term/impl/http/TtyServerInitializer.java
index 6693a7dc8..ff4cba921 100644
--- a/core/src/main/java/com/taobao/arthas/core/shell/term/impl/http/TtyServerInitializer.java
+++ b/core/src/main/java/com/taobao/arthas/core/shell/term/impl/http/TtyServerInitializer.java
@@ -8,6 +8,7 @@ import io.netty.handler.codec.http.HttpObjectAggregator;
 import io.netty.handler.codec.http.HttpServerCodec;
 import io.netty.handler.codec.http.websocketx.WebSocketServerProtocolHandler;
 import io.netty.handler.stream.ChunkedWriteHandler;
+import io.netty.handler.timeout.IdleStateHandler;
 import io.netty.util.concurrent.EventExecutorGroup;
 import io.termd.core.function.Consumer;
 import io.termd.core.tty.TtyConnection;
@@ -41,6 +42,7 @@ public class TtyServerInitializer extends ChannelInitializer<SocketChannel> {
     pipeline.addLast(new HttpObjectAggregator(ArthasConstants.MAX_HTTP_CONTENT_LENGTH));
     pipeline.addLast(workerGroup, "HttpRequestHandler", new HttpRequestHandler("/ws"));
     pipeline.addLast(new WebSocketServerProtocolHandler("/ws"));
+    pipeline.addLast(new IdleStateHandler(0, 0, ArthasConstants.WEBSOCKET_IDLE_SECONDS));
     pipeline.addLast(new TtyWebSocketFrameHandler(group, handler));
   }
 }
diff --git a/core/src/main/java/com/taobao/arthas/core/shell/term/impl/http/TtyWebSocketFrameHandler.java b/core/src/main/java/com/taobao/arthas/core/shell/term/impl/http/TtyWebSocketFrameHandler.java
index cf7c6a32c..e3702ea31 100644
--- a/core/src/main/java/com/taobao/arthas/core/shell/term/impl/http/TtyWebSocketFrameHandler.java
+++ b/core/src/main/java/com/taobao/arthas/core/shell/term/impl/http/TtyWebSocketFrameHandler.java
@@ -21,8 +21,10 @@ import io.netty.buffer.Unpooled;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.channel.SimpleChannelInboundHandler;
 import io.netty.channel.group.ChannelGroup;
+import io.netty.handler.codec.http.websocketx.PingWebSocketFrame;
 import io.netty.handler.codec.http.websocketx.TextWebSocketFrame;
 import io.netty.handler.codec.http.websocketx.WebSocketServerProtocolHandler;
+import io.netty.handler.timeout.IdleStateEvent;
 import io.termd.core.function.Consumer;
 import io.termd.core.http.HttpTtyConnection;
 import io.termd.core.tty.TtyConnection;
@@ -87,6 +89,8 @@ public class TtyWebSocketFrameHandler extends SimpleChannelInboundHandler<TextWe
         }
       };
       handler.accept(conn);
+    } else if (evt instanceof IdleStateEvent) {
+      ctx.writeAndFlush(new PingWebSocketFrame());
     } else {
       super.userEventTriggered(ctx, evt);
     }
@@ -105,6 +109,7 @@ public class TtyWebSocketFrameHandler extends SimpleChannelInboundHandler<TextWe
     }
   }
 
+  @Override
   public void channelRead0(ChannelHandlerContext ctx, TextWebSocketFrame msg) throws Exception {
     conn.writeToDecoder(msg.text());
   }
diff --git a/core/src/main/java/com/taobao/arthas/core/shell/term/impl/httptelnet/ProtocolDetectHandler.java b/core/src/main/java/com/taobao/arthas/core/shell/term/impl/httptelnet/ProtocolDetectHandler.java
index 09fbf2e18..77cb80672 100644
--- a/core/src/main/java/com/taobao/arthas/core/shell/term/impl/httptelnet/ProtocolDetectHandler.java
+++ b/core/src/main/java/com/taobao/arthas/core/shell/term/impl/httptelnet/ProtocolDetectHandler.java
@@ -16,6 +16,7 @@ import io.netty.handler.codec.http.HttpObjectAggregator;
 import io.netty.handler.codec.http.HttpServerCodec;
 import io.netty.handler.codec.http.websocketx.WebSocketServerProtocolHandler;
 import io.netty.handler.stream.ChunkedWriteHandler;
+import io.netty.handler.timeout.IdleStateHandler;
 import io.netty.util.concurrent.EventExecutorGroup;
 import io.netty.util.concurrent.ScheduledFuture;
 import io.termd.core.function.Consumer;
@@ -88,6 +89,7 @@ public class ProtocolDetectHandler extends ChannelInboundHandlerAdapter {
             pipeline.addLast(new HttpObjectAggregator(ArthasConstants.MAX_HTTP_CONTENT_LENGTH));
             pipeline.addLast(workerGroup, "HttpRequestHandler", new HttpRequestHandler("/ws"));
             pipeline.addLast(new WebSocketServerProtocolHandler("/ws"));
+            pipeline.addLast(new IdleStateHandler(0, 0, ArthasConstants.WEBSOCKET_IDLE_SECONDS));
             pipeline.addLast(new TtyWebSocketFrameHandler(channelGroup, ttyConnectionFactory));
             ctx.fireChannelActive();
         }
diff --git a/tunnel-client/src/main/java/com/alibaba/arthas/tunnel/client/TunnelClient.java b/tunnel-client/src/main/java/com/alibaba/arthas/tunnel/client/TunnelClient.java
index fc9aa6eb5..f79a067e7 100644
--- a/tunnel-client/src/main/java/com/alibaba/arthas/tunnel/client/TunnelClient.java
+++ b/tunnel-client/src/main/java/com/alibaba/arthas/tunnel/client/TunnelClient.java
@@ -34,6 +34,7 @@ import io.netty.handler.codec.http.websocketx.WebSocketVersion;
 import io.netty.handler.ssl.SslContext;
 import io.netty.handler.ssl.SslContextBuilder;
 import io.netty.handler.ssl.util.InsecureTrustManagerFactory;
+import io.netty.handler.timeout.IdleStateHandler;
 import io.netty.util.concurrent.DefaultThreadFactory;
 
 /**
@@ -129,6 +130,7 @@ public class TunnelClient {
                         }
 
                         p.addLast(new HttpClientCodec(), new HttpObjectAggregator(ArthasConstants.MAX_HTTP_CONTENT_LENGTH), websocketClientHandler,
+                                new IdleStateHandler(0, 0, ArthasConstants.WEBSOCKET_IDLE_SECONDS),
                                 handler);
                     }
                 });
diff --git a/tunnel-client/src/main/java/com/alibaba/arthas/tunnel/client/TunnelClientSocketClientHandler.java b/tunnel-client/src/main/java/com/alibaba/arthas/tunnel/client/TunnelClientSocketClientHandler.java
index fdf04dea5..ef5f87266 100644
--- a/tunnel-client/src/main/java/com/alibaba/arthas/tunnel/client/TunnelClientSocketClientHandler.java
+++ b/tunnel-client/src/main/java/com/alibaba/arthas/tunnel/client/TunnelClientSocketClientHandler.java
@@ -20,11 +20,13 @@ import io.netty.channel.ChannelHandlerContext;
 import io.netty.channel.ChannelPromise;
 import io.netty.channel.SimpleChannelInboundHandler;
 import io.netty.handler.codec.base64.Base64;
-import io.netty.handler.codec.base64.Base64Encoder;
 import io.netty.handler.codec.http.QueryStringDecoder;
 import io.netty.handler.codec.http.QueryStringEncoder;
+import io.netty.handler.codec.http.websocketx.PingWebSocketFrame;
 import io.netty.handler.codec.http.websocketx.TextWebSocketFrame;
 import io.netty.handler.codec.http.websocketx.WebSocketFrame;
+import io.netty.handler.timeout.IdleState;
+import io.netty.handler.timeout.IdleStateEvent;
 import io.netty.util.CharsetUtil;
 
 /**
@@ -152,6 +154,15 @@ public class TunnelClientSocketClientHandler extends SimpleChannelInboundHandler
         }, tunnelClient.getReconnectDelay(), TimeUnit.SECONDS);
     }
 
+    @Override
+    public void userEventTriggered(ChannelHandlerContext ctx, Object evt) throws Exception {
+        if (evt instanceof IdleStateEvent) {
+            ctx.writeAndFlush(new PingWebSocketFrame());
+        } else {
+            super.userEventTriggered(ctx, evt);
+        }
+    }
+
     @Override
     public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
         if (!registerPromise.isDone()) {