Add check the client connection status (drogonframework#2191)

joestarzxh · Oct 25, 2024 · 23c561f · 23c561f
1 parent 284d14b
commit 23c561f
Show file tree

Hide file tree

Showing 5 changed files with 28 additions and 2 deletions.
diff --git a/examples/helloworld/main.cc b/examples/helloworld/main.cc
@@ -1,3 +1,4 @@
+#include <trantor/utils/Logger.h>
 #ifdef _WIN32
 #include <ws2tcpip.h>
 #else
@@ -15,8 +16,10 @@ int main()
     // sent to Drogon
     app().registerHandler(
         "/",
-        [](const HttpRequestPtr &,
+        [](const HttpRequestPtr &request,
            std::function<void(const HttpResponsePtr &)> &&callback) {
+            LOG_INFO << "connected:"
+                     << (request->connected() ? "true" : "false");
             auto resp = HttpResponse::newHttpResponse();
             resp->setBody("Hello, World!");
             callback(resp);

diff --git a/lib/inc/drogon/HttpRequest.h b/lib/inc/drogon/HttpRequest.h
@@ -504,6 +504,8 @@ class DROGON_EXPORT HttpRequest
     virtual void setContentTypeString(const char *typeString,
                                       size_t typeStringLength) = 0;
 
+    virtual bool connected() const noexcept = 0;
+
     virtual ~HttpRequest()
     {
     }

diff --git a/lib/src/HttpRequestImpl.cc b/lib/src/HttpRequestImpl.cc
@@ -600,6 +600,7 @@ void HttpRequestImpl::swap(HttpRequestImpl &that) noexcept
     swap(streamFinishCb_, that.streamFinishCb_);
     swap(streamExceptionPtr_, that.streamExceptionPtr_);
     swap(startProcessing_, that.startProcessing_);
+    swap(connPtr_, that.connPtr_);
 }
 
 const char *HttpRequestImpl::versionString() const

diff --git a/lib/src/HttpRequestImpl.h b/lib/src/HttpRequestImpl.h
@@ -26,9 +26,12 @@
 #include <trantor/utils/Logger.h>
 #include <trantor/utils/MsgBuffer.h>
 #include <trantor/utils/NonCopyable.h>
+#include <trantor/net/TcpConnection.h>
 #include <algorithm>
+#include <functional>
+#include <memory>
 #include <string>
-#include <thread>
+#include <future>
 #include <unordered_map>
 #include <assert.h>
 #include <stdio.h>
@@ -97,6 +100,7 @@ class HttpRequestImpl : public HttpRequest
         streamFinishCb_ = nullptr;
         streamExceptionPtr_ = nullptr;
         startProcessing_ = false;
+        connPtr_.reset();
     }
 
     trantor::EventLoop *getLoop()
@@ -326,6 +330,11 @@ class HttpRequestImpl : public HttpRequest
         peerCertificate_ = cert;
     }
 
+    void setConnectionPtr(const std::shared_ptr<trantor::TcpConnection> &ptr)
+    {
+        connPtr_ = ptr;
+    }
+
     void addHeader(const char *start, const char *colon, const char *end);
 
     void removeHeader(std::string key) override
@@ -554,6 +563,15 @@ class HttpRequestImpl : public HttpRequest
         return keepAlive_;
     }
 
+    bool connected() const noexcept override
+    {
+        if (auto conn = connPtr_.lock())
+        {
+            return conn->connected();
+        }
+        return false;
+    }
+
     bool isOnSecureConnection() const noexcept override
     {
         return isOnSecureConnection_;
@@ -705,6 +723,7 @@ class HttpRequestImpl : public HttpRequest
     RequestStreamReaderPtr streamReaderPtr_;
     std::exception_ptr streamExceptionPtr_;
     bool startProcessing_{false};
+    std::weak_ptr<trantor::TcpConnection> connPtr_;
 
   protected:
     std::string content_;

diff --git a/lib/src/HttpServer.cc b/lib/src/HttpServer.cc
@@ -220,6 +220,7 @@ void HttpServer::onMessage(const TcpConnectionPtr &conn, MsgBuffer *buf)
             req->setCreationDate(trantor::Date::date());
             req->setSecure(conn->isSSLConnection());
             req->setPeerCertificate(conn->peerCertificate());
+            req->setConnectionPtr(conn);
             // TODO: maybe call onRequests() directly in stream mode
             requests.push_back(req);
         }