소스 검색

Buffer the WebSocket handshake before writing it to the socket

Follow-up to #2514. The rebuilt handshake wrote the request line straight
to the socket, so a header rejected by check_and_write_headers left a
truncated "GET /ws HTTP/1.1" sitting in the peer's buffer before the
connection was torn down, and every header cost its own small write.

Build the request into a BufferStream and flush it in one go, matching
ClientImpl::write_request. The new test drives a raw listener and asserts
the peer sees a clean EOF with zero bytes; without this change it observes
18.

Also fold WebSocketTest.HostHeaderInHandshake into
WebSocketTest.DefaultHeadersInHandshake, which covers the same Host
assertion through the capture helper #2514 introduced.
yhirose 1 주 전
부모
커밋
447b9c4a29
2개의 변경된 파일71개의 추가작업 그리고 43개의 파일을 삭제
  1. 11 2
      httplib.h
  2. 60 41
      test/test.cc

+ 11 - 2
httplib.h

@@ -9074,13 +9074,22 @@ inline bool perform_websocket_handshake(Stream &strm, Request &req,
   req.headers.emplace("Sec-WebSocket-Key", client_key);
   req.headers.emplace("Sec-WebSocket-Version", "13");
 
-  if (write_request_line(strm, req.method, req.path) < 0) { return false; }
+  // Build the request in memory first, like ClientImpl::write_request does.
+  // Writing straight to the socket would leak a request line onto the wire
+  // before check_and_write_headers gets a chance to reject an invalid header,
+  // and would emit one small write per header.
+  BufferStream bstrm;
+
+  if (write_request_line(bstrm, req.method, req.path) < 0) { return false; }
 
   auto error = Error::Success;
-  if (!check_and_write_headers(strm, req.headers, write_headers, error)) {
+  if (!check_and_write_headers(bstrm, req.headers, write_headers, error)) {
     return false;
   }
 
+  const auto &data = bstrm.get_buffer();
+  if (!write_data(strm, data.data(), data.size())) { return false; }
+
   // Verify 101 response and Sec-WebSocket-Accept header
   auto expected_accept = websocket_accept_key(client_key);
   return read_websocket_upgrade_response(strm, expected_accept,

+ 60 - 41
test/test.cc

@@ -19777,47 +19777,6 @@ TEST(WebSocketTest, QueryStringInHandshake) {
   t.join();
 }
 
-TEST(WebSocketTest, HostHeaderInHandshake) {
-  Server svr;
-
-  std::mutex mtx;
-  std::string received_host;
-
-  svr.WebSocket("/ws", [&](const Request &req, ws::WebSocket &ws) {
-    {
-      std::lock_guard<std::mutex> lock(mtx);
-      received_host = req.get_header_value("Host");
-    }
-    std::string msg;
-    while (ws.read(msg)) {
-      ws.send(msg);
-    }
-  });
-
-  auto port = svr.bind_to_any_port("localhost");
-  std::thread t([&]() { svr.listen_after_bind(); });
-  svr.wait_until_ready();
-
-  ws::WebSocketClient client("ws://localhost:" + std::to_string(port) + "/ws");
-  ASSERT_TRUE(client.connect());
-  // Round-trip ensures the handler has run and captured the request.
-  ASSERT_TRUE(client.send("hello"));
-  std::string msg;
-  ASSERT_TRUE(client.read(msg));
-  client.close();
-
-  {
-    std::lock_guard<std::mutex> lock(mtx);
-    // Non-default port must be present in the Host header. Default ports
-    // (80/443) are omitted; that logic is covered by
-    // MakeHostAndPortStringTest.
-    EXPECT_EQ("localhost:" + std::to_string(port), received_host);
-  }
-
-  svr.stop();
-  t.join();
-}
-
 // Run a handshake against a throwaway server and hand the request the server
 // received back to the caller, so tests can assert on the headers the client
 // actually put on the wire.
@@ -19864,6 +19823,9 @@ static void capture_websocket_handshake_request(
 
 TEST(WebSocketTest, DefaultHeadersInHandshake) {
   capture_websocket_handshake_request({}, [](const Request &req, int port) {
+    // Non-default port must be present in the Host header. Default ports
+    // (80/443) are omitted; that logic is covered by
+    // MakeHostAndPortStringTest.
     EXPECT_EQ("localhost:" + std::to_string(port),
               req.get_header_value("Host"));
     EXPECT_EQ(std::string("cpp-httplib/") + CPPHTTPLIB_VERSION,
@@ -19907,6 +19869,63 @@ TEST(WebSocketTest, MandatoryHeadersInHandshakeAreEnforced) {
       });
 }
 
+TEST(WebSocketTest, InvalidHeaderInHandshakeWritesNothing) {
+  // A header the client refuses to send must abort the handshake before any
+  // part of it reaches the wire; a lone request line would otherwise sit in
+  // the peer's buffer as a truncated request.
+  auto srv = ::socket(AF_INET, SOCK_STREAM, 0);
+  ASSERT_NE(INVALID_SOCKET, srv);
+  auto se_srv = detail::scope_exit([&] { detail::close_socket(srv); });
+
+  sockaddr_in addr{};
+  addr.sin_family = AF_INET;
+  addr.sin_port = 0; // ephemeral, so parallel shards don't collide
+  ::inet_pton(AF_INET, "127.0.0.1", &addr.sin_addr);
+  ASSERT_EQ(0, ::bind(srv, reinterpret_cast<sockaddr *>(&addr), sizeof(addr)));
+  ASSERT_EQ(0, ::listen(srv, 1));
+
+  sockaddr_in bound{};
+  socklen_t bound_len = sizeof(bound);
+  ASSERT_EQ(
+      0, ::getsockname(srv, reinterpret_cast<sockaddr *>(&bound), &bound_len));
+  auto port = ntohs(bound.sin_port);
+
+  ssize_t received = -1;
+  std::thread t([&] {
+    // Bound every blocking call so a regression fails the test with a bad
+    // value instead of hanging the suite.
+    fd_set rfds;
+    FD_ZERO(&rfds);
+    FD_SET(srv, &rfds);
+    timeval tv{5, 0};
+    if (::select(static_cast<int>(srv + 1), &rfds, nullptr, nullptr, &tv) <=
+        0) {
+      return;
+    }
+
+    sockaddr_in cli_addr{};
+    socklen_t cli_len = sizeof(cli_addr);
+    auto cli = ::accept(srv, reinterpret_cast<sockaddr *>(&cli_addr), &cli_len);
+    if (cli == INVALID_SOCKET) { return; }
+    auto se_cli = detail::scope_exit([&] { detail::close_socket(cli); });
+
+    detail::set_socket_opt_time(cli, SOL_SOCKET, SO_RCVTIMEO, 5, 0);
+    char buf[4096];
+    received = ::recv(cli, buf, sizeof(buf), 0);
+  });
+  // The CR/LF makes the value invalid, so check_and_write_headers rejects it.
+  // connect() has already shut the socket down by the time it returns false,
+  // so the peer sees EOF without waiting for the client to be destroyed.
+  ws::WebSocketClient client("ws://127.0.0.1:" + std::to_string(port) + "/ws",
+                             {{"X-Bad", "a\r\nInjected: 1"}});
+  EXPECT_FALSE(client.connect());
+
+  t.join();
+
+  // 0 means the peer saw a clean EOF without a single byte of the handshake.
+  EXPECT_EQ(0, received);
+}
+
 TEST(WebSocketTest, HostHeaderOverUnixSocket) {
   // The socket path doubles as the URL host, so it must not contain '/'.
   const char *shard = getenv("GTEST_SHARD_INDEX");