From 23c561f0725d77c78b0629c0e441c1cdd385e9d9 Mon Sep 17 00:00:00 2001
From: fantasy-peak <82742316+fantasy-peak@users.noreply.github.com>
Date: Fri, 25 Oct 2024 16:08:03 +0800
Subject: [PATCH] Add check the client connection status (#2191)

---
 examples/helloworld/main.cc  |  5 ++++-
 lib/inc/drogon/HttpRequest.h |  2 ++
 lib/src/HttpRequestImpl.cc   |  1 +
 lib/src/HttpRequestImpl.h    | 21 ++++++++++++++++++++-
 lib/src/HttpServer.cc        |  1 +
 5 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/examples/helloworld/main.cc b/examples/helloworld/main.cc
index e92d27ca..fd25d524 100644
--- 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
index 68c4271b..1bc9eabb 100644
--- 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
index 83fc771c..dc6f68f7 100644
--- 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
index 412aecf9..bd0a4689 100644
--- 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
index 1cf60ad3..257976b1 100644
--- 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);
         }