diff --git a/Makefile b/Makefile index 2b7f50464..34817430b 100644 --- a/Makefile +++ b/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 diff --git a/http/WebSocketParser.cpp b/http/WebSocketParser.cpp index 5df60f95d..bdb3655c9 100644 --- a/http/WebSocketParser.cpp +++ b/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,10 +45,14 @@ 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; } @@ -40,6 +60,12 @@ static int on_frame_body(websocket_parser* parser, const char * at, size_t lengt 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; diff --git a/http/WebSocketParser.h b/http/WebSocketParser.h index c551550fd..f1bd6d3e1 100644 --- a/http/WebSocketParser.h +++ b/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 onMessage; WebSocketParser(); diff --git a/scripts/unittest.sh b/scripts/unittest.sh index 5a4d05668..ca1046e8d 100755 --- a/scripts/unittest.sh +++ b/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 diff --git a/unittest/CMakeLists.txt b/unittest/CMakeLists.txt index d0a1327fc..d3a67a4ee 100644 --- a/unittest/CMakeLists.txt +++ b/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 diff --git a/unittest/websocket_parser_test.cpp b/unittest/websocket_parser_test.cpp new file mode 100644 index 000000000..5cdfe4a38 --- /dev/null +++ b/unittest/websocket_parser_test.cpp @@ -0,0 +1,106 @@ +#include + +#include +#include +#include + +#include "WebSocketParser.h" +#include "wsdef.h" + +typedef std::pair Message; +typedef std::vector 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; +}