Close the connection on an answered upgrade body in process_request (#2637)

When pre_routing_handler or pre_request_handler answers a WebSocket upgrade request, Server::process_request returns before the drain of an unread body and closes the connection only for a 4xx/5xx. A 2xx/3xx answer to an upgrade request with a framed body left that body to be parsed as the next request on the keep-alive connection. Those two early returns now close the connection when the request has a framed body, as the Expect rejection already does.

httplib.h

@@ -15208,10 +15208,15 @@ Server::process_request(Stream &strm, const std::string &remote_addr,
   // An upgrade request has been through pre-routing before it reaches routing()
   auto is_upgrade = detail::is_websocket_upgrade(req);
   if (is_upgrade) {
+    // A response given here leaves the body unread, and it would be parsed as
+    // the next request.
+    auto close_if_answered =
+        close_after_response || detail::has_framed_body(req);
+
     if (pre_routing_handler_ &&
         pre_routing_handler_(req, res) == HandlerResponse::Handled) {
       if (res.status == -1) { res.status = StatusCode::OK_200; }
-      return write_response_with_content(strm, close_after_response, req, res);
+      return write_response_with_content(strm, close_if_answered, req, res);
     }
     // Find matching WebSocket handler
     for (const auto &entry : websocket_handlers_) {
@@ -15220,8 +15225,7 @@ Server::process_request(Stream &strm, const std::string &remote_addr,
         if (pre_request_handler_ &&
             pre_request_handler_(req, res) == HandlerResponse::Handled) {
           if (res.status == -1) { res.status = StatusCode::OK_200; }
-          return write_response_with_content(strm, close_after_response, req,
+          return write_response_with_content(strm, close_if_answered, req, res);
-                                             res);
         }
 
         // Compute accept key

test/test.cc

@@ -27701,6 +27701,62 @@ TEST(WebSocketPreRoutingTest, RunsOnceForUnmatchedUpgrade) {
   EXPECT_EQ(1, pre_request_calls);
 }
 
+// An upgrade request answered by a hook leaves its body unread, so that body
+// must not be served as the next request on the connection.
+TEST(WebSocketPreRoutingTest, AnsweredUpgradeBodyIsNotNextRequest) {
+  Server svr;
+
+  svr.set_pre_routing_handler([](const Request &req, Response &res) {
+    if (req.path == "/pre-routing") {
+      res.set_content("pre-routing", "text/plain");
+      return Server::HandlerResponse::Handled;
+    }
+    return Server::HandlerResponse::Unhandled;
+  });
+  svr.set_pre_request_handler([](const Request &req, Response &res) {
+    if (req.matched_route == "/pre-request") {
+      res.set_content("pre-request", "text/plain");
+      return Server::HandlerResponse::Handled;
+    }
+    return Server::HandlerResponse::Unhandled;
+  });
+
+  svr.Get("/smuggled", [](const Request &, Response &res) {
+    res.set_content("smuggled", "text/plain");
+  });
+  svr.WebSocket("/pre-request",
+                [](const Request &, ws::WebSocket &ws) { ws.send("ws"); });
+
+  auto port = svr.bind_to_any_port(HOST);
+  std::thread t([&]() { svr.listen_after_bind(); });
+  auto se = detail::scope_exit([&] {
+    svr.stop();
+    t.join();
+  });
+  svr.wait_until_ready();
+
+  const std::string body = "GET /smuggled HTTP/1.1\r\nHost: localhost\r\n\r\n";
+
+  for (auto path : {"/pre-routing", "/pre-request"}) {
+    std::string req = std::string("GET ") + path +
+                      " HTTP/1.1\r\n"
+                      "Host: localhost\r\n"
+                      "Upgrade: websocket\r\n"
+                      "Connection: Upgrade\r\n"
+                      "Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n"
+                      "Sec-WebSocket-Version: 13\r\n"
+                      "Content-Length: " +
+                      std::to_string(body.size()) + "\r\n\r\n" + body;
+
+    std::string res;
+    ASSERT_TRUE(send_request(5, req, &res, port));
+    EXPECT_EQ(0u, res.find("HTTP/1.1 200 OK\r\n")) << path;
+    EXPECT_NE(std::string::npos, res.find("Connection: close\r\n")) << path;
+    // Only one response, and the body was not served as a request
+    EXPECT_EQ(std::string::npos, res.find("HTTP/1.1", 1)) << path;
+  }
+}
+
 TEST(WebSocketPreRequestTest, RejectWithoutAuth) {
   Server svr;