From ac0a1b873e1b0c8f6ef281a8f907a1cfe4240a74 Mon Sep 17 00:00:00 2001
From: An Tao <antao2002@gmail.com>
Date: Wed, 2 Oct 2024 16:42:40 +0800
Subject: [PATCH] Fix a bug after removing content-length header in some
 responses (#2176)

---
 lib/src/HttpResponseImpl.cc | 22 +++++++++++++++++-----
 lib/src/HttpResponseImpl.h  | 20 ++++++++++----------
 lib/src/HttpServer.cc       |  4 ++++
 3 files changed, 31 insertions(+), 15 deletions(-)

diff --git a/lib/src/HttpResponseImpl.cc b/lib/src/HttpResponseImpl.cc
index 1e0cc1e1..5716e34e 100644
--- a/lib/src/HttpResponseImpl.cc
+++ b/lib/src/HttpResponseImpl.cc
@@ -526,7 +526,18 @@ void HttpResponseImpl::makeHeaderString(trantor::MsgBuffer &buffer)
     if (!passThrough_)
     {
         buffer.ensureWritableBytes(64);
-        if (streamCallback_ || asyncStreamCallback_)
+        if (!contentLengthIsAllowed())
+        {
+            len = 0;
+            if ((bodyPtr_ && bodyPtr_->length() > 0) ||
+                !sendfileName_.empty() || streamCallback_ ||
+                asyncStreamCallback_)
+            {
+                LOG_ERROR << "The body should be empty when the content-length "
+                             "is not allowed!";
+            }
+        }
+        else if (streamCallback_ || asyncStreamCallback_)
         {
             // When the headers are created, it is time to set the transfer
             // encoding to chunked if the contents size is not specified
@@ -538,7 +549,7 @@ void HttpResponseImpl::makeHeaderString(trantor::MsgBuffer &buffer)
             }
             len = 0;
         }
-        else if (sendfileName_.empty() && contentLengthIsAllowed())
+        else if (sendfileName_.empty())
         {
             auto bodyLength = bodyPtr_ ? bodyPtr_->length() : 0;
             len = snprintf(buffer.beginWrite(),
@@ -546,7 +557,7 @@ void HttpResponseImpl::makeHeaderString(trantor::MsgBuffer &buffer)
                            contentLengthFormatString<decltype(bodyLength)>(),
                            bodyLength);
         }
-        else if (contentLengthIsAllowed())
+        else
         {
             auto bodyLength = sendfileRange_.second;
             len = snprintf(buffer.beginWrite(),
@@ -629,7 +640,7 @@ void HttpResponseImpl::renderToBuffer(trantor::MsgBuffer &buffer)
     {
         buffer.append("\r\n");
     }
-    if (bodyPtr_)
+    if (bodyPtr_ && contentLengthIsAllowed())
         buffer.append(bodyPtr_->data(), bodyPtr_->length());
 }
 
@@ -958,7 +969,8 @@ bool HttpResponseImpl::shouldBeCompressed() const
 {
     if (streamCallback_ || asyncStreamCallback_ || !sendfileName_.empty() ||
         contentType() >= CT_APPLICATION_OCTET_STREAM ||
-        getBody().length() < 1024 || !(getHeaderBy("content-encoding").empty()))
+        getBody().length() < 1024 ||
+        !(getHeaderBy("content-encoding").empty()) || !contentLengthIsAllowed())
     {
         return false;
     }
diff --git a/lib/src/HttpResponseImpl.h b/lib/src/HttpResponseImpl.h
index 61c10e1b..d6b949c8 100644
--- a/lib/src/HttpResponseImpl.h
+++ b/lib/src/HttpResponseImpl.h
@@ -402,6 +402,16 @@ class DROGON_EXPORT HttpResponseImpl : public HttpResponse
             addHeader("content-length", std::to_string(bodyPtr_->length()));
         }
     }
+
+    bool contentLengthIsAllowed() const
+    {
+        int statusCode =
+            customStatusCode_ >= 0 ? customStatusCode_ : statusCode_;
+
+        // return false if status code is 1xx or 204
+        return (statusCode >= k200OK || statusCode < k100Continue) &&
+               statusCode != k204NoContent;
+    }
 #ifdef USE_BROTLI
     void brDecompress()
     {
@@ -537,16 +547,6 @@ class DROGON_EXPORT HttpResponseImpl : public HttpResponse
     {
         statusMessage_ = message;
     }
-
-    bool contentLengthIsAllowed()
-    {
-        int statusCode =
-            customStatusCode_ >= 0 ? customStatusCode_ : statusCode_;
-
-        // return false if status code is 1xx or 204
-        return (statusCode >= k200OK || statusCode < k100Continue) &&
-               statusCode != k204NoContent;
-    }
 };
 
 using HttpResponseImplPtr = std::shared_ptr<HttpResponseImpl>;
diff --git a/lib/src/HttpServer.cc b/lib/src/HttpServer.cc
index 3d8ca852..1cf60ad3 100644
--- a/lib/src/HttpServer.cc
+++ b/lib/src/HttpServer.cc
@@ -966,6 +966,8 @@ void HttpServer::sendResponse(const TcpConnectionPtr &conn,
     {
         auto httpString = respImplPtr->renderToBuffer();
         conn->send(httpString);
+        if (!respImplPtr->contentLengthIsAllowed())
+            return;
         auto &asyncStreamCallback = respImplPtr->asyncStreamCallback();
         if (asyncStreamCallback)
         {
@@ -1046,6 +1048,8 @@ void HttpServer::sendResponses(
         {
             // Not HEAD method
             respImplPtr->renderToBuffer(buffer);
+            if (!respImplPtr->contentLengthIsAllowed())
+                continue;
             auto &asyncStreamCallback = respImplPtr->asyncStreamCallback();
             if (asyncStreamCallback)
             {