From 7880800beea977eee6f140e3a02f7630b9e76393 Mon Sep 17 00:00:00 2001
From: Yihao Deng <luomoxyz@126.com>
Date: Tue, 23 Jul 2024 11:25:36 +0000
Subject: [PATCH] refactor transport

---
 source/libs/transport/src/transCli.c  |  5 +++-
 source/libs/transport/src/transComm.c | 35 ++++++++++++++++++---------
 source/libs/transport/src/transSvr.c  |  6 ++++-
 3 files changed, 33 insertions(+), 13 deletions(-)

diff --git a/source/libs/transport/src/transCli.c b/source/libs/transport/src/transCli.c
index 10c643f644..dc5e16e671 100644
--- a/source/libs/transport/src/transCli.c
+++ b/source/libs/transport/src/transCli.c
@@ -870,7 +870,10 @@ static int32_t specifyConnRef(SCliConn* conn, bool update, int64_t handle) {
 static void cliAllocRecvBufferCb(uv_handle_t* handle, size_t suggested_size, uv_buf_t* buf) {
   SCliConn*    conn = handle->data;
   SConnBuffer* pBuf = &conn->readBuf;
-  transAllocBuffer(pBuf, buf);
+  int32_t      code = transAllocBuffer(pBuf, buf);
+  if (code < 0) {
+    cliDestroyConn(conn, true);
+  }
 }
 static void cliRecvCb(uv_stream_t* handle, ssize_t nread, const uv_buf_t* buf) {
   // impl later
diff --git a/source/libs/transport/src/transComm.c b/source/libs/transport/src/transComm.c
index 2818a767b8..8a64b7d808 100644
--- a/source/libs/transport/src/transComm.c
+++ b/source/libs/transport/src/transComm.c
@@ -115,17 +115,20 @@ int32_t transInitBuffer(SConnBuffer* buf) {
   buf->invalid = 0;
   return 0;
 }
-int transDestroyBuffer(SConnBuffer* p) {
+int32_t transDestroyBuffer(SConnBuffer* p) {
   taosMemoryFree(p->buf);
   p->buf = NULL;
   return 0;
 }
 
-int transClearBuffer(SConnBuffer* buf) {
+int32_t transClearBuffer(SConnBuffer* buf) {
   SConnBuffer* p = buf;
   if (p->cap > BUFFER_CAP) {
     p->cap = BUFFER_CAP;
     p->buf = taosMemoryRealloc(p->buf, BUFFER_CAP);
+    if (p->buf == NULL) {
+      return TSDB_CODE_OUT_OF_MEMORY;
+    }
   }
   p->left = -1;
   p->len = 0;
@@ -134,27 +137,31 @@ int transClearBuffer(SConnBuffer* buf) {
   return 0;
 }
 
-int transDumpFromBuffer(SConnBuffer* connBuf, char** buf, int8_t resetBuf) {
+int32_t transDumpFromBuffer(SConnBuffer* connBuf, char** buf, int8_t resetBuf) {
   static const int HEADSIZE = sizeof(STransMsgHead);
-
-  SConnBuffer* p = connBuf;
+  int32_t          code = 0;
+  SConnBuffer*     p = connBuf;
   if (p->left != 0 || p->total <= 0) {
-    return -1;
+    return TSDB_CODE_INVALID_MSG;
   }
   int total = p->total;
   if (total >= HEADSIZE && !p->invalid) {
     *buf = taosMemoryCalloc(1, total);
+    if (*buf == NULL) {
+      return TSDB_CODE_OUT_OF_MEMORY;
+    }
     memcpy(*buf, p->buf, total);
-    if (transResetBuffer(connBuf, resetBuf) < 0) {
-      return -1;
+    if ((code = transResetBuffer(connBuf, resetBuf)) < 0) {
+      return code;
     }
   } else {
     total = -1;
+    return TSDB_CODE_INVALID_MSG;
   }
   return total;
 }
 
-int transResetBuffer(SConnBuffer* connBuf, int8_t resetBuf) {
+int32_t transResetBuffer(SConnBuffer* connBuf, int8_t resetBuf) {
   SConnBuffer* p = connBuf;
   if (p->total < p->len) {
     int left = p->len - p->total;
@@ -170,15 +177,18 @@ int transResetBuffer(SConnBuffer* connBuf, int8_t resetBuf) {
       if (resetBuf) {
         p->cap = BUFFER_CAP;
         p->buf = taosMemoryRealloc(p->buf, p->cap);
+        if (p->buf == NULL) {
+          return TSDB_CODE_OUT_OF_MEMORY;
+        }
       }
     }
   } else {
     ASSERTS(0, "invalid read from sock buf");
-    return -1;
+    return TSDB_CODE_INVALID_MSG;
   }
   return 0;
 }
-int transAllocBuffer(SConnBuffer* connBuf, uv_buf_t* uvBuf) {
+int32_t transAllocBuffer(SConnBuffer* connBuf, uv_buf_t* uvBuf) {
   /*
    * formate of data buffer:
    * |<--------------------------data from socket------------------------------->|
@@ -195,6 +205,9 @@ int transAllocBuffer(SConnBuffer* connBuf, uv_buf_t* uvBuf) {
     } else {
       p->cap = p->left + p->len;
       p->buf = taosMemoryRealloc(p->buf, p->cap);
+      if (p->buf == NULL) {
+        return TSDB_CODE_OUT_OF_MEMORY;
+      }
       uvBuf->base = p->buf + p->len;
       uvBuf->len = p->left;
     }
diff --git a/source/libs/transport/src/transSvr.c b/source/libs/transport/src/transSvr.c
index a5ef2341be..abffdac7a4 100644
--- a/source/libs/transport/src/transSvr.c
+++ b/source/libs/transport/src/transSvr.c
@@ -202,7 +202,11 @@ static int32_t addHandleToAcceptloop(void* arg);
 void uvAllocRecvBufferCb(uv_handle_t* handle, size_t suggested_size, uv_buf_t* buf) {
   SSvrConn*    conn = handle->data;
   SConnBuffer* pBuf = &conn->readBuf;
-  transAllocBuffer(pBuf, buf);
+  int32_t      code = transAllocBuffer(pBuf, buf);
+  if (code < 0) {
+    tError("conn %p failed to alloc buffer, since %s", conn, tstrerror(code));
+    destroyConn(conn, true);
+  }
 }
 
 // refers specifically to query or insert timeout