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..6575ca9c05b 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; } + + if (peer->in_buffer.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,18 @@ 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) { + if (peer->in_buffer.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 +662,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..2596cbcff17 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); @@ -102,6 +103,8 @@ class WSLPeer : public WebSocketPeer { String close_reason; uint8_t was_string = 0; + + // WebSocket configuration. bool use_tls = true; Ref tls_options; @@ -122,6 +125,8 @@ class WSLPeer : public WebSocketPeer { void _clear(); public: + int length_needed = 0; + static void initialize(); static void deinitialize();