From c45634dda4ecf0e9a344a79019bb238b7031ea59 Mon Sep 17 00:00:00 2001
From: zhou-hao <admin@hsweb.me>
Date: Mon, 16 Nov 2020 20:09:44 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8Dtcp=E4=BF=9D=E6=8C=81?=
 =?UTF-8?q?=E4=BC=9A=E8=AF=9D=E6=96=AD=E5=BC=80=E8=BF=9E=E6=8E=A5=E5=90=8E?=
 =?UTF-8?q?=E9=87=8D=E6=96=B0=E8=BF=9E=E6=8E=A5=E6=97=A0=E6=B3=95=E6=81=A2?=
 =?UTF-8?q?=E5=A4=8Dsession?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../community/network/tcp/device/TcpDeviceSession.java    | 3 +++
 .../network/tcp/device/TcpServerDeviceGateway.java        | 8 +++++++-
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/jetlinks-components/network-component/tcp-component/src/main/java/org/jetlinks/community/network/tcp/device/TcpDeviceSession.java b/jetlinks-components/network-component/tcp-component/src/main/java/org/jetlinks/community/network/tcp/device/TcpDeviceSession.java
index 5d200fe9..d0ab63dd 100644
--- a/jetlinks-components/network-component/tcp-component/src/main/java/org/jetlinks/community/network/tcp/device/TcpDeviceSession.java
+++ b/jetlinks-components/network-component/tcp-component/src/main/java/org/jetlinks/community/network/tcp/device/TcpDeviceSession.java
@@ -1,6 +1,7 @@
 package org.jetlinks.community.network.tcp.device;
 
 import lombok.Getter;
+import lombok.Setter;
 import org.jetlinks.core.device.DeviceOperator;
 import org.jetlinks.core.message.codec.EncodedMessage;
 import org.jetlinks.core.message.codec.Transport;
@@ -19,8 +20,10 @@ class TcpDeviceSession implements DeviceSession {
     private String id;
 
     @Getter
+    @Setter
     private DeviceOperator operator;
 
+    @Setter
     private TcpClient client;
 
     @Getter
diff --git a/jetlinks-components/network-component/tcp-component/src/main/java/org/jetlinks/community/network/tcp/device/TcpServerDeviceGateway.java b/jetlinks-components/network-component/tcp-component/src/main/java/org/jetlinks/community/network/tcp/device/TcpServerDeviceGateway.java
index 50a037a8..717c7c14 100644
--- a/jetlinks-components/network-component/tcp-component/src/main/java/org/jetlinks/community/network/tcp/device/TcpServerDeviceGateway.java
+++ b/jetlinks-components/network-component/tcp-component/src/main/java/org/jetlinks/community/network/tcp/device/TcpServerDeviceGateway.java
@@ -209,9 +209,15 @@ class TcpServerDeviceGateway implements DeviceGateway, MonitorSupportDeviceGatew
                         } else {
                             client.onDisconnect(() -> sessionManager.unregister(device.getDeviceId()));
                         }
-                        sessionRef.set(fSession);
                         sessionManager.register(fSession);
                     }
+                    try {
+                        fSession.unwrap(TcpDeviceSession.class)
+                                .setClient(client);
+                    } catch (Throwable ignore) {
+
+                    }
+                    sessionRef.set(fSession);
                     fSession.keepAlive();
                     Duration timeout = message.getHeader(Headers.keepOnlineTimeoutSeconds).map(Duration::ofSeconds).orElse(keepaliveTimeout.get());
                     if (timeout != null) {