From cc8e5e157eefa53971f4d5d78a3dae38f9f09b21 Mon Sep 17 00:00:00 2001
From: dapan1121 <wpan@taosdata.com>
Date: Tue, 7 Dec 2021 14:48:14 +0800
Subject: [PATCH] add msg protection

---
 src/client/src/tscServer.c | 2 +-
 src/dnode/src/dnodeShell.c | 8 ++++++++
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/src/client/src/tscServer.c b/src/client/src/tscServer.c
index 685e4210d1..215f828d1e 100644
--- a/src/client/src/tscServer.c
+++ b/src/client/src/tscServer.c
@@ -506,7 +506,7 @@ void tscProcessMsgFromServer(SRpcMsg *rpcMsg, SRpcEpSet *pEpSet) {
     }
   }
 
-  if (pRes->code == TSDB_CODE_SUCCESS && tscProcessMsgRsp[pCmd->command]) {
+  if (pRes->code == TSDB_CODE_SUCCESS && pCmd->command < TSDB_SQL_MAX && tscProcessMsgRsp[pCmd->command]) {
     rpcMsg->code = (*tscProcessMsgRsp[pCmd->command])(pSql);
   }
 
diff --git a/src/dnode/src/dnodeShell.c b/src/dnode/src/dnodeShell.c
index f62e0c0df4..7676343b37 100644
--- a/src/dnode/src/dnodeShell.c
+++ b/src/dnode/src/dnodeShell.c
@@ -120,6 +120,14 @@ static void dnodeProcessMsgFromShell(SRpcMsg *pMsg, SRpcEpSet *pEpSet) {
 
   if (pMsg->pCont == NULL) return;
 
+  if (pMsg->msgType >= TSDB_MSG_TYPE_MAX) {
+    dError("RPC %p, shell msg type:%d is not processed", pMsg->handle, pMsg->msgType);
+    rpcMsg.code = TSDB_CODE_DND_MSG_NOT_PROCESSED;
+    rpcSendResponse(&rpcMsg);
+    rpcFreeCont(pMsg->pCont);
+    return;
+  }
+
   SRunStatus dnodeStatus = dnodeGetRunStatus();
   if (dnodeStatus == TSDB_RUN_STATUS_STOPPED) {
     dError("RPC %p, shell msg:%s is ignored since dnode exiting", pMsg->handle, taosMsg[pMsg->msgType]);
-- 
GitLab