From 8ede724abb360a844342270c66ffa57bb5bb3eab Mon Sep 17 00:00:00 2001
From: lynn <lynn@promptness.cn>
Date: Tue, 8 Mar 2022 02:57:48 +0800
Subject: [PATCH] auto read

---
 .../rdp/client/handler/LocalHandler.java      | 32 +++++++++----------
 .../rdp/server/handler/RemoteHandler.java     |  1 +
 2 files changed, 17 insertions(+), 16 deletions(-)

diff --git a/rdp-client/src/main/java/cn/promptness/rdp/client/handler/LocalHandler.java b/rdp-client/src/main/java/cn/promptness/rdp/client/handler/LocalHandler.java
index 85ddbfd..df2467f 100644
--- a/rdp-client/src/main/java/cn/promptness/rdp/client/handler/LocalHandler.java
+++ b/rdp-client/src/main/java/cn/promptness/rdp/client/handler/LocalHandler.java
@@ -29,23 +29,17 @@ public LocalHandler(Channel channel, ClientConfig clientConfig) {
     public void channelActive(ChannelHandlerContext ctx) throws Exception {
         RemoteConfig remoteConfig = clientConfig.getConfig().get(0);
         logger.info("客户端建立本地连接成功,本地绑定IP:{},本地绑定端口:{}", remoteConfig.getLocalIp(), remoteConfig.getLocalPort());
-        Message message = new Message();
-        message.setType(MessageType.TYPE_CONNECTED);
-        message.setClientConfig(clientConfig);
-        message.setData(new byte[0]);
-        channel.writeAndFlush(message);
+        ctx.channel().config().setAutoRead(false);
+        send(MessageType.TYPE_CONNECTED, new byte[0]);
+        ctx.channel().config().setAutoRead(true);
     }
 
+
     @Override
-    protected void channelRead0(ChannelHandlerContext channelHandlerContext, byte[] bytes) throws Exception {
+    protected void channelRead0(ChannelHandlerContext ctx, byte[] bytes) throws Exception {
         RemoteConfig remoteConfig = clientConfig.getConfig().get(0);
         logger.debug("收到本地{}:{}的数据,数据量为:{}字节", remoteConfig.getLocalIp(), remoteConfig.getLocalPort(), bytes.length);
-        Message message = new Message();
-        message.setType(MessageType.TYPE_DATA);
-        message.setClientConfig(clientConfig);
-        message.setData(bytes);
-        // 收到内网服务器响应后返回给服务器端
-        channel.writeAndFlush(message);
+        send(MessageType.TYPE_DATA, bytes);
     }
 
 
@@ -56,10 +50,8 @@ protected void channelRead0(ChannelHandlerContext channelHandlerContext, byte[]
     public void channelInactive(ChannelHandlerContext ctx) throws Exception {
         RemoteConfig remoteConfig = clientConfig.getConfig().get(0);
         logger.info("客户端本地连接断开,本地绑定IP:{},本地绑定端口:{}", remoteConfig.getLocalIp(), remoteConfig.getLocalPort());
-        Message message = new Message();
-        message.setType(MessageType.TYPE_DISCONNECTED);
-        message.setClientConfig(clientConfig);
-        channel.writeAndFlush(message);
+        ctx.channel().config().setAutoRead(true);
+        send(MessageType.TYPE_DISCONNECTED, new byte[0]);
     }
 
     /**
@@ -72,5 +64,13 @@ public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws E
         ctx.channel().close();
     }
 
+    private void send(MessageType type, byte[] data) {
+        Message message = new Message();
+        message.setType(type);
+        message.setClientConfig(clientConfig);
+        message.setData(data);
+        // 收到内网服务器响应后返回给服务器端
+        channel.writeAndFlush(message);
+    }
 
 }
diff --git a/rdp-server/src/main/java/cn/promptness/rdp/server/handler/RemoteHandler.java b/rdp-server/src/main/java/cn/promptness/rdp/server/handler/RemoteHandler.java
index e87d4fb..a068aa6 100644
--- a/rdp-server/src/main/java/cn/promptness/rdp/server/handler/RemoteHandler.java
+++ b/rdp-server/src/main/java/cn/promptness/rdp/server/handler/RemoteHandler.java
@@ -52,6 +52,7 @@ protected void channelRead0(ChannelHandlerContext ctx, byte[] bytes) throws Exce
     @Override
     public void channelInactive(ChannelHandlerContext ctx) throws Exception {
         logger.info("服务端端口[{}]连接断开", remoteConfig.getRemotePort());
+        ctx.channel().config().setAutoRead(true);
         send(MessageType.TYPE_DISCONNECTED, new byte[]{}, ctx);
     }