From 820715cd62d2c2a8a5dd562772808e6bccbbb3c2 Mon Sep 17 00:00:00 2001
From: An Tao <antao2002@gmail.com>
Date: Fri, 31 Jul 2020 09:25:20 +0800
Subject: [PATCH] Set session to requests for websockets (#521)

---
 lib/src/HttpAppFrameworkImpl.cc | 30 +++++++++++++++++-------------
 lib/src/HttpAppFrameworkImpl.h  |  2 ++
 2 files changed, 19 insertions(+), 13 deletions(-)

diff --git a/lib/src/HttpAppFrameworkImpl.cc b/lib/src/HttpAppFrameworkImpl.cc
index 3fd5f72e..a10cfec7 100644
--- a/lib/src/HttpAppFrameworkImpl.cc
+++ b/lib/src/HttpAppFrameworkImpl.cc
@@ -514,7 +514,6 @@ void HttpAppFrameworkImpl::run()
                                                      // TODO: new plugin
                                                  });
     }
-
     getLoop()->queueInLoop([this]() {
         // Let listener event loops run when everything is ready.
         listenerManagerPtr_->startListening();
@@ -614,11 +613,27 @@ HttpAppFramework &HttpAppFrameworkImpl::setUploadPath(
     }
     return *this;
 }
+void HttpAppFrameworkImpl::findSessionForRequest(const HttpRequestImplPtr &req)
+{
+    if (useSession_)
+    {
+        std::string sessionId = req->getCookie("JSESSIONID");
+        bool needSetJsessionid = false;
+        if (sessionId.empty())
+        {
+            sessionId = utils::getUuid();
+            needSetJsessionid = true;
+        }
+        req->setSession(
+            sessionManagerPtr_->getSession(sessionId, needSetJsessionid));
+    }
+}
 void HttpAppFrameworkImpl::onNewWebsockRequest(
     const HttpRequestImplPtr &req,
     std::function<void(const HttpResponsePtr &)> &&callback,
     const WebSocketConnectionImplPtr &wsConnPtr)
 {
+    findSessionForRequest(req);
     // Route to controller
     if (!preRoutingObservers_.empty())
     {
@@ -744,18 +759,7 @@ void HttpAppFrameworkImpl::onAsyncRequest(
         callback(resp);
         return;
     }
-    if (useSession_)
-    {
-        std::string sessionId = req->getCookie("JSESSIONID");
-        bool needSetJsessionid = false;
-        if (sessionId.empty())
-        {
-            sessionId = utils::getUuid();
-            needSetJsessionid = true;
-        }
-        req->setSession(
-            sessionManagerPtr_->getSession(sessionId, needSetJsessionid));
-    }
+    findSessionForRequest(req);
     // Route to controller
     if (!preRoutingObservers_.empty())
     {
diff --git a/lib/src/HttpAppFrameworkImpl.h b/lib/src/HttpAppFrameworkImpl.h
index 0fc0ac1a..9b49bf27 100644
--- a/lib/src/HttpAppFrameworkImpl.h
+++ b/lib/src/HttpAppFrameworkImpl.h
@@ -512,6 +512,8 @@ class HttpAppFrameworkImpl : public HttpAppFramework
                      const std::vector<HttpMethod> &validMethods,
                      const std::vector<std::string> &filters);
 
+    void findSessionForRequest(const HttpRequestImplPtr &req);
+
     // We use a uuid string as session id;
     // set sessionTimeout_=0 to make location session valid forever based on
     // cookies;