From c94c43951e427b65b842a5c5175f69ab1436e2bf Mon Sep 17 00:00:00 2001
From: slzhou <slzhou@taosdata.com>
Date: Tue, 17 May 2022 11:13:11 +0800
Subject: [PATCH] fix: sanitize udf memory usage

---
 source/libs/function/src/tudf.c    | 17 +++++++++++------
 source/libs/function/test/runUdf.c |  6 +++++-
 2 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/source/libs/function/src/tudf.c b/source/libs/function/src/tudf.c
index 388ec28b76..bf2fbce79e 100644
--- a/source/libs/function/src/tudf.c
+++ b/source/libs/function/src/tudf.c
@@ -795,7 +795,6 @@ int32_t convertScalarParamToDataBlock(SScalarParam *input, int32_t numOfCols, SS
   }
   output->info.hasVarCol = hasVarCol;
 
-  //TODO: free the array output->pDataBlock
   output->pDataBlock = taosArrayInit(numOfCols, sizeof(SColumnInfoData));
   for (int32_t i = 0; i < numOfCols; ++i) {
     taosArrayPush(output->pDataBlock, (input + i)->columnData);
@@ -809,8 +808,12 @@ int32_t convertDataBlockToScalarParm(SSDataBlock *input, SScalarParam *output) {
     return -1;
   }
   output->numOfRows = input->info.rows;
-  //TODO: memory
-  output->columnData = taosArrayGet(input->pDataBlock, 0);
+
+  output->columnData = taosMemoryMalloc(sizeof(SColumnInfoData));
+  memcpy(output->columnData,
+         taosArrayGet(input->pDataBlock, 0),
+         sizeof(SColumnInfoData));
+
   return 0;
 }
 
@@ -1427,7 +1430,10 @@ int32_t doCallUdfScalarFunc(UdfcFuncHandle handle, SScalarParam *input, int32_t
   int32_t err = callUdf(handle, callType, &inputBlock, NULL, NULL, &resultBlock, NULL);
   if (err == 0) {
     convertDataBlockToScalarParm(&resultBlock, output);
+    taosArrayDestroy(resultBlock.pDataBlock);
   }
+
+  taosArrayDestroy(inputBlock.pDataBlock);
   return err;
 }
 
@@ -1508,16 +1514,15 @@ int32_t doTeardownUdf(UdfcFuncHandle handle) {
 
   udfcRunUdfUvTask(task, UV_TASK_REQ_RSP);
 
-  SUdfTeardownResponse *rsp = &task->_teardown.rsp;
   int32_t err = task->errCode;
 
   udfcRunUdfUvTask(task, UV_TASK_DISCONNECT);
 
+  fnInfo("tear down udf. udf name: %s, udf func handle: %p", session->udfName, handle);
+
   taosMemoryFree(task->session);
   taosMemoryFree(task);
 
-  fnInfo("tear down udf. udf name: %s, udf func handle: %p", session->udfName, handle);
-
   return err;
 }
 
diff --git a/source/libs/function/test/runUdf.c b/source/libs/function/test/runUdf.c
index 329da87fcf..4490c2d5e3 100644
--- a/source/libs/function/test/runUdf.c
+++ b/source/libs/function/test/runUdf.c
@@ -77,11 +77,15 @@ int main(int argc, char *argv[]) {
   input.columnData = taosArrayGet(pBlock->pDataBlock, 0);
   SScalarParam output = {0};
   doCallUdfScalarFunc(handle, &input, 1, &output);
-
+  taosArrayDestroy(pBlock->pDataBlock);
   SColumnInfoData *col = output.columnData;
   for (int32_t i = 0; i < output.numOfRows; ++i) {
     fprintf(stderr, "%d\t%d\n", i, *(int32_t *)(col->pData + i * sizeof(int32_t)));
   }
+
+  colDataDestroy(output.columnData);
+  taosMemoryFree(output.columnData);
+
   doTeardownUdf(handle);
   udfcClose();
 }
-- 
GitLab