From 7d6db536d0f16b61eaf99e377f08b8464b0d5d06 Mon Sep 17 00:00:00 2001
From: slguan <slguan@taosdata.com>
Date: Wed, 11 Dec 2019 16:19:44 +0800
Subject: [PATCH] [TBASE-1241]

---
 src/inc/taosmsg.h                 |  2 ++
 src/rpc/src/trpc.c                |  1 +
 src/system/detail/src/mgmtShell.c | 11 +++++------
 3 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/src/inc/taosmsg.h b/src/inc/taosmsg.h
index b4baaff745..db36309222 100644
--- a/src/inc/taosmsg.h
+++ b/src/inc/taosmsg.h
@@ -222,6 +222,7 @@ typedef struct {
 
   // internal part
   uint32_t destId;
+  uint32_t destIp;
   char     meterId[TSDB_UNI_LEN];
   uint16_t port;  // for UDP only
   char     empty[1];
@@ -663,6 +664,7 @@ typedef struct {
 // internal message
 typedef struct {
   uint32_t destId;
+  uint32_t destIp;
   char     meterId[TSDB_UNI_LEN];
   char     empty[3];
   char     msgType;
diff --git a/src/rpc/src/trpc.c b/src/rpc/src/trpc.c
index 35740acc14..21f2ae4dc0 100755
--- a/src/rpc/src/trpc.c
+++ b/src/rpc/src/trpc.c
@@ -1219,6 +1219,7 @@ int taosSendMsgToPeerH(void *thandle, char *pCont, int contLen, void *ahandle) {
   pServer = pConn->pServer;
   pChann = pServer->channList + pConn->chann;
   pHeader = (STaosHeader *)(pCont - sizeof(STaosHeader));
+  pHeader->destIp = pConn->peerIp;
   msg = (char *)pHeader;
 
   if ((pHeader->msgType & 1U) == 0 && pConn->localPort) pHeader->port = pConn->localPort;
diff --git a/src/system/detail/src/mgmtShell.c b/src/system/detail/src/mgmtShell.c
index f78791ea75..c56939e0dd 100644
--- a/src/system/detail/src/mgmtShell.c
+++ b/src/system/detail/src/mgmtShell.c
@@ -1224,7 +1224,7 @@ int mgmtProcessConnectMsg(char *pMsg, int msgLen, SConnObj *pConn) {
   uint32_t peerIp = taosGetRpcLocalIp(pConn->thandle);
   pConn->usePublicIp = (peerIp == tsPublicIpInt ? 1 : 0);
   mgmtEstablishConn(pConn);
-  
+
 _rsp:
   pStart = taosBuildRspMsgWithSize(pConn->thandle, TSDB_MSG_TYPE_CONNECT_RSP, 128);
   if (pStart == NULL) return 0;
@@ -1298,13 +1298,12 @@ void *mgmtProcessMsgFromShell(char *msg, void *ahandle, void *thandle) {
     pConn = connList + pMsg->destId;
     pConn->thandle = thandle;
     strcpy(pConn->user, pMsg->meterId);
-
-    uint32_t peerIp = taosGetRpcLocalIp(thandle);
-    pConn->usePublicIp = (peerIp == tsPublicIpInt ? 1 : 0);
-    mPrint("pConn:%p is rebuild, peerIp:%s publicIp:%s usePublicIp:%u",
-            pConn, taosIpStr(peerIp), taosIpStr(tsPublicIpInt), pConn->usePublicIp);
   }
 
+  pConn->usePublicIp = (pMsg->destIp == tsPublicIpInt ? 1 : 0);
+  mTrace("pConn:%p, destIp:%s publicIp:%s usePublicIp:%u",
+          pConn, taosIpStr(pMsg->destIp), taosIpStr(tsPublicIpInt), pConn->usePublicIp);
+
   if (pMsg->msgType == TSDB_MSG_TYPE_CONNECT) {
     (*mgmtProcessShellMsg[pMsg->msgType])((char *)pMsg->content, pMsg->msgLen - sizeof(SIntMsg), pConn);
   } else {
-- 
GitLab