From aa0a588ed53d4a33705fff803e3486bb5349dd31 Mon Sep 17 00:00:00 2001 From: yhirose Date: Fri, 16 Oct 2026 21:38:40 -0400 Subject: [PATCH] Answer WebSocket control frames from one place The read loop and the loop over the fragments of a message each answered Ping, Pong and Close the same way. Both go through one handler now. --- httplib.h | 81 +++++++++++++++++++++++----------------------------- test/test.cc | 38 ++++++++++++++---------- 2 files changed, 59 insertions(+), 60 deletions(-) diff --git a/httplib.h b/httplib.h index 0768d87c..8de6f87d 100644 --- a/httplib.h +++ b/httplib.h @@ -23046,18 +23046,37 @@ 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(CloseStatus::ProtocolError); - reply = {static_cast(code >> 8), static_cast(code & 0xff)}; + // Answers a control frame; any other opcode here is a protocol error. + // Returns false when the connection is done. + auto handle_control = [&](Opcode op, const std::string &payload) -> bool { + switch (op) { + case Opcode::Ping: { + std::lock_guard lock(write_mutex_); + detail::write_websocket_frame(strm_, Opcode::Pong, payload.data(), + payload.size(), true, !is_server_); + return true; + } + case Opcode::Pong: { + std::lock_guard lock(ping_mutex_); + unacked_pings_ = 0; + return true; + } + case Opcode::Close: { + // RFC 6455 §5.5.1: a Close is echoed back, unless its payload is + // malformed, which is answered with a protocol error instead. + if (closed_.exchange(true)) { return false; } + auto reply = payload; + if (!impl::is_valid_close_payload(payload)) { + auto code = static_cast(CloseStatus::ProtocolError); + reply = {static_cast(code >> 8), static_cast(code & 0xff)}; + } + std::lock_guard lock(write_mutex_); + detail::write_websocket_frame(strm_, Opcode::Close, reply.data(), + reply.size(), true, !is_server_); + return false; + } + default: closed_ = true; return false; } - std::lock_guard lock(write_mutex_); - detail::write_websocket_frame(strm_, Opcode::Close, reply.data(), - reply.size(), true, !is_server_); }; std::unique_lock read_lock(read_mutex_); @@ -23081,21 +23100,6 @@ inline ReadResult WebSocket::read(std::string &msg) { } switch (opcode) { - case Opcode::Ping: { - std::lock_guard lock(write_mutex_); - detail::write_websocket_frame(strm_, Opcode::Pong, payload.data(), - payload.size(), true, !is_server_); - continue; - } - case Opcode::Pong: { - std::lock_guard lock(ping_mutex_); - unacked_pings_ = 0; - continue; - } - case Opcode::Close: { - answer_close(payload); - return Fail; - } case Opcode::Text: case Opcode::Binary: { auto result = opcode == Opcode::Text ? Text : Binary; @@ -23125,25 +23129,10 @@ inline ReadResult WebSocket::read(std::string &msg) { closed_ = true; return Fail; } - if (cont_opcode == Opcode::Ping) { - std::lock_guard lock(write_mutex_); - detail::write_websocket_frame( - strm_, Opcode::Pong, cont_payload.data(), cont_payload.size(), - true, !is_server_); - continue; - } - if (cont_opcode == Opcode::Pong) { - std::lock_guard lock(ping_mutex_); - unacked_pings_ = 0; - continue; - } - if (cont_opcode == Opcode::Close) { - answer_close(cont_payload); - return Fail; - } - // RFC 6455: continuation frames must use opcode 0x0 + // RFC 6455 5.4: the fragments are continuation frames, with only + // control frames between them. if (cont_opcode != Opcode::Continuation) { - closed_ = true; + if (handle_control(cont_opcode, cont_payload)) { continue; } return Fail; } msg += cont_payload; @@ -23165,7 +23154,9 @@ inline ReadResult WebSocket::read(std::string &msg) { } return result; } - default: closed_ = true; return Fail; + default: + if (handle_control(opcode, payload)) { continue; } + return Fail; } } return Fail; diff --git a/test/test.cc b/test/test.cc index 5d1ffb81..815c9674 100644 --- a/test/test.cc +++ b/test/test.cc @@ -27294,6 +27294,11 @@ TEST(WebSocketTest, AssemblesFragmentsAroundAPing) { } ASSERT_TRUE(done); EXPECT_EQ("Hello", received); + + // The Ping between the fragments is answered with a Pong. + char pong[3]; + ASSERT_EQ(3, ::recv(sock, pong, sizeof(pong), MSG_WAITALL)); + EXPECT_EQ(std::string("\x8a\x01p", 3), std::string(pong, 3)); } TEST(WebSocketTest, AnswersAMalformedCloseWithProtocolError) { @@ -27320,23 +27325,26 @@ TEST(WebSocketTest, AnswersAMalformedCloseWithProtocolError) { {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(close_frame.size()), 0); + // Also between the fragments of a message. + for (const auto &prefix : {std::string(), masked_frame(0x02, "a")}) { + for (const auto &c : cases) { + auto sock = connect_raw_websocket(port, "/ws"); + ASSERT_NE(INVALID_SOCKET, sock); + auto frames = prefix + masked_frame(0x88, c.first); + ::send(sock, frames.data(), static_cast(frames.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(n)); + 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(n)); + } + detail::close_socket(sock); + EXPECT_EQ(std::string("\x88") + static_cast(c.second.size()) + + c.second, + reply); } - detail::close_socket(sock); - EXPECT_EQ(std::string("\x88") + static_cast(c.second.size()) + - c.second, - reply); } }