From b0c5331bc11a67014f4bd38f4b4237d4d123d558 Mon Sep 17 00:00:00 2001
From: Chad Barth <chadebarth@gmail.com>
Date: Wed, 25 Sep 2024 04:52:41 -0500
Subject: [PATCH] Remove content-length header from 101 Switching Protocols
 response (#2164)

---
 lib/src/HttpResponseImpl.cc                  |  5 ++-
 lib/src/HttpResponseImpl.h                   | 10 +++++
 lib/tests/CMakeLists.txt                     |  1 +
 lib/tests/unittests/WebsocketResponseTest.cc | 40 ++++++++++++++++++++
 4 files changed, 54 insertions(+), 2 deletions(-)
 create mode 100644 lib/tests/unittests/WebsocketResponseTest.cc

diff --git a/lib/src/HttpResponseImpl.cc b/lib/src/HttpResponseImpl.cc
index 163102ce38..1e0cc1e125 100644
--- a/lib/src/HttpResponseImpl.cc
+++ b/lib/src/HttpResponseImpl.cc
@@ -516,6 +516,7 @@ void HttpResponseImpl::makeHeaderString(trantor::MsgBuffer &buffer)
                            statusCode_);
         }
     }
+
     buffer.hasWritten(len);
 
     if (!statusMessage_.empty())
@@ -537,7 +538,7 @@ void HttpResponseImpl::makeHeaderString(trantor::MsgBuffer &buffer)
             }
             len = 0;
         }
-        else if (sendfileName_.empty())
+        else if (sendfileName_.empty() && contentLengthIsAllowed())
         {
             auto bodyLength = bodyPtr_ ? bodyPtr_->length() : 0;
             len = snprintf(buffer.beginWrite(),
@@ -545,7 +546,7 @@ void HttpResponseImpl::makeHeaderString(trantor::MsgBuffer &buffer)
                            contentLengthFormatString<decltype(bodyLength)>(),
                            bodyLength);
         }
-        else
+        else if (contentLengthIsAllowed())
         {
             auto bodyLength = sendfileRange_.second;
             len = snprintf(buffer.beginWrite(),
diff --git a/lib/src/HttpResponseImpl.h b/lib/src/HttpResponseImpl.h
index 6d02bb3736..61c10e1ba2 100644
--- a/lib/src/HttpResponseImpl.h
+++ b/lib/src/HttpResponseImpl.h
@@ -537,6 +537,16 @@ 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/tests/CMakeLists.txt b/lib/tests/CMakeLists.txt
index 143380186b..a73e9930c2 100644
--- a/lib/tests/CMakeLists.txt
+++ b/lib/tests/CMakeLists.txt
@@ -28,6 +28,7 @@ set(UNITTEST_SOURCES
     unittests/SlashRemoverTest.cc
     unittests/UtilitiesTest.cc
     unittests/UuidUnittest.cc
+    unittests/WebsocketResponseTest.cc
 )
 
 if(DROGON_CXX_STANDARD GREATER_EQUAL 20 AND HAS_COROUTINE)
diff --git a/lib/tests/unittests/WebsocketResponseTest.cc b/lib/tests/unittests/WebsocketResponseTest.cc
new file mode 100644
index 0000000000..7d881c6479
--- /dev/null
+++ b/lib/tests/unittests/WebsocketResponseTest.cc
@@ -0,0 +1,40 @@
+#include <drogon/drogon_test.h>
+
+#include "../../lib/src/HttpRequestImpl.h"
+#include "../../lib/src/HttpResponseImpl.h"
+#include "../../lib/src/HttpControllerBinder.h"
+
+using namespace drogon;
+
+DROGON_TEST(WebsocketReponseTest)
+{
+    WebsocketControllerBinder binder;
+
+    auto reqPtr = std::make_shared<HttpRequestImpl>(nullptr);
+
+    // Value from rfc6455-1.3
+    reqPtr->addHeader("sec-websocket-key", "dGhlIHNhbXBsZSBub25jZQ==");
+
+    binder.handleRequest(reqPtr, [&](const HttpResponsePtr &resp) {
+        CHECK(resp->statusCode() == k101SwitchingProtocols);
+        CHECK(resp->headers().size() == 3);
+        CHECK(resp->getHeader("Upgrade") == "websocket");
+        CHECK(resp->getHeader("Connection") == "Upgrade");
+
+        // Value from rfc6455-1.3
+        CHECK(resp->getHeader("Sec-WebSocket-Accept") ==
+              "s3pPLMBiTxaQ9kYGzzhZRbK+xOo=");
+
+        auto implPtr = std::dynamic_pointer_cast<HttpResponseImpl>(resp);
+
+        implPtr->makeHeaderString();
+        auto buffer = implPtr->renderToBuffer();
+        auto str = std::string{buffer->peek(), buffer->readableBytes()};
+
+        CHECK(str.find("upgrade: websocket") != std::string::npos);
+        CHECK(str.find("connection: Upgrade") != std::string::npos);
+        CHECK(str.find("sec-websocket-accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=") !=
+              std::string::npos);
+        CHECK(str.find("content-length:") == std::string::npos);
+    });
+}