fix UAF

include/libs/stream/tstream.h

@@ -397,6 +397,7 @@ typedef struct SStreamMeta {
   SMetaHbInfo   hbInfo;
   int32_t       closedTask;
   int32_t       chkptNotReadyTasks;
+  int64_t       rid;
 
   int64_t  chkpId;
   SArray*  chkpSaved;
@@ -727,8 +728,8 @@ int32_t streamTaskReleaseState(SStreamTask* pTask);
 int32_t streamTaskReloadState(SStreamTask* pTask);
 int32_t streamAlignTransferState(SStreamTask* pTask);
 
-int32_t streamAddCheckpointSourceRspMsg(SStreamCheckpointSourceReq* pReq, SRpcHandleInfo* pRpcInfo,
+int32_t streamAddCheckpointSourceRspMsg(SStreamCheckpointSourceReq* pReq, SRpcHandleInfo* pRpcInfo, SStreamTask* pTask,
-                                        SStreamTask* pTask, int8_t isSucceed);
+                                        int8_t isSucceed);
 int32_t buildCheckpointSourceRsp(SStreamCheckpointSourceReq* pReq, SRpcHandleInfo* pRpcInfo, SRpcMsg* pMsg,
                                  int8_t isSucceed);
 #ifdef __cplusplus

source/libs/stream/src/stream.c

@@ -66,8 +66,8 @@ static void streamSchedByTimer(void* param, void* tmrId) {
   qDebug("s-task:%s in scheduler, trigger status:%d, next:%dms", pTask->id.idStr, status, (int32_t)pTask->triggerParam);
 
   if (streamTaskShouldStop(&pTask->status) || streamTaskShouldPause(&pTask->status)) {
-    streamMetaReleaseTask(NULL, pTask);
     qDebug("s-task:%s jump out of schedTimer", pTask->id.idStr);
+    streamMetaReleaseTask(NULL, pTask);
     return;
   }
 
@@ -262,7 +262,8 @@ int32_t streamTaskOutputResultBlock(SStreamTask* pTask, SStreamDataBlock* pBlock
 //   return status;
 // }
 
-//static int32_t buildDispatchRsp(const SStreamTask* pTask, const SStreamDispatchReq* pReq, int32_t status, void** pBuf) {
+// static int32_t buildDispatchRsp(const SStreamTask* pTask, const SStreamDispatchReq* pReq, int32_t status, void**
+// pBuf) {
 //   *pBuf = rpcMallocCont(sizeof(SMsgHead) + sizeof(SStreamDispatchRsp));
 //   if (*pBuf == NULL) {
 //     return TSDB_CODE_OUT_OF_MEMORY;
@@ -300,7 +301,8 @@ int32_t streamProcessDispatchMsg(SStreamTask* pTask, SStreamDispatchReq* pReq, S
             pReq->upstreamTaskId);
       status = TASK_INPUT_STATUS__BLOCKED;
     } else {
-      // Current task has received the checkpoint req from the upstream task, from which the message should all be blocked
+      // Current task has received the checkpoint req from the upstream task, from which the message should all be
+      // blocked
       if (pReq->type == STREAM_INPUT__CHECKPOINT_TRIGGER) {
         streamTaskCloseUpstreamInput(pTask, pReq->upstreamTaskId);
         qDebug("s-task:%s close inputQ for upstream:0x%x", pTask->id.idStr, pReq->upstreamTaskId);
@@ -393,10 +395,11 @@ int32_t tAppendDataToInputQueue(SStreamTask* pTask, SStreamQueueItem* pItem) {
       destroyStreamDataBlock((SStreamDataBlock*)pItem);
       return code;
     }
-  } else if (type == STREAM_INPUT__CHECKPOINT || type == STREAM_INPUT__CHECKPOINT_TRIGGER || type == STREAM_INPUT__TRANS_STATE) {
+  } else if (type == STREAM_INPUT__CHECKPOINT || type == STREAM_INPUT__CHECKPOINT_TRIGGER ||
+             type == STREAM_INPUT__TRANS_STATE) {
     taosWriteQitem(pTask->inputQueue->queue, pItem);
-    qDebug("s-task:%s level:%d %s blockdata enqueue, total in queue:%d, size:%.2fMiB",
+    qDebug("s-task:%s level:%d %s blockdata enqueue, total in queue:%d, size:%.2fMiB", pTask->id.idStr,
-           pTask->id.idStr, pTask->info.taskLevel, streamGetBlockTypeStr(type), total, size);
+           pTask->info.taskLevel, streamGetBlockTypeStr(type), total, size);
   } else if (type == STREAM_INPUT__GET_RES) {
     // use the default memory limit, refactor later.
     taosWriteQitem(pTask->inputQueue->queue, pItem);

source/libs/stream/src/streamMeta.c

@@ -25,22 +25,28 @@
 #define META_HB_SEND_IDLE_COUNTER 25  // send hb every 5 sec
 
 static TdThreadOnce streamMetaModuleInit = PTHREAD_ONCE_INIT;
+
 int32_t streamBackendId = 0;
 int32_t streamBackendCfWrapperId = 0;
+int32_t streamMetaId = 0;
 
 static int64_t streamGetLatestCheckpointId(SStreamMeta* pMeta);
 static void    metaHbToMnode(void* param, void* tmrId);
 static void    streamMetaClear(SStreamMeta* pMeta);
 
+void        streamMetaCloseImpl(void* arg);
 static void streamMetaEnvInit() {
   streamBackendId = taosOpenRef(64, streamBackendCleanup);
   streamBackendCfWrapperId = taosOpenRef(64, streamBackendHandleCleanup);
+
+  streamMetaId = taosOpenRef(64, streamMetaCloseImpl);
 }
 
 void streamMetaInit() { taosThreadOnce(&streamMetaModuleInit, streamMetaEnvInit); }
 void streamMetaCleanup() {
   taosCloseRef(streamBackendId);
   taosCloseRef(streamBackendCfWrapperId);
+  taosCloseRef(streamMetaId);
 }
 
 SStreamMeta* streamMetaOpen(const char* path, void* ahandle, FTaskExpand expandFunc, int32_t vgId, int64_t stage) {
@@ -92,7 +98,11 @@ SStreamMeta* streamMetaOpen(const char* path, void* ahandle, FTaskExpand expandF
   pMeta->stage = stage;
 
   // send heartbeat every 5sec.
-  pMeta->hbInfo.hbTmr = taosTmrStart(metaHbToMnode, META_HB_CHECK_INTERVAL, pMeta, streamEnv.timer);
+  pMeta->rid = taosAddRef(streamMetaId, pMeta);
+  int64_t* pRid = taosMemoryMalloc(sizeof(int64_t));
+  *pRid = pMeta->rid;
+
+  pMeta->hbInfo.hbTmr = taosTmrStart(metaHbToMnode, META_HB_CHECK_INTERVAL, pRid, streamEnv.timer);
   pMeta->hbInfo.tickCounter = 0;
   pMeta->hbInfo.stopFlag = 0;
 
@@ -116,9 +126,6 @@ SStreamMeta* streamMetaOpen(const char* path, void* ahandle, FTaskExpand expandF
     }
   }
 
-  // if (pMeta->streamBackend == NULL) {
-  //   goto _err;
-  // }
   pMeta->streamBackendRid = taosAddRef(streamBackendId, pMeta->streamBackend);
 
   code = streamBackendLoadCheckpointInfo(pMeta);
@@ -207,11 +214,15 @@ void streamMetaClear(SStreamMeta* pMeta) {
 
 void streamMetaClose(SStreamMeta* pMeta) {
   qDebug("start to close stream meta");
+  taosRemoveRef(streamMetaId, pMeta->rid);
+}
+
+void streamMetaCloseImpl(void* arg) {
+  SStreamMeta* pMeta = arg;
+  qDebug("start to do-close stream meta");
   if (pMeta == NULL) {
     return;
   }
-  // TODO, add ref to resolve race: timer thread cannot stop
-  taosTmrStop(pMeta->hbInfo.hbTmr);
 
   streamMetaClear(pMeta);
 
@@ -625,18 +636,26 @@ static bool readyToSendHb(SMetaHbInfo* pInfo) {
 }
 
 void metaHbToMnode(void* param, void* tmrId) {
-  SStreamMeta* pMeta = param;
+  int64_t rid = *(int64_t*)param;
+
   SStreamHbMsg hbMsg = {0};
+  SStreamMeta* pMeta = taosAcquireRef(streamMetaId, rid);
+  if (pMeta == NULL) {
+    taosMemoryFree(param);
+    return;
+  }
 
   // need to stop, stop now
   if (pMeta->hbInfo.stopFlag == STREAM_META_WILL_STOP) {
     pMeta->hbInfo.stopFlag = STREAM_META_OK_TO_STOP;
     qDebug("vgId:%d jump out of meta timer", pMeta->vgId);
+    taosReleaseRef(streamMetaId, rid);
     return;
   }
 
   if (!readyToSendHb(&pMeta->hbInfo)) {
-    taosTmrReset(metaHbToMnode, META_HB_CHECK_INTERVAL, pMeta, streamEnv.timer, &pMeta->hbInfo.hbTmr);
+    taosTmrReset(metaHbToMnode, META_HB_CHECK_INTERVAL, param, streamEnv.timer, &pMeta->hbInfo.hbTmr);
+    taosReleaseRef(streamMetaId, rid);
     return;
   }
 
@@ -675,6 +694,7 @@ void metaHbToMnode(void* param, void* tmrId) {
   if (code < 0) {
     qError("vgId:%d encode stream hb msg failed, code:%s", pMeta->vgId, tstrerror(code));
     taosArrayDestroy(hbMsg.pTaskStatus);
+    taosReleaseRef(streamMetaId, rid);
     return;
   }
 
@@ -682,6 +702,7 @@ void metaHbToMnode(void* param, void* tmrId) {
   if (buf == NULL) {
     qError("vgId:%d encode stream hb msg failed, code:%s", pMeta->vgId, tstrerror(TSDB_CODE_OUT_OF_MEMORY));
     taosArrayDestroy(hbMsg.pTaskStatus);
+    taosReleaseRef(streamMetaId, rid);
     return;
   }
 
@@ -691,6 +712,7 @@ void metaHbToMnode(void* param, void* tmrId) {
     rpcFreeCont(buf);
     qError("vgId:%d encode stream hb msg failed, code:%s", pMeta->vgId, tstrerror(code));
     taosArrayDestroy(hbMsg.pTaskStatus);
+    taosReleaseRef(streamMetaId, rid);
     return;
   }
   tEncoderClear(&encoder);
@@ -704,5 +726,6 @@ void metaHbToMnode(void* param, void* tmrId) {
   qDebug("vgId:%d, build and send hb to mnode", pMeta->vgId);
 
   tmsgSendReq(&epset, &msg);
-  taosTmrReset(metaHbToMnode, META_HB_CHECK_INTERVAL, pMeta, streamEnv.timer, &pMeta->hbInfo.hbTmr);
+  taosTmrReset(metaHbToMnode, META_HB_CHECK_INTERVAL, param, streamEnv.timer, &pMeta->hbInfo.hbTmr);
+  taosReleaseRef(streamMetaId, rid);
 }