From 4b7f93333dacb6a8d56e62d1043944bf2270e2cc Mon Sep 17 00:00:00 2001
From: fantasy-peak <1356346239@qq.com>
Date: Sun, 9 Feb 2025 13:09:57 +0800
Subject: [PATCH] Add setting chunk callback to http request

---
 examples/client_example/main.cc | 23 +++++++++++++++++++++++
 lib/inc/drogon/HttpRequest.h    |  7 +++++++
 lib/src/HttpClientImpl.cc       |  3 ++-
 lib/src/HttpRequestImpl.h       | 12 ++++++++++++
 lib/src/HttpResponseParser.cc   | 10 +++++++++-
 lib/src/HttpResponseParser.h    |  6 +++++-
 6 files changed, 58 insertions(+), 3 deletions(-)

diff --git a/examples/client_example/main.cc b/examples/client_example/main.cc
index d783f54e0c..2ea65cac62 100644
--- a/examples/client_example/main.cc
+++ b/examples/client_example/main.cc
@@ -14,6 +14,7 @@ int nth_resp = 0;
 
 int main()
 {
+#if 1
     trantor::Logger::setLogLevel(trantor::Logger::kTrace);
     {
         auto client = HttpClient::newHttpClient("http://www.baidu.com");
@@ -75,6 +76,28 @@ int main()
         std::cout << "requestsBufferSize:" << client->requestsBufferSize()
                   << std::endl;
     }
+#else
+    {
+        auto client = HttpClient::newHttpClient("http://127.0.0.1:8848/stream");
 
+        auto req = HttpRequest::newHttpRequest();
+        req->setMethod(drogon::Get);
+        req->setPath("/stream");
+        req->setChunkCallback([](std::string data) {
+            std::cout << "recv chunk:" << data << std::endl;
+        });
+        client->sendRequest(
+            req, [](ReqResult result, const HttpResponsePtr &response) {
+                if (result != ReqResult::Ok)
+                {
+                    std::cout
+                        << "error while sending request to server! result: "
+                        << result << std::endl;
+                    return;
+                }
+                std::cout << "recv chunk done" << std::endl;
+            });
+    }
+#endif
     app().run();
 }
diff --git a/lib/inc/drogon/HttpRequest.h b/lib/inc/drogon/HttpRequest.h
index 24d8c7579d..1a0fc83cec 100644
--- a/lib/inc/drogon/HttpRequest.h
+++ b/lib/inc/drogon/HttpRequest.h
@@ -25,6 +25,7 @@
 #include <trantor/net/InetAddress.h>
 #include <trantor/net/Certificate.h>
 #include <trantor/utils/Date.h>
+#include <functional>
 #include <memory>
 #include <string>
 #include <unordered_map>
@@ -510,6 +511,12 @@ class DROGON_EXPORT HttpRequest
     virtual const std::weak_ptr<trantor::TcpConnection> &getConnectionPtr()
         const noexcept = 0;
 
+    virtual void setChunkCallback(
+        std::function<void(std::string)>) noexcept = 0;
+
+    virtual const std::function<void(std::string)> &getChunkCallback()
+        const noexcept = 0;
+
     virtual ~HttpRequest()
     {
     }
diff --git a/lib/src/HttpClientImpl.cc b/lib/src/HttpClientImpl.cc
index b8cf7596c9..5de6dd498d 100644
--- a/lib/src/HttpClientImpl.cc
+++ b/lib/src/HttpClientImpl.cc
@@ -618,7 +618,8 @@ void HttpClientImpl::onRecvMessage(const trantor::TcpConnectionPtr &connPtr,
         {
             responseParser->setForHeadMethod();
         }
-        if (!responseParser->parseResponse(msg))
+        if (!responseParser->parseResponse(msg,
+                                           firstReq.first->getChunkCallback()))
         {
             onError(ReqResult::BadResponse);
             bytesReceived_ += (msgSize - msg->readableBytes());
diff --git a/lib/src/HttpRequestImpl.h b/lib/src/HttpRequestImpl.h
index 966ef92b78..0e90ed7039 100644
--- a/lib/src/HttpRequestImpl.h
+++ b/lib/src/HttpRequestImpl.h
@@ -579,6 +579,17 @@ class HttpRequestImpl : public HttpRequest
         return connPtr_;
     }
 
+    void setChunkCallback(std::function<void(std::string)> cb) noexcept override
+    {
+        chunkCb_ = std::move(cb);
+    }
+
+    const std::function<void(std::string)> &getChunkCallback()
+        const noexcept override
+    {
+        return chunkCb_;
+    }
+
     bool isOnSecureConnection() const noexcept override
     {
         return isOnSecureConnection_;
@@ -731,6 +742,7 @@ class HttpRequestImpl : public HttpRequest
     std::exception_ptr streamExceptionPtr_;
     bool startProcessing_{false};
     std::weak_ptr<trantor::TcpConnection> connPtr_;
+    std::function<void(std::string)> chunkCb_;
 
   protected:
     std::string content_;
diff --git a/lib/src/HttpResponseParser.cc b/lib/src/HttpResponseParser.cc
index 4e0d3ad7c1..c12a4458f3 100644
--- a/lib/src/HttpResponseParser.cc
+++ b/lib/src/HttpResponseParser.cc
@@ -17,6 +17,7 @@
 #include <trantor/utils/Logger.h>
 #include <trantor/utils/MsgBuffer.h>
 #include <algorithm>
+#include <string>
 
 using namespace trantor;
 using namespace drogon;
@@ -84,7 +85,9 @@ bool HttpResponseParser::parseResponseOnClose()
 }
 
 // return false if any error
-bool HttpResponseParser::parseResponse(MsgBuffer *buf)
+bool HttpResponseParser::parseResponse(
+    MsgBuffer *buf,
+    const std::function<void(std::string)> &chunkCallback)
 {
     bool ok = true;
     bool hasMore = true;
@@ -277,6 +280,11 @@ bool HttpResponseParser::parseResponse(MsgBuffer *buf)
                         responsePtr_->bodyPtr_ =
                             std::make_shared<HttpMessageStringBody>();
                     }
+                    if (chunkCallback)
+                    {
+                        chunkCallback(
+                            std::string(buf->peek(), currentChunkLength_));
+                    }
                     responsePtr_->bodyPtr_->append(buf->peek(),
                                                    currentChunkLength_);
                     buf->retrieve(currentChunkLength_ + 2);
diff --git a/lib/src/HttpResponseParser.h b/lib/src/HttpResponseParser.h
index 6f31c1838f..3a6ba0a289 100644
--- a/lib/src/HttpResponseParser.h
+++ b/lib/src/HttpResponseParser.h
@@ -18,6 +18,7 @@
 #include <trantor/utils/NonCopyable.h>
 #include <trantor/net/TcpConnection.h>
 #include <trantor/utils/MsgBuffer.h>
+#include <functional>
 #include <list>
 #include <mutex>
 
@@ -43,7 +44,10 @@ class HttpResponseParser : public trantor::NonCopyable
     // default copy-ctor, dtor and assignment are fine
 
     // return false if any error
-    bool parseResponse(trantor::MsgBuffer *buf);
+    bool parseResponse(
+        trantor::MsgBuffer *buf,
+        const std::function<void(std::string)> &chunkCallback = nullptr);
+
     bool parseResponseOnClose();
 
     bool gotAll() const