From 1d5f0323b12bebb6399325e4053da8edab54cf9d Mon Sep 17 00:00:00 2001 From: yhirose Date: Sun, 18 Oct 2026 12:41:23 -0400 Subject: [PATCH] Run the pre-routing handler once for an unmatched WebSocket upgrade An upgrade request goes through the pre-routing handler before the WebSocket routes are matched. When none matched, it was routed as a plain HTTP request, and routing() ran the handler a second time. Skip it there for an upgrade request. --- httplib.h | 18 +++++++----- test/test.cc | 79 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 90 insertions(+), 7 deletions(-) diff --git a/httplib.h b/httplib.h index 7bcad4a4..eff46786 100644 --- a/httplib.h +++ b/httplib.h @@ -2473,7 +2473,8 @@ private: int bind_internal(const std::string &host, int port, int socket_flags); bool listen_internal(); - bool routing(Request &req, Response &res, Stream &strm); + bool routing(Request &req, Response &res, Stream &strm, + bool skip_pre_routing); bool handle_file_request(Request &req, Response &res); bool check_if_not_modified(const Request &req, Response &res, const std::string &etag, time_t mtime) const; @@ -14612,8 +14613,9 @@ inline bool Server::listen_internal() { return ret; } -inline bool Server::routing(Request &req, Response &res, Stream &strm) { - if (pre_routing_handler_ && +inline bool Server::routing(Request &req, Response &res, Stream &strm, + bool skip_pre_routing) { + if (!skip_pre_routing && pre_routing_handler_ && pre_routing_handler_(req, res) == HandlerResponse::Handled) { return true; } @@ -15171,7 +15173,9 @@ Server::process_request(Stream &strm, const std::string &remote_addr, // The response says so when this connection closes after it. auto close_after_response = close_connection || connection_closed; - if (detail::is_websocket_upgrade(req)) { + // An upgrade request has been through pre-routing before it reaches routing() + auto is_upgrade = detail::is_websocket_upgrade(req); + if (is_upgrade) { if (pre_routing_handler_ && pre_routing_handler_(req, res) == HandlerResponse::Handled) { if (res.status == -1) { res.status = StatusCode::OK_200; } @@ -15265,16 +15269,16 @@ Server::process_request(Stream &strm, const std::string &remote_addr, return true; } } - // No matching handler - fall through to 404 + // No WebSocket route: route it as a plain HTTP request } // Routing auto routed = false; #ifdef CPPHTTPLIB_NO_EXCEPTIONS - routed = routing(req, res, strm); + routed = routing(req, res, strm, is_upgrade); #else try { - routed = routing(req, res, strm); + routed = routing(req, res, strm, is_upgrade); } catch (std::exception &) { if (exception_handler_) { auto ep = std::current_exception(); diff --git a/test/test.cc b/test/test.cc index c75aa984..7bedb460 100644 --- a/test/test.cc +++ b/test/test.cc @@ -27435,6 +27435,85 @@ TEST(WebSocketPreRoutingTest, RejectWithoutAuth) { t.join(); } +TEST(WebSocketPreRoutingTest, RunsOnceForUnmatchedUpgrade) { + Server svr; + + std::atomic pre_routing_calls{0}; + std::atomic pre_request_calls{0}; + svr.set_pre_routing_handler([&](const Request &req, Response &res) { + pre_routing_calls++; + if (req.path == "/handled") { + res.set_content("handled", "text/plain"); + return Server::HandlerResponse::Handled; + } + return Server::HandlerResponse::Unhandled; + }); + svr.set_pre_request_handler([&](const Request &, Response &) { + pre_request_calls++; + return Server::HandlerResponse::Unhandled; + }); + + svr.Get("/http", [](const Request &, Response &res) { + res.set_content("http", "text/plain"); + }); + svr.WebSocket("/ws", + [](const Request &, ws::WebSocket &ws) { ws.send("ws"); }); + + auto port = svr.bind_to_any_port("localhost"); + std::thread t([&]() { svr.listen_after_bind(); }); + auto se = detail::scope_exit([&] { + svr.stop(); + t.join(); + }); + svr.wait_until_ready(); + + Headers upgrade = {{"Upgrade", "websocket"}, + {"Connection", "Upgrade"}, + {"Sec-WebSocket-Key", "dGhlIHNhbXBsZSBub25jZQ=="}, + {"Sec-WebSocket-Version", "13"}}; + + auto get = [&](const char *path) { + pre_routing_calls = 0; + return Client("localhost", port).Get(path, upgrade); + }; + + // No WebSocket route, but an HTTP one + auto res = get("/http"); + ASSERT_TRUE(res); + EXPECT_EQ(StatusCode::OK_200, res->status); + EXPECT_EQ("http", res->body); + EXPECT_EQ(1, pre_routing_calls); + EXPECT_EQ(1, pre_request_calls); + + // No route at all + res = get("/none"); + ASSERT_TRUE(res); + EXPECT_EQ(StatusCode::NotFound_404, res->status); + EXPECT_EQ(1, pre_routing_calls); + + // Answered by the pre-routing handler + res = get("/handled"); + ASSERT_TRUE(res); + EXPECT_EQ(StatusCode::OK_200, res->status); + EXPECT_EQ("handled", res->body); + EXPECT_EQ(1, pre_routing_calls); + + // A WebSocket route + pre_routing_calls = 0; + pre_request_calls = 0; + { + ws::WebSocketClient client("ws://localhost:" + std::to_string(port) + + "/ws"); + ASSERT_TRUE(client.connect()); + std::string msg; + ASSERT_TRUE(client.read(msg)); + EXPECT_EQ("ws", msg); + client.close(); + } + EXPECT_EQ(1, pre_routing_calls); + EXPECT_EQ(1, pre_request_calls); +} + TEST(WebSocketPreRequestTest, RejectWithoutAuth) { Server svr;