Parcourir la source

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.
yhirose il y a 2 jours
Parent
commit
1d5f0323b1
2 fichiers modifiés avec 90 ajouts et 7 suppressions
  1. 11 7
      httplib.h
  2. 79 0
      test/test.cc

+ 11 - 7
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();

+ 79 - 0
test/test.cc

@@ -27435,6 +27435,85 @@ TEST(WebSocketPreRoutingTest, RejectWithoutAuth) {
   t.join();
 }
 
+TEST(WebSocketPreRoutingTest, RunsOnceForUnmatchedUpgrade) {
+  Server svr;
+
+  std::atomic<int> pre_routing_calls{0};
+  std::atomic<int> 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;