Remove content-length header from 101 Switching Protocols response (#…

…2164)
drogonframework · Sep 25, 2024 · b0c5331 · b0c5331
1 parent c9f5754
commit b0c5331
Show file tree

Hide file tree

Showing 4 changed files with 54 additions and 2 deletions.
diff --git 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,15 +538,15 @@ 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(),
                            buffer.writableBytes(),
                            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
@@ -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
@@ -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
@@ -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);
+    });
+}