Skip to content

Commit

Permalink
Fix websocket dropping messages for WSLPeer
Browse files Browse the repository at this point in the history
Allows WSLPeer to grab the length of incoming messages
	Only accept them if they fit in in_buffer

(cherry picked from commit e11d581)
  • Loading branch information
RadenTheFolf authored and Spartan322 committed Oct 14, 2024
1 parent 50c9bfb commit b0fffe2
Show file tree
Hide file tree
Showing 3 changed files with 24 additions and 1 deletion.
4 changes: 4 additions & 0 deletions modules/websocket/packet_buffer.h
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,10 @@ class PacketBuffer {
return _queued;
}

int space_left() const {
return _payload.space_left();
}

void clear() {
_payload.resize(0);
_packets.resize(0);
Expand Down
18 changes: 17 additions & 1 deletion modules/websocket/wsl_peer.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -565,10 +565,15 @@ Error WSLPeer::connect_to_url(const String &p_url, Ref<TLSOptions> 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<StreamPeer> 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) {
Expand All @@ -583,6 +588,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<StreamPeer> 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<StreamPeer> conn = peer->connection;
Expand Down Expand Up @@ -644,7 +660,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
Expand Down
3 changes: 3 additions & 0 deletions modules/websocket/wsl_peer.h
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -122,6 +123,8 @@ class WSLPeer : public WebSocketPeer {
void _clear();

public:
uint64_t length_needed = 0;

static void initialize();
static void deinitialize();

Expand Down

0 comments on commit b0fffe2

Please sign in to comment.