diff --git a/modules/websocket/packet_buffer.h b/modules/websocket/packet_buffer.h index cec946204b9..d67add73b72 100644 --- a/modules/websocket/packet_buffer.h +++ b/modules/websocket/packet_buffer.h @@ -106,6 +106,10 @@ class PacketBuffer { return _queued; } + int space_left() const { + return _payload.space_left(); + } + void clear() { _payload.resize(0); _packets.resize(0); diff --git a/modules/websocket/wsl_peer.cpp b/modules/websocket/wsl_peer.cpp index dff452d1258..f3ab8116089 100644 --- a/modules/websocket/wsl_peer.cpp +++ b/modules/websocket/wsl_peer.cpp @@ -566,10 +566,15 @@ Error WSLPeer::connect_to_url(const String &p_url, Ref p_options) { ssize_t WSLPeer::_wsl_recv_callback(wslay_event_context_ptr ctx, uint8_t *data, size_t len, int flags, void *user_data) { WSLPeer *peer = (WSLPeer *)user_data; Ref conn = peer->connection; + if (conn.is_null()) { wslay_event_set_error(ctx, WSLAY_ERR_CALLBACK_FAILURE); return -1; } + uint64_t space_left = peer->in_buffer.space_left(); + if (space_left < peer->length_needed) { + return WSLAY_ERR_NOMEM; + } int read = 0; Error err = conn->get_partial_data(data, len, read); if (err != OK) { @@ -584,6 +589,17 @@ ssize_t WSLPeer::_wsl_recv_callback(wslay_event_context_ptr ctx, uint8_t *data, return read; } +void WSLPeer::_wsl_on_frame_start_callback(wslay_event_context_ptr ctx, const wslay_event_on_frame_recv_start_arg *arg, void *user_data) { + WSLPeer *peer = (WSLPeer *)user_data; + Ref conn = peer->connection; + if (arg->opcode == WSLAY_TEXT_FRAME || arg->opcode == WSLAY_BINARY_FRAME) { + uint64_t space_left = peer->in_buffer.space_left(); + if (space_left < arg->payload_length) { + peer->length_needed = arg->payload_length; + } + } +} + 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 conn = peer->connection; @@ -645,7 +661,7 @@ wslay_event_callbacks WSLPeer::_wsl_callbacks = { _wsl_recv_callback, _wsl_send_callback, _wsl_genmask_callback, - nullptr, /* on_frame_recv_start_callback */ + _wsl_on_frame_start_callback, /* on_frame_recv_start_callback */ nullptr, /* on_frame_recv_callback */ nullptr, /* on_frame_recv_end_callback */ _wsl_msg_recv_callback diff --git a/modules/websocket/wsl_peer.h b/modules/websocket/wsl_peer.h index 72214476830..bd720be5e22 100644 --- a/modules/websocket/wsl_peer.h +++ b/modules/websocket/wsl_peer.h @@ -55,6 +55,7 @@ class WSLPeer : public WebSocketPeer { // Callbacks. static ssize_t _wsl_recv_callback(wslay_event_context_ptr ctx, uint8_t *data, size_t len, int flags, void *user_data); + static void _wsl_on_frame_start_callback(wslay_event_context_ptr ctx, const struct wslay_event_on_frame_recv_start_arg *arg, void *user_data); static ssize_t _wsl_send_callback(wslay_event_context_ptr ctx, const uint8_t *data, size_t len, int flags, void *user_data); static int _wsl_genmask_callback(wslay_event_context_ptr ctx, uint8_t *buf, size_t len, void *user_data); static void _wsl_msg_recv_callback(wslay_event_context_ptr ctx, const struct wslay_event_on_msg_recv_arg *arg, void *user_data); @@ -122,6 +123,8 @@ class WSLPeer : public WebSocketPeer { void _clear(); public: + uint64_t length_needed = 0; + static void initialize(); static void deinitialize();