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

@@ -23046,10 +23046,25 @@ 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,
+  // Answers a control frame; any other opcode here is a protocol error.
-  // which is answered with a protocol error instead.
+  // Returns false when the connection is done.
-  auto answer_close = [&](const std::string &payload) {
+  auto handle_control = [&](Opcode op, const std::string &payload) -> bool {
-    if (closed_.exchange(true)) { return; }
+    switch (op) {
+    case Opcode::Ping: {
+      std::lock_guard<std::mutex> 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<std::mutex> 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<uint16_t>(CloseStatus::ProtocolError);
@@ -23058,6 +23073,10 @@ inline ReadResult WebSocket::read(std::string &msg) {
       std::lock_guard<std::mutex> 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::unique_lock<std::mutex> read_lock(read_mutex_);
@@ -23081,21 +23100,6 @@ inline ReadResult WebSocket::read(std::string &msg) {
     }
 
     switch (opcode) {
-    case Opcode::Ping: {
-      std::lock_guard<std::mutex> lock(write_mutex_);
-      detail::write_websocket_frame(strm_, Opcode::Pong, payload.data(),
-                                    payload.size(), true, !is_server_);
-      continue;
-    }
-    case Opcode::Pong: {
-      std::lock_guard<std::mutex> 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) {
+          // RFC 6455 5.4: the fragments are continuation frames, with only
-            std::lock_guard<std::mutex> lock(write_mutex_);
+          // control frames between them.
-            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<std::mutex> 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
           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;

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,11 +27325,13 @@ TEST(WebSocketTest, AnswersAMalformedCloseWithProtocolError) {
       {std::string("\x03\xed"), std::string("\x03\xea")},     // 1005
       {std::string("\x03\xe8\xff"), std::string("\x03\xea")}, // bad UTF-8
   };
+  // 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 close_frame = masked_frame(0x88, c.first);
+      auto frames = prefix + masked_frame(0x88, c.first);
-    ::send(sock, close_frame.data(), static_cast<int>(close_frame.size()), 0);
+      ::send(sock, frames.data(), static_cast<int>(frames.size()), 0);
 
       std::string reply;
       char buf[256];
@@ -27339,6 +27346,7 @@ TEST(WebSocketTest, AnswersAMalformedCloseWithProtocolError) {
                 reply);
     }
   }
+}
 
 #ifdef CPPHTTPLIB_OPENSSL_SUPPORT
 class WebSocketSSLIntegrationTest : public ::testing::Test {