diff options
Diffstat (limited to 'modules/websocket/wsl_peer.cpp')
| -rw-r--r-- | modules/websocket/wsl_peer.cpp | 123 |
1 files changed, 104 insertions, 19 deletions
diff --git a/modules/websocket/wsl_peer.cpp b/modules/websocket/wsl_peer.cpp index 38de729eb6..f5cfad597e 100644 --- a/modules/websocket/wsl_peer.cpp +++ b/modules/websocket/wsl_peer.cpp @@ -297,6 +297,7 @@ Error WSLPeer::_do_server_handshake() { resolver.stop(); // Response sent, initialize wslay context. wslay_event_context_server_init(&wsl_ctx, &_wsl_callbacks, this); + wslay_event_config_set_no_buffering(wsl_ctx, 1); wslay_event_config_set_max_recv_msg_length(wsl_ctx, inbound_buffer_size); in_buffer.resize(nearest_shift(inbound_buffer_size), max_queued_packets); packet_buffer.resize(inbound_buffer_size); @@ -405,6 +406,7 @@ void WSLPeer::_do_client_handshake() { ERR_FAIL_MSG("Invalid response headers."); } wslay_event_context_client_init(&wsl_ctx, &_wsl_callbacks, this); + wslay_event_config_set_no_buffering(wsl_ctx, 1); wslay_event_config_set_max_recv_msg_length(wsl_ctx, inbound_buffer_size); in_buffer.resize(nearest_shift(inbound_buffer_size), max_queued_packets); packet_buffer.resize(inbound_buffer_size); @@ -570,8 +572,15 @@ ssize_t WSLPeer::_wsl_recv_callback(wslay_event_context_ptr ctx, uint8_t *data, wslay_event_set_error(ctx, WSLAY_ERR_CALLBACK_FAILURE); return -1; } + // Make sure we don't read more than what our buffer can hold. + size_t buffer_limit = MIN(peer->in_buffer.payload_space_left(), peer->in_buffer.packets_space_left() * 2); // The minimum size of a websocket message is 2 bytes. + size_t to_read = MIN(len, buffer_limit); + if (to_read == 0) { + wslay_event_set_error(ctx, WSLAY_ERR_WOULDBLOCK); + return -1; + } int read = 0; - Error err = conn->get_partial_data(data, len, read); + Error err = conn->get_partial_data(data, to_read, read); if (err != OK) { print_verbose("Websocket get data error: " + itos(err) + ", read (should be 0!): " + itos(read)); wslay_event_set_error(ctx, WSLAY_ERR_CALLBACK_FAILURE); @@ -584,6 +593,37 @@ ssize_t WSLPeer::_wsl_recv_callback(wslay_event_context_ptr ctx, uint8_t *data, return read; } +void WSLPeer::_wsl_recv_start_callback(wslay_event_context_ptr ctx, const struct wslay_event_on_frame_recv_start_arg *arg, void *user_data) { + WSLPeer *peer = (WSLPeer *)user_data; + uint8_t op = arg->opcode; + if (op == WSLAY_TEXT_FRAME || op == WSLAY_BINARY_FRAME) { + // Get ready to process a data package. + PendingMessage &pm = peer->pending_message; + pm.opcode = op; + pm.payload_size = arg->payload_length; + } +} + +void WSLPeer::_wsl_frame_recv_chunk_callback(wslay_event_context_ptr ctx, const struct wslay_event_on_frame_recv_chunk_arg *arg, void *user_data) { + WSLPeer *peer = (WSLPeer *)user_data; + PendingMessage &pm = peer->pending_message; + if (pm.opcode != 0) { + // Only write the payload. + peer->in_buffer.write_packet(arg->data, arg->data_length, nullptr); + } +} + +void WSLPeer::_wsl_frame_recv_end_callback(wslay_event_context_ptr ctx, void *user_data) { + WSLPeer *peer = (WSLPeer *)user_data; + PendingMessage &pm = peer->pending_message; + if (pm.opcode != 0) { + // Only write the packet (since it's now completed). + uint8_t is_string = pm.opcode == WSLAY_TEXT_FRAME ? 1 : 0; + peer->in_buffer.write_packet(nullptr, pm.payload_size, &is_string); + pm.clear(); + } +} + ssize_t WSLPeer::_wsl_send_callback(wslay_event_context_ptr ctx, const uint8_t *data, size_t len, int flags, void *user_data) { WSLPeer *peer = (WSLPeer *)user_data; Ref<StreamPeer> conn = peer->connection; @@ -629,25 +669,19 @@ void WSLPeer::_wsl_msg_recv_callback(wslay_event_context_ptr ctx, const struct w return; } - if (peer->ready_state == STATE_CLOSING) { - return; - } - - if (op == WSLAY_TEXT_FRAME || op == WSLAY_BINARY_FRAME) { - // Message. - uint8_t is_string = arg->opcode == WSLAY_TEXT_FRAME ? 1 : 0; - peer->in_buffer.write_packet(arg->msg, arg->msg_length, &is_string); + if (op == WSLAY_PONG) { + peer->heartbeat_waiting = false; } - // Ping or pong. + // Ping, or message (already parsed in chunks). } wslay_event_callbacks WSLPeer::_wsl_callbacks = { _wsl_recv_callback, _wsl_send_callback, _wsl_genmask_callback, - nullptr, /* on_frame_recv_start_callback */ - nullptr, /* on_frame_recv_callback */ - nullptr, /* on_frame_recv_end_callback */ + _wsl_recv_start_callback, + _wsl_frame_recv_chunk_callback, + _wsl_frame_recv_end_callback, _wsl_msg_recv_callback }; @@ -682,7 +716,31 @@ void WSLPeer::poll() { if (ready_state == STATE_OPEN || ready_state == STATE_CLOSING) { ERR_FAIL_NULL(wsl_ctx); + uint64_t ticks = OS::get_singleton()->get_ticks_msec(); int err = 0; + if (heartbeat_interval_msec != 0 && ticks - last_heartbeat > heartbeat_interval_msec && ready_state == STATE_OPEN) { + if (heartbeat_waiting) { + wslay_event_context_free(wsl_ctx); + wsl_ctx = nullptr; + close(-1); + return; + } + heartbeat_waiting = true; + struct wslay_event_msg msg; + msg.opcode = WSLAY_PING; + msg.msg = nullptr; + msg.msg_length = 0; + err = wslay_event_queue_msg(wsl_ctx, &msg); + if (err == 0) { + last_heartbeat = ticks; + } else { + print_verbose("Websocket (wslay) failed to send ping: " + itos(err)); + wslay_event_context_free(wsl_ctx); + wsl_ctx = nullptr; + close(-1); + return; + } + } if ((err = wslay_event_recv(wsl_ctx)) != 0 || (err = wslay_event_send(wsl_ctx)) != 0) { // Error close. print_verbose("Websocket (wslay) poll error: " + itos(err)); @@ -691,12 +749,37 @@ void WSLPeer::poll() { close(-1); return; } - if (wslay_event_get_close_sent(wsl_ctx) && wslay_event_get_close_received(wsl_ctx)) { - // Clean close. - wslay_event_context_free(wsl_ctx); - wsl_ctx = nullptr; - close(-1); - return; + if (wslay_event_get_close_sent(wsl_ctx)) { + if (wslay_event_get_close_received(wsl_ctx)) { + // Clean close. + wslay_event_context_free(wsl_ctx); + wsl_ctx = nullptr; + close(-1); + return; + } else if (!wslay_event_get_read_enabled(wsl_ctx)) { + // Some protocol error caused wslay to stop processing incoming events, we'll never receive a close from the other peer. + close_code = wslay_event_get_status_code_sent(wsl_ctx); + switch (close_code) { + case WSLAY_CODE_MESSAGE_TOO_BIG: + close_reason = "Message too big"; + break; + case WSLAY_CODE_PROTOCOL_ERROR: + close_reason = "Protocol error"; + break; + case WSLAY_CODE_ABNORMAL_CLOSURE: + close_reason = "Abnormal closure"; + break; + case WSLAY_CODE_INVALID_FRAME_PAYLOAD_DATA: + close_reason = "Invalid frame payload data"; + break; + default: + close_reason = "Unknown"; + } + wslay_event_context_free(wsl_ctx); + wsl_ctx = nullptr; + close(-1); + return; + } } } } @@ -783,8 +866,10 @@ void WSLPeer::close(int p_code, String p_reason) { } } + heartbeat_waiting = false; in_buffer.clear(); packet_buffer.resize(0); + pending_message.clear(); } IPAddress WSLPeer::get_connected_host() const { |