From 01e348cc0182111863b2d819700719255c721726 Mon Sep 17 00:00:00 2001
From: Bomin Zhang <bmzhang@taosdata.com>
Date: Wed, 22 Jul 2020 16:23:52 +0800
Subject: [PATCH] fix td-980

---
 src/query/src/qExecutor.c |  6 ++++-
 src/tsdb/src/tsdbMeta.c   |  5 +++-
 src/tsdb/src/tsdbRead.c   | 53 ++++++++++++++++++++++++++-------------
 3 files changed, 45 insertions(+), 19 deletions(-)

diff --git a/src/query/src/qExecutor.c b/src/query/src/qExecutor.c
index 906d0cfe67..1220c5ca31 100644
--- a/src/query/src/qExecutor.c
+++ b/src/query/src/qExecutor.c
@@ -6419,8 +6419,12 @@ int32_t qDumpRetrieveResult(qinfo_t qinfo, SRetrieveTableRsp **pRsp, int32_t *co
   size += sizeof(STableIdInfo) * taosArrayGetSize(pQInfo->arrTableIdInfo);
   *contLen = size + sizeof(SRetrieveTableRsp);
 
-  // todo handle failed to allocate memory
+  // todo proper handle failed to allocate memory,
+  // current solution only avoid crash, but cannot return error code to client
   *pRsp = (SRetrieveTableRsp *)rpcMallocCont(*contLen);
+  if (*pRsp == NULL) {
+    return TSDB_CODE_QRY_OUT_OF_MEMORY;
+  }
   (*pRsp)->numOfRows = htonl(pQuery->rec.rows);
 
   int32_t code = pQInfo->code;
diff --git a/src/tsdb/src/tsdbMeta.c b/src/tsdb/src/tsdbMeta.c
index c1923f5235..b25e734694 100644
--- a/src/tsdb/src/tsdbMeta.c
+++ b/src/tsdb/src/tsdbMeta.c
@@ -123,7 +123,10 @@ int tsdbCreateTable(TSDB_REPO_T *repo, STableCfg *pCfg) {
   int   tlen2 = tsdbGetTableEncodeSize(TSDB_UPDATE_META, table);
   int   tlen = tlen1 + tlen2;
   void *buf = tsdbAllocBytes(pRepo, tlen);
-  ASSERT(buf != NULL);
+  if (buf == NULL) {
+    goto _err;
+  }
+
   if (newSuper) {
     void *pBuf = tsdbInsertTableAct(pRepo, TSDB_UPDATE_META, buf, super);
     ASSERT(POINTER_DISTANCE(pBuf, buf) == tlen1);
diff --git a/src/tsdb/src/tsdbRead.c b/src/tsdb/src/tsdbRead.c
index b2204948d0..fc67c64e2e 100644
--- a/src/tsdb/src/tsdbRead.c
+++ b/src/tsdb/src/tsdbRead.c
@@ -201,18 +201,34 @@ TsdbQueryHandleT* tsdbQueryTables(TSDB_REPO_T* tsdb, STsdbQueryCond* pCond, STab
   int32_t numOfCols = pCond->numOfCols;
   
   pQueryHandle->statis = calloc(numOfCols, sizeof(SDataStatis));
+  if (pQueryHandle->statis == NULL) {
+    tsdbCleanupQueryHandle(pQueryHandle);
+    return NULL;
+  }
   pQueryHandle->pColumns = taosArrayInit(numOfCols, sizeof(SColumnInfoData));  // todo: use list instead of array?
+  if (pQueryHandle->pColumns == NULL) {
+    tsdbCleanupQueryHandle(pQueryHandle);
+    return NULL;
+  }
   
   for (int32_t i = 0; i < numOfCols; ++i) {
     SColumnInfoData  colInfo = {{0}, 0};
   
     colInfo.info = pCond->colList[i];
     colInfo.pData = calloc(1, EXTRA_BYTES + pQueryHandle->outputCapacity * pCond->colList[i].bytes);
+    if (colInfo.pData == NULL) {
+      tsdbCleanupQueryHandle(pQueryHandle);
+      return NULL;
+    }
     taosArrayPush(pQueryHandle->pColumns, &colInfo);
     pQueryHandle->statis[i].colId = colInfo.info.colId;
   }
   
   pQueryHandle->pTableCheckInfo = taosArrayInit(groupList->numOfTables, sizeof(STableCheckInfo));
+  if (pQueryHandle->pTableCheckInfo == NULL) {
+    tsdbCleanupQueryHandle(pQueryHandle);
+    return NULL;
+  }
   STsdbMeta* pMeta = tsdbGetMeta(tsdb);
   assert(pMeta != NULL);
   
@@ -2372,28 +2388,31 @@ void tsdbCleanupQueryHandle(TsdbQueryHandleT queryHandle) {
     return;
   }
   
-  size_t size = taosArrayGetSize(pQueryHandle->pTableCheckInfo);
-  for (int32_t i = 0; i < size; ++i) {
-    STableCheckInfo* pTableCheckInfo = taosArrayGet(pQueryHandle->pTableCheckInfo, i);
-    destroyTableMemIterator(pTableCheckInfo);
+  if (pQueryHandle->pTableCheckInfo != NULL) {
+    size_t size = taosArrayGetSize(pQueryHandle->pTableCheckInfo);
+    for (int32_t i = 0; i < size; ++i) {
+      STableCheckInfo* pTableCheckInfo = taosArrayGet(pQueryHandle->pTableCheckInfo, i);
+      destroyTableMemIterator(pTableCheckInfo);
 
-    if (pTableCheckInfo->pDataCols != NULL) {
-      tfree(pTableCheckInfo->pDataCols->buf);
-    }
+      if (pTableCheckInfo->pDataCols != NULL) {
+        tfree(pTableCheckInfo->pDataCols->buf);
+      }
 
-    tfree(pTableCheckInfo->pDataCols);
-    tfree(pTableCheckInfo->pCompInfo);
+      tfree(pTableCheckInfo->pDataCols);
+      tfree(pTableCheckInfo->pCompInfo);
+    }
+    taosArrayDestroy(pQueryHandle->pTableCheckInfo);
   }
 
-  taosArrayDestroy(pQueryHandle->pTableCheckInfo);
-
-   size_t cols = taosArrayGetSize(pQueryHandle->pColumns);
-   for (int32_t i = 0; i < cols; ++i) {
-     SColumnInfoData* pColInfo = taosArrayGet(pQueryHandle->pColumns, i);
-     tfree(pColInfo->pData);
-   }
+  if (pQueryHandle->pColumns != NULL) {
+    size_t cols = taosArrayGetSize(pQueryHandle->pColumns);
+    for (int32_t i = 0; i < cols; ++i) {
+      SColumnInfoData* pColInfo = taosArrayGet(pQueryHandle->pColumns, i);
+      tfree(pColInfo->pData);
+    }
+    taosArrayDestroy(pQueryHandle->pColumns);
+  }
 
-  taosArrayDestroy(pQueryHandle->pColumns);
   taosArrayDestroy(pQueryHandle->defaultLoadColumn);
   tfree(pQueryHandle->pDataBlockInfo);
   tfree(pQueryHandle->statis);
-- 
GitLab