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.

httplib.h

@@ -5832,6 +5832,19 @@ inline FrameRead read_websocket_frame(Stream &strm, Opcode &opcode,
   return FrameRead::Ok;
 }
 
+// RFC 6455 §5.5.1, §7.4: a Close payload is empty, or a status code that may
+// be sent followed by a UTF-8 reason.
+inline bool is_valid_close_payload(const std::string &payload) {
+  if (payload.empty()) { return true; }
+  if (payload.size() < 2) { return false; }
+  auto code = (static_cast<uint8_t>(payload[0]) << 8) |
+              static_cast<uint8_t>(payload[1]);
+  auto may_be_sent = (1000 <= code && code <= 1003) ||
+                     (1007 <= code && code <= 1014) ||
+                     (3000 <= code && code <= 4999);
+  return may_be_sent && is_valid_utf8(payload.substr(2));
+}
+
 } // namespace impl
 } // namespace ws
 
@@ -22907,6 +22920,20 @@ inline bool WebSocket::send_frame(Opcode op, const char *data, size_t len,
 }
 
 inline ReadResult WebSocket::read(std::string &msg) {
+  // RFC 6455 §5.5.1: a Close is echoed back, unless its payload is malformed,
+  // which is answered with a protocol error instead.
+  auto answer_close = [&](const std::string &payload) {
+    if (closed_.exchange(true)) { return; }
+    auto reply = payload;
+    if (!impl::is_valid_close_payload(payload)) {
+      auto code = static_cast<uint16_t>(CloseStatus::ProtocolError);
+      reply = {static_cast<char>(code >> 8), static_cast<char>(code & 0xff)};
+    }
+    std::lock_guard<std::mutex> lock(write_mutex_);
+    detail::write_websocket_frame(strm_, Opcode::Close, reply.data(),
+                                  reply.size(), true, !is_server_);
+  };
+
   std::unique_lock<std::mutex> read_lock(read_mutex_);
   while (!closed_) {
     Opcode opcode;
@@ -22940,12 +22967,7 @@ inline ReadResult WebSocket::read(std::string &msg) {
       continue;
     }
     case Opcode::Close: {
-      if (!closed_.exchange(true)) {
-        // Echo close frame back
-        std::lock_guard<std::mutex> lock(write_mutex_);
-        detail::write_websocket_frame(strm_, Opcode::Close, payload.data(),
-                                      payload.size(), true, !is_server_);
-      }
+      answer_close(payload);
       return Fail;
     }
     case Opcode::Text:
@@ -22990,12 +23012,7 @@ inline ReadResult WebSocket::read(std::string &msg) {
             continue;
           }
           if (cont_opcode == Opcode::Close) {
-            if (!closed_.exchange(true)) {
-              std::lock_guard<std::mutex> lock(write_mutex_);
-              detail::write_websocket_frame(
-                  strm_, Opcode::Close, cont_payload.data(),
-                  cont_payload.size(), true, !is_server_);
-            }
+            answer_close(cont_payload);
             return Fail;
           }
           // RFC 6455: continuation frames must use opcode 0x0

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: