Validate a WebSocket Close before echoing it

Any Close payload was echoed back, a 1-byte one, a code that may not be
sent and a reason that is not UTF-8 included. Echo a valid one and
answer the others with 1002, from one place for both read loops.

test/test.cc

@@ -26779,6 +26779,53 @@ TEST(WebSocketTest, CloseDoesNotStealBytesFromConcurrentRead) {
   }
 }
 
+// Connects to 127.0.0.1:port and completes a WebSocket handshake for path,
+// returning the socket, or INVALID_SOCKET.
+static socket_t connect_raw_websocket(int port, const std::string &path) {
+  auto sock = ::socket(AF_INET, SOCK_STREAM, 0);
+  if (sock == INVALID_SOCKET) { return sock; }
+  detail::set_socket_opt_time(sock, SOL_SOCKET, SO_RCVTIMEO, 5, 0);
+
+  sockaddr_in addr{};
+  addr.sin_family = AF_INET;
+  addr.sin_port = htons(static_cast<uint16_t>(port));
+  ::inet_pton(AF_INET, "127.0.0.1", &addr.sin_addr);
+  std::string req = "GET " + path +
+                    " HTTP/1.1\r\n"
+                    "Host: 127.0.0.1\r\n"
+                    "Upgrade: websocket\r\n"
+                    "Connection: Upgrade\r\n"
+                    "Sec-WebSocket-Key: AAAAAAAAAAAAAAAAAAAAAA==\r\n"
+                    "Sec-WebSocket-Version: 13\r\n"
+                    "\r\n";
+  std::string res;
+  if (::connect(sock, reinterpret_cast<sockaddr *>(&addr), sizeof(addr)) == 0 &&
+      ::send(sock, req.data(), static_cast<int>(req.size()), 0) ==
+          static_cast<int>(req.size())) {
+    char c;
+    while (res.find("\r\n\r\n") == std::string::npos &&
+           ::recv(sock, &c, 1, 0) == 1) {
+      res += c;
+    }
+  }
+  if (res.find(" 101 ") == std::string::npos) {
+    detail::close_socket(sock);
+    return INVALID_SOCKET;
+  }
+  return sock;
+}
+
+// A client frame with an all-zero mask key, so the payload goes out verbatim.
+// The payload must fit the 7-bit length field.
+static std::string masked_frame(uint8_t first_byte,
+                                const std::string &payload) {
+  std::string f;
+  f += static_cast<char>(first_byte);
+  f += static_cast<char>(0x80 | payload.size());
+  f.append(4, '\0');
+  return f + payload;
+}
+
 TEST(WebSocketTest, AssemblesFragmentsAroundAPing) {
   std::atomic<bool> done{false};
   std::string received;
@@ -26799,41 +26846,14 @@ TEST(WebSocketTest, AssemblesFragmentsAroundAPing) {
   });
   svr.wait_until_ready();
 
-  auto sock = ::socket(AF_INET, SOCK_STREAM, 0);
+  auto sock = connect_raw_websocket(port, "/ws");
   ASSERT_NE(INVALID_SOCKET, sock);
   auto se_sock = detail::scope_exit([&] { detail::close_socket(sock); });
-  detail::set_socket_opt_time(sock, SOL_SOCKET, SO_RCVTIMEO, 5, 0);
 
-  sockaddr_in addr{};
-  addr.sin_family = AF_INET;
-  addr.sin_port = htons(static_cast<uint16_t>(port));
-  ::inet_pton(AF_INET, "127.0.0.1", &addr.sin_addr);
-  ASSERT_EQ(0,
-            ::connect(sock, reinterpret_cast<sockaddr *>(&addr), sizeof(addr)));
-
-  // Frames with an all-zero mask key, so the payload goes out verbatim.
-  auto frame = [](uint8_t first_byte, const std::string &payload) {
-    std::string f;
-    f += static_cast<char>(first_byte);
-    f += static_cast<char>(0x80 | payload.size());
-    f.append(4, '\0');
-    return f + payload;
-  };
-  std::string data = "GET /ws HTTP/1.1\r\n"
-                     "Host: 127.0.0.1\r\n"
-                     "Upgrade: websocket\r\n"
-                     "Connection: Upgrade\r\n"
-                     "Sec-WebSocket-Key: AAAAAAAAAAAAAAAAAAAAAA==\r\n"
-                     "Sec-WebSocket-Version: 13\r\n"
-                     "\r\n";
-  data += frame(0x02, "Hel"); // Binary, more to come
-  data += frame(0x89, "p");   // FIN + Ping
-  data += frame(0x80, "lo");  // FIN + Continuation
-#ifdef _WIN32
-  ::send(sock, data.data(), static_cast<int>(data.size()), 0);
-#else
-  ::send(sock, data.data(), data.size(), 0);
-#endif
+  auto frames = masked_frame(0x02, "Hel") + // Binary, more to come
+                masked_frame(0x89, "p") +   // FIN + Ping
+                masked_frame(0x80, "lo");   // FIN + Continuation
+  ::send(sock, frames.data(), static_cast<int>(frames.size()), 0);
 
   for (int i = 0; i < 500 && !done; i++) {
     std::this_thread::sleep_for(std::chrono::milliseconds(10));
@@ -26842,6 +26862,50 @@ TEST(WebSocketTest, AssemblesFragmentsAroundAPing) {
   EXPECT_EQ("Hello", received);
 }
 
+TEST(WebSocketTest, AnswersAMalformedCloseWithProtocolError) {
+  Server svr;
+  svr.set_websocket_ping_interval(0);
+  svr.WebSocket("/ws", [&](const Request &, ws::WebSocket &ws) {
+    std::string msg;
+    while (ws.read(msg)) {}
+  });
+
+  auto port = svr.bind_to_any_port("127.0.0.1");
+  std::thread t([&]() { svr.listen_after_bind(); });
+  auto se = detail::scope_exit([&] {
+    svr.stop();
+    t.join();
+  });
+  svr.wait_until_ready();
+
+  // The Close each payload is answered with: an echo when it is valid, and
+  // 1002 otherwise.
+  std::vector<std::pair<std::string, std::string>> cases = {
+      {std::string("\x03\xe8") + "bye", std::string("\x03\xe8") + "bye"},
+      {std::string("\x03"), std::string("\x03\xea")},         // 1 byte
+      {std::string("\x03\xed"), std::string("\x03\xea")},     // 1005
+      {std::string("\x03\xe8\xff"), std::string("\x03\xea")}, // bad UTF-8
+  };
+  for (const auto &c : cases) {
+    auto sock = connect_raw_websocket(port, "/ws");
+    ASSERT_NE(INVALID_SOCKET, sock);
+    auto close_frame = masked_frame(0x88, c.first);
+    ::send(sock, close_frame.data(), static_cast<int>(close_frame.size()), 0);
+
+    std::string reply;
+    char buf[256];
+    while (reply.size() < 2 + c.second.size()) {
+      auto n = ::recv(sock, buf, sizeof(buf), 0);
+      if (n <= 0) { break; }
+      reply.append(buf, static_cast<size_t>(n));
+    }
+    detail::close_socket(sock);
+    EXPECT_EQ(std::string("\x88") + static_cast<char>(c.second.size()) +
+                  c.second,
+              reply);
+  }
+}
+
 #ifdef CPPHTTPLIB_OPENSSL_SUPPORT
 class WebSocketSSLIntegrationTest : public ::testing::Test {
 protected: