Merge pull request #52 from pingfang-tx/master

Support server cookie

servant/libservant/AsyncProcThread.cpp

@@ -109,6 +109,9 @@ void AsyncProcThread::callback(ReqMessage * msg)
 	pServantProxyThreadData->_dyeing  = msg->bDyeing;
 	pServantProxyThreadData->_dyeingKey = msg->sDyeingKey;
 
+	pServantProxyThreadData->_hasCookie  = msg->hasCookie;
+	pServantProxyThreadData->_cookie = msg->cookie;
+
 	if(msg->adapter)
 	{
 		pServantProxyThreadData->_szHost = msg->adapter->endpoint().desc();

servant/libservant/ServantHandle.cpp

@@ -22,6 +22,7 @@
 #include "servant/AppProtocol.h"
 #include "servant/BaseF.h"
 #include "servant/TarsNodeF.h"
+#include "servant/TarsCookie.h"
 #ifdef _USE_OPENTRACKING
 #include "servant/text_map_carrier.h"
 #endif
@@ -704,6 +705,24 @@ bool ServantHandle::processDye(const TarsCurrentPtr &current, string& dyeingKey)
     return false;
 }
 
+
+bool ServantHandle::processCookie(const TarsCurrentPtr &current, map<string, string> &cookie)
+{
+    map<string, string>::const_iterator cookieIt = current->getContext().find(ServantProxy::STATUS_COOKIE);
+    if (cookieIt != current->getContext().end())
+    {
+        TLOGINFO("[TARS] cookie:" << cookieIt->second << endl);
+
+        Cookie stCookie;
+        stCookie.readFromJsonString(cookieIt->second);
+        cookie = stCookie.cookie;
+        return true;
+    }
+
+    return false;
+
+}
+
 bool ServantHandle::checkValidSetInvoke(const TarsCurrentPtr &current)
 {
     /*是否允许检查合法性*/
@@ -804,11 +823,21 @@ void ServantHandle::handleTarsProtocol(const TarsCurrentPtr &current)
     //处理染色消息
     string dyeingKey = "";
     TarsDyeingSwitch dyeSwitch;
-
     if (processDye(current, dyeingKey))
     {
         dyeSwitch.enableDyeing(dyeingKey);
     }
+
+    //处理cookie
+    map<string, string> cookie;
+    TarsCookieOp cookieOp;
+    if (processCookie(current, cookie))
+    {
+        cookieOp.setCookie(cookie);
+        current->setCookie(cookie);
+    }
+
+
 #ifdef _USE_OPENTRACKING
     //处理tracking信息
     processTracking(current);

servant/libservant/ServantProxy.cpp

@@ -233,6 +233,9 @@ string ServantProxy::TARS_MASTER_KEY      = "TARS_MASTER_KEY";
 
 string ServantProxy::STATUS_TRACK_KEY     = "STATUS_TRACK_KEY";
 
+string ServantProxy::STATUS_COOKIE       = "STATUS_COOKIE";
+
+
 ////////////////////////////////////
 ServantProxy::ServantProxy(Communicator * pCommunicator, ObjectProxy ** ppObjectProxy, size_t iClientThreadNum)
 : _communicator(pCommunicator)
@@ -570,6 +573,9 @@ void ServantProxy::invoke(ReqMessage * msg, bool bCoroAsync)
         TLOGTARS("[TARS][ServantProxy::invoke, set dyeing, key=" << pSptd->_dyeingKey << endl);
     }
 
+    msg->hasCookie    = pSptd->_hasCookie;
+    msg->cookie       = pSptd->_cookie;
+
 
 #ifdef _USE_OPENTRACKING
     msg->trackInfoMap = pSptd->_trackInfoMap;
@@ -779,6 +785,8 @@ void ServantProxy::tars_invoke_async(char  cPacketType,
 
     checkDye(msg->request);
 
+    checkCookie(msg->request);
+
     invoke(msg, bCoro);
 }
 
@@ -812,6 +820,8 @@ shared_ptr<ResponsePacket> ServantProxy::tars_invoke(char  cPacketType,
 
     checkDye(msg->request);
 
+    checkCookie(msg->request);
+
     invoke(msg);
 
     shared_ptr<ResponsePacket> rsp = msg->response;
@@ -977,6 +987,19 @@ void ServantProxy::checkDye(RequestPacket& req)
     }
 }
 
+void ServantProxy::checkCookie(RequestPacket& req)
+{
+    //线程私有数据
+    ServantProxyThreadData * pSptd = ServantProxyThreadData::getData();
+    assert(pSptd != NULL);
+    if(pSptd->_hasCookie)
+    {
+        Cookie stCookie;
+        stCookie.cookie = pSptd->_cookie;
+        req.context[ServantProxy::STATUS_COOKIE] = stCookie.writeToJsonString();
+    }
+}
+
 
 void ServantProxy::tars_endpoints(vector<EndpointInfo> &activeEndPoint, vector<EndpointInfo> &inactiveEndPoint)
 {

servant/servant/Message.h

@@ -203,6 +203,7 @@ struct ReqMessage : public TC_HandleBase
     , bCoroFlag(false)
     , sched(NULL)
     , iCoroId(0)
+    , hasCookie(0)
     {
     }
 
@@ -250,6 +251,7 @@ struct ReqMessage : public TC_HandleBase
         bCoroFlag      = false;
         sched          = NULL;
         iCoroId        = 0;
+        hasCookie      = false;
     }
 
 
@@ -292,6 +294,9 @@ struct ReqMessage : public TC_HandleBase
 #ifdef _USE_OPENTRACKING
     std::unordered_map<std::string, std::string> trackInfoMap; //调用链信息
 #endif
+
+    bool                            hasCookie;       // 是否包含cookie
+    map<string, string>             cookie;          // cookie内容 
 };
 
 typedef TC_AutoPtr<ReqMessage>          ReqMessagePtr;

servant/servant/ServantHandle.h

@@ -186,6 +186,11 @@ protected:
      */
     bool processDye(const TarsCurrentPtr &current, string& dyeingKey);
 
+    /**
+     * 处理cookie
+     */
+    bool processCookie(const TarsCurrentPtr &current, map<string, string> &cookie);
+
     /**
      * 检查set调用合法性
      *

servant/servant/ServantProxy.h

@@ -151,6 +151,12 @@ public:
 #ifdef _USE_OPENTRACKING
     std::unordered_map<std::string, std::string> _trackInfoMap;
 #endif
+
+    /**
+     * cookie
+     */
+    bool                            _hasCookie;       // 是否包含cookie
+    map<string, string>             _cookie;          // cookie内容 
 };
 
 
@@ -406,6 +412,8 @@ public:
 
     static string STATUS_TRACK_KEY; //track信息
 
+    static string STATUS_COOKIE; //cookie信息
+
     /**
      * 缺省的同步调用超时时间
      * 超时后不保证消息不会被服务端处理
@@ -716,6 +724,13 @@ private:
      * @param  req
      */
     void checkDye(RequestPacket& req);
+
+    /**
+     * 检查是否需要设置cookie
+     * @param  req
+     */
+    void checkCookie(RequestPacket &req);
+
 private:
     friend class ObjectProxy;
     friend class AdapterProxy;

servant/servant/TarsCookie.h

@@ -0,0 +1,85 @@
+/**
+ * Tencent is pleased to support the open source community by making Tars available.
+ *
+ * Copyright (C) 2016THL A29 Limited, a Tencent company. All rights reserved.
+ *
+ * Licensed under the BSD 3-Clause License (the "License"); you may not use this file except 
+ * in compliance with the License. You may obtain a copy of the License at
+ *
+ * https://opensource.org/licenses/BSD-3-Clause
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed 
+ * under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR 
+ * CONDITIONS OF ANY KIND, either express or implied. See the License for the 
+ * specific language governing permissions and limitations under the License.
+ */
+
+#ifndef __TARS_COOKIE_H__
+#define __TARS_COOKIE_H__
+
+
+/**
+ * cookie操作类
+ */
+class TarsCookieOp
+{
+public:
+    /**
+     * 构造函数
+     */
+    TarsCookieOp()
+    :_hasCookie(false)
+    {
+    }
+
+    /**
+     * 析构函数，清理掉已设置的cookie
+     */
+    ~TarsCookieOp()
+    {
+        if(_hasCookie)
+        {
+            ServantProxyThreadData * td = ServantProxyThreadData::getData();
+            assert(NULL != td);
+            if (td)
+            {
+                td->_hasCookie = false;
+                td->_cookie.clear();
+            }
+            _hasCookie = false;
+        }
+    }
+
+    /**
+     * 获取cookie
+     */
+    static map<string, string> & getCookie()
+    {
+        ServantProxyThreadData * td = ServantProxyThreadData::getData();
+        assert(NULL != td);
+
+        return td->_cookie;
+    }
+
+    /**
+     * 设置cookie
+     */
+    void setCookie(const map<string, string> &cookie)
+    {
+        ServantProxyThreadData * td = ServantProxyThreadData::getData();
+        assert(NULL != td);
+        if(td)
+        {
+            td->_hasCookie = true;
+            td->_cookie = cookie;
+        }
+
+        _hasCookie = true;
+
+    }
+
+protected:
+    bool _hasCookie;
+};
+
+#endif

servant/servant/TarsCurrent.h

@@ -229,6 +229,22 @@ public:
      */
     void sendResponse(const char* buff, uint32_t len);
 
+    /**
+     * 设置cookie
+     */
+    void setCookie(const map<string, string> &cookie)
+    {
+        _cookie = cookie;
+    }
+
+    /**
+     * 获取cookie
+     */
+    map<string, string> & getCookie()
+    {
+        return _cookie;
+    }
+
 protected:
 
     friend class ServantHandle;
@@ -303,6 +319,11 @@ protected:
      * 设置额外返回的内容
      */
     map<std::string, std::string> _responseContext;
+
+    /**
+     * cookie
+     */
+    map<string, string>             _cookie;
 };
 //////////////////////////////////////////////////////////////
 }

servant/tup/RequestF.tars

@@ -44,4 +44,10 @@ module tars
         8 optional string        sResultDesc;
         9 optional map<string, string> context;
     };
+
+    // Cookie, 在服务调用链上透传
+    struct Cookie
+    {
+        1 optional map<string, string>         cookie;
+    };
 };