fix(http): keep WebSocket control frames out of fragmented messages

Makefile

@@ -414,6 +414,9 @@ unittest: prepare libhv
 ifeq ($(WITH_EVPP), yes)
 	$(CXX) -g -Wall -O0 -std=c++11 -I. -Ibase -Issl -Ievent -Icpputil -Ievpp -o bin/tcpclient_dns_test unittest/tcpclient_dns_test.cpp -Llib -lhv -pthread
 endif
+ifeq ($(WITH_HTTP), yes)
+	$(CXX) -g -Wall -O0 -std=c++11 -I. -Ibase -Ihttp -o bin/websocket_parser_test unittest/websocket_parser_test.cpp -Llib -lhv -pthread
+endif
 ifeq ($(WITH_EVPP), yes)
 ifeq ($(WITH_REDIS), yes)
 	$(CXX) -g -Wall -O0 -std=c++11 -I. -Ibase -Ievent -Icpputil -Iredis -o bin/redis_protocol_test unittest/redis_protocol_test.cpp redis/RedisMessage.cpp

http/WebSocketParser.cpp

@@ -4,11 +4,27 @@
 #include "hdef.h"
 
 #define MAX_PAYLOAD_LENGTH  (1 << 24)   // 16M
+#define MAX_CONTROL_PAYLOAD_LENGTH  125
+
+// Control frames may be injected between the fragments of a data message
+// (RFC 6455 5.4), so they must not touch the data message state.
+static bool is_control_frame(websocket_parser* parser) {
+    return (parser->flags & WS_OP_MASK) >= WS_OP_CLOSE;
+}
 
 static int on_frame_header(websocket_parser* parser) {
     WebSocketParser* wp = (WebSocketParser*)parser->data;
     int opcode = parser->flags & WS_OP_MASK;
     // printf("on_frame_header opcode=%d\n", opcode);
+    if (is_control_frame(parser)) {
+        // RFC 6455 5.5: control frames MUST NOT be fragmented and MUST have
+        // a payload length of 125 bytes or less.
+        if (!(parser->flags & WS_FIN) || parser->length > MAX_CONTROL_PAYLOAD_LENGTH) {
+            return -1;
+        }
+        wp->control_message.clear();
+        return 0;
+    }
     if (opcode != WS_OP_CONTINUE) {
         wp->opcode = opcode;
     }
@@ -29,17 +45,27 @@ static int on_frame_header(websocket_parser* parser) {
 static int on_frame_body(websocket_parser* parser, const char * at, size_t length) {
     // printf("on_frame_body length=%d\n", (int)length);
     WebSocketParser* wp = (WebSocketParser*)parser->data;
-    wp->state = WS_FRAME_BODY;
     if (wp->parser->flags & WS_HAS_MASK) {
         websocket_parser_decode((char*)at, at, length, wp->parser);
     }
+    if (is_control_frame(parser)) {
+        wp->control_message.append(at, length);
+        return 0;
+    }
+    wp->state = WS_FRAME_BODY;
     wp->message.append(at, length);
     return 0;
 }
 
 static int on_frame_end(websocket_parser* parser) {
     // printf("on_frame_end\n");
     WebSocketParser* wp = (WebSocketParser*)parser->data;
+    if (is_control_frame(parser)) {
+        if (wp->onMessage) {
+            wp->onMessage(parser->flags & WS_OP_MASK, wp->control_message);
+        }
+        return 0;
+    }
     wp->state = WS_FRAME_END;
     if (wp->parser->flags & WS_FIN) {
         wp->state = WS_FRAME_FIN;

http/WebSocketParser.h

@@ -22,6 +22,7 @@ class HV_EXPORT WebSocketParser {
     websocket_parser_state              state;
     int                                 opcode;
     std::string                         message;
+    std::string                         control_message;
     std::function<void(int opcode, const std::string& msg)> onMessage;
 
     WebSocketParser();

scripts/unittest.sh

@@ -32,6 +32,9 @@ bin/socketpair_test
 # bin/objectpool_test
 bin/sizeof_test
 bin/http_router_test
+if [ -x bin/websocket_parser_test ]; then
+    bin/websocket_parser_test || exit $?
+fi
 if [ -x bin/tlv_test ]; then
     bin/tlv_test
 fi

unittest/CMakeLists.txt

@@ -90,6 +90,13 @@ target_include_directories(sendmail PRIVATE .. ../base ../protocol ../util)
 add_executable(http_router_test http_router_test.cpp)
 target_include_directories(http_router_test PRIVATE ../http/server)
 
+if(WITH_HTTP)
+add_executable(websocket_parser_test websocket_parser_test.cpp)
+target_include_directories(websocket_parser_test PRIVATE .. ../base ../http)
+target_link_libraries(websocket_parser_test ${HV_LIBRARIES})
+set(HTTP_UNITTEST_TARGETS websocket_parser_test)
+endif()
+
 if(WITH_LUA)
 add_executable(lua_binding_test lua_binding_test.cpp)
 target_include_directories(lua_binding_test PRIVATE .. ../base ../ssl ../event ../cpputil ../evpp ../lua)
@@ -221,6 +228,7 @@ add_custom_target(unittest DEPENDS
     ftp
     sendmail
     http_router_test
+    ${HTTP_UNITTEST_TARGETS}
     ${HTTP_LUA_UNITTEST_TARGETS}
     ${HTTP_JS_UNITTEST_TARGETS}
     hdns_test

unittest/websocket_parser_test.cpp

@@ -0,0 +1,106 @@
+#include <assert.h>
+
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "WebSocketParser.h"
+#include "wsdef.h"
+
+typedef std::pair<int, std::string> Message;
+typedef std::vector<Message> Messages;
+
+static std::string frame(enum ws_opcode opcode, const std::string& payload, bool fin = true, bool has_mask = true) {
+    static const char mask[4] = {0x11, 0x22, 0x33, 0x44};
+    std::string buf(ws_calc_frame_size((int)payload.size(), has_mask), '\0');
+    ws_build_frame(&buf[0], payload.data(), (int)payload.size(), mask, has_mask, opcode, fin);
+    return buf;
+}
+
+static size_t feed(const std::string& data, Messages* messages, bool byte_by_byte = false) {
+    WebSocketParser parser;
+    parser.onMessage = [messages](int opcode, const std::string& msg) {
+        messages->push_back(Message(opcode, msg));
+    };
+    if (!byte_by_byte) {
+        return (size_t)parser.FeedRecvData(data.data(), data.size());
+    }
+    size_t nfeed = 0;
+    for (size_t i = 0; i < data.size(); ++i) {
+        if (parser.FeedRecvData(data.data() + i, 1) != 1) break;
+        ++nfeed;
+    }
+    return nfeed;
+}
+
+static void test_ping_between_fragments() {
+    std::string data = frame(WS_OPCODE_TEXT, "Hel", false) +
+                       frame(WS_OPCODE_PING, "x") +
+                       frame(WS_OPCODE_CONTINUE, "lo");
+    Messages messages;
+    assert(feed(data, &messages) == data.size());
+    assert(messages.size() == 2);
+    assert(messages[0] == Message(WS_OPCODE_PING, "x"));
+    assert(messages[1] == Message(WS_OPCODE_TEXT, "Hello"));
+}
+
+static void test_control_frames_between_fragments_byte_by_byte() {
+    std::string data = frame(WS_OPCODE_BINARY, "ab", false, false) +
+                       frame(WS_OPCODE_PONG, "", true, false) +
+                       frame(WS_OPCODE_CONTINUE, "cd", false, false) +
+                       frame(WS_OPCODE_PING, "ping", true, false) +
+                       frame(WS_OPCODE_CONTINUE, "ef", true, false);
+    Messages messages;
+    assert(feed(data, &messages, true) == data.size());
+    assert(messages.size() == 3);
+    assert(messages[0] == Message(WS_OPCODE_PONG, ""));
+    assert(messages[1] == Message(WS_OPCODE_PING, "ping"));
+    assert(messages[2] == Message(WS_OPCODE_BINARY, "abcdef"));
+}
+
+static void test_control_frame_between_messages() {
+    std::string data = frame(WS_OPCODE_TEXT, "a") +
+                       frame(WS_OPCODE_PING, "p") +
+                       frame(WS_OPCODE_TEXT, "b", false) +
+                       frame(WS_OPCODE_CONTINUE, "c");
+    Messages messages;
+    assert(feed(data, &messages) == data.size());
+    assert(messages.size() == 3);
+    assert(messages[0] == Message(WS_OPCODE_TEXT, "a"));
+    assert(messages[1] == Message(WS_OPCODE_PING, "p"));
+    assert(messages[2] == Message(WS_OPCODE_TEXT, "bc"));
+}
+
+static void test_max_control_payload() {
+    std::string payload(125, 'A');
+    std::string data = frame(WS_OPCODE_PING, payload);
+    Messages messages;
+    assert(feed(data, &messages) == data.size());
+    assert(messages.size() == 1);
+    assert(messages[0] == Message(WS_OPCODE_PING, payload));
+}
+
+static void test_reject_fragmented_control_frame() {
+    std::string data = frame(WS_OPCODE_PING, "pi", false) +
+                       frame(WS_OPCODE_CONTINUE, "ng");
+    Messages messages;
+    assert(feed(data, &messages) < data.size());
+    assert(messages.empty());
+}
+
+static void test_reject_oversized_control_frame() {
+    std::string data = frame(WS_OPCODE_PING, std::string(126, 'A'));
+    Messages messages;
+    assert(feed(data, &messages) < data.size());
+    assert(messages.empty());
+}
+
+int main() {
+    test_ping_between_fragments();
+    test_control_frames_between_fragments_byte_by_byte();
+    test_control_frame_between_messages();
+    test_max_control_payload();
+    test_reject_fragmented_control_frame();
+    test_reject_oversized_control_frame();
+    return 0;
+}