From c3130c7f44ab1cb2cb1f59e13802decb493139ef Mon Sep 17 00:00:00 2001
From: Jeff Tao <jhtao@taosdata.com>
Date: Sun, 1 Mar 2020 10:48:22 +0800
Subject: [PATCH] optimize the authentication part

---
 src/rpc/src/rpcMain.c  | 24 +++++++++++++++---------
 src/rpc/test/rclient.c |  7 ++++---
 src/rpc/test/rserver.c |  2 +-
 3 files changed, 20 insertions(+), 13 deletions(-)

diff --git a/src/rpc/src/rpcMain.c b/src/rpc/src/rpcMain.c
index b661b42f5b..17e8a94408 100755
--- a/src/rpc/src/rpcMain.c
+++ b/src/rpc/src/rpcMain.c
@@ -94,6 +94,7 @@ typedef struct _RpcConn {
   char      encrypt; // encryption, 0:1 
   char      secret[TSDB_KEY_LEN]; // secret for the link
   char      ckey[TSDB_KEY_LEN];   // ciphering key 
+  char      secured;              // if set to 1, no authentication
   uint16_t  localPort;      // for UDP only
   uint32_t  peerUid;        // peer UID
   uint32_t  peerIp;         // peer IP
@@ -264,7 +265,7 @@ void *rpcOpen(SRpcInit *pInit) {
       return NULL;
     }
   } else {
-    pRpc->pCache = rpcOpenConnCache(pRpc->sessions, rpcCloseConn, pRpc->tmrCtrl, tsShellActivityTimer*1000); 
+    pRpc->pCache = rpcOpenConnCache(pRpc->sessions, rpcCloseConn, pRpc->tmrCtrl, pRpc->idleTime); 
     if ( pRpc->pCache == NULL ) {
       tError("%s failed to init connection cache", pRpc->label);
       rpcClose(pRpc);
@@ -417,6 +418,7 @@ void rpcSendResponse(void *handle, int32_t code, void *pCont, int contLen) {
 
   taosTmrStopA(&pConn->pTimer);
   rpcSendMsgToPeer(pConn, msg, msgLen);
+  pConn->secured = 1; // connection shall be secured
 
   return;
 }
@@ -811,7 +813,8 @@ static void *rpcProcessMsgFromPeer(SRecvInfo *pRecv) {
         pRecv->msgLen, pHead->sourceId, pHead->destId, pHead->tranId, pHead->port);
   }
 
-  if (pConn && pRpc->idleTime) {
+  if (pRpc->connType == TAOS_CONN_SERVER && pConn && pRpc->idleTime) {  
+    // only for server, starts the idle timer. For client, it is started by cache mgmt
     taosTmrReset(rpcProcessIdleTimer, pRpc->idleTime, pConn, pRpc->tmrCtrl, &pConn->pIdleTimer);
   }
 
@@ -1023,8 +1026,8 @@ static void rpcProcessRetryTimer(void *param, void *tmrId) {
     pConn->retry++;
 
     if (pConn->retry < 4) {
-      tTrace("%s %p, re-send msg:%s to %s:%hu retry:%d", pRpc->label, pConn, 
-             taosMsg[pConn->outType], pConn->peerIpstr, pConn->peerPort, pConn->retry);
+      tTrace("%s %p, re-send msg:%s to %s:%hud", pRpc->label, pConn, 
+             taosMsg[pConn->outType], pConn->peerIpstr, pConn->peerPort);
       rpcSendMsgToPeer(pConn, pConn->pReqMsg, pConn->reqMsgLen);      
       taosTmrReset(rpcProcessRetryTimer, tsRpcTimer, pConn, pRpc->tmrCtrl, &pConn->pTimer);
     } else {
@@ -1176,7 +1179,7 @@ static void rpcBuildAuthHead(void *pMsg, int msgLen, void *pAuth, void *pKey) {
 static int rpcAddAuthPart(SRpcConn *pConn, char *msg, int msgLen) {
   SRpcHead *pHead = (SRpcHead *)msg;
 
-  if (pConn->spi) {
+  if (pConn->spi && pConn->secured == 0) {
     // add auth part
     pHead->spi = pConn->spi;
     SRpcDigest *pDigest = (SRpcDigest *)(msg + msgLen);
@@ -1185,6 +1188,7 @@ static int rpcAddAuthPart(SRpcConn *pConn, char *msg, int msgLen) {
     pHead->msgLen = (int32_t)htonl((uint32_t)msgLen);
     rpcBuildAuthHead(pHead, msgLen - TSDB_AUTH_LEN, pDigest->auth, pConn->secret);
   } else {
+    pHead->spi = 0;
     pHead->msgLen = (int32_t)htonl((uint32_t)msgLen);
   }
 
@@ -1194,9 +1198,10 @@ static int rpcAddAuthPart(SRpcConn *pConn, char *msg, int msgLen) {
 static int rpcCheckAuthentication(SRpcConn *pConn, char *msg, int msgLen) {
   SRpcHead *pHead = (SRpcHead *)msg;
   SRpcInfo *pRpc = pConn->pRpc;
-  int32_t   code = 0;
+  int       code = 0;
 
-  if (pConn->spi == 0) {
+  if ((pConn->secured && pHead->spi == 0) || (pHead->spi == 0 && pConn->spi == 0)){
+    // secured link, or no authentication 
     pHead->msgLen = (int32_t)htonl((uint32_t)pHead->msgLen);
     return 0;
   }
@@ -1211,7 +1216,6 @@ static int rpcCheckAuthentication(SRpcConn *pConn, char *msg, int msgLen) {
   }
  
   code = 0;
-
   if (pHead->spi == pConn->spi) {
     // authentication
     SRpcDigest *pDigest = (SRpcDigest *)((char *)pHead + msgLen - sizeof(SRpcDigest));
@@ -1228,6 +1232,8 @@ static int rpcCheckAuthentication(SRpcConn *pConn, char *msg, int msgLen) {
         code = TSDB_CODE_AUTH_FAILURE;
       } else {
         pHead->msgLen = (int32_t)htonl((uint32_t)pHead->msgLen) - sizeof(SRpcDigest);
+        if ( !rpcIsReq(pHead->msgType) ) pConn->secured = 1;  // link is secured for client
+        tTrace("%s %p, message is authenticated", pRpc->label, pConn);
       }
     }
   } else {
@@ -1240,7 +1246,7 @@ static int rpcCheckAuthentication(SRpcConn *pConn, char *msg, int msgLen) {
 
 static void rpcLockConn(SRpcConn *pConn) {
   int64_t tid = taosGetPthreadId();
-  int       i = 0;
+  int     i = 0;
   while (atomic_val_compare_exchange_64(&(pConn->lockedBy), 0, tid) != 0) {
     if (++i % 1000 == 0) {
       sched_yield();
diff --git a/src/rpc/test/rclient.c b/src/rpc/test/rclient.c
index 63c23ce7bc..aa97535e31 100644
--- a/src/rpc/test/rclient.c
+++ b/src/rpc/test/rclient.c
@@ -106,11 +106,12 @@ int main(int argc, char *argv[]) {
   rpcInit.cfp          = processResponse;
   rpcInit.ufp          = processUpdateIpSet;
   rpcInit.sessions     = 100;
-  rpcInit.idleTime     = 2000;
+  rpcInit.idleTime     = tsShellActivityTimer*1000;
   rpcInit.user         = "michael";
   rpcInit.secret       = "mypassword";
   rpcInit.ckey         = "key";
   rpcInit.spi          = 1;
+  rpcInit.connType     = TAOS_CONN_CLIENT;
 
   for (int i=1; i<argc; ++i) { 
     if (strcmp(argv[i], "-p")==0 && i < argc-1) {
@@ -159,8 +160,8 @@ int main(int argc, char *argv[]) {
     }
   }
 
-  rpcInit.connType = TAOS_CONN_CLIENT;
   taosInitLog("client.log", 100000, 10);
+  tPrint("rpcDebugFlag:%d", rpcDebugFlag);
 
   void *pRpc = rpcOpen(&rpcInit);
   if (pRpc == NULL) {
@@ -200,7 +201,7 @@ int main(int argc, char *argv[]) {
   tPrint("it takes %.3f mseconds to send %d requests to server", usedTime, numOfReqs*appThreads);
   tPrint("Performance: %.3f requests per second, msgSize:%d bytes", 1000.0*numOfReqs*appThreads/usedTime, msgSize);
 
-  taosCloseLog();
+  taosCloseLogger();
 
   return 0;
 }
diff --git a/src/rpc/test/rserver.c b/src/rpc/test/rserver.c
index 84f0d50d44..718b0c5b6e 100644
--- a/src/rpc/test/rserver.c
+++ b/src/rpc/test/rserver.c
@@ -110,7 +110,7 @@ int main(int argc, char *argv[]) {
   rpcInit.numOfThreads = 1;
   rpcInit.cfp          = processRequestMsg;
   rpcInit.sessions     = 1000;
-  rpcInit.idleTime     = 2000;
+  rpcInit.idleTime     = tsShellActivityTimer*1500; 
   rpcInit.afp          = retrieveAuthInfo;
 
   for (int i=1; i<argc; ++i) {
-- 
GitLab