diff --git a/modules/websocket/packet_buffer.h b/modules/websocket/packet_buffer.h index f98ee12ef9a..b1a7f70eb6d 100644 --- a/modules/websocket/packet_buffer.h +++ b/modules/websocket/packet_buffer.h @@ -104,6 +104,10 @@ public: 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 0c0a046805c..8587912b15e 100644 --- a/modules/websocket/wsl_peer.cpp +++ b/modules/websocket/wsl_peer.cpp @@ -564,10 +564,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) { @@ -582,6 +587,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; @@ -643,7 +659,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 fb01da7ce29..ec8a213e0a9 100644 --- a/modules/websocket/wsl_peer.h +++ b/modules/websocket/wsl_peer.h @@ -53,6 +53,7 @@ private: // 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); @@ -120,6 +121,8 @@ private: void _clear(); public: + uint64_t length_needed = 0; + static void initialize(); static void deinitialize();