Allowed wsl_peer to grab the lenth of incoming messages and only acce…

…pt them if they fit in in_buffer
RadenTheFolf · Oct 13, 2024 · bd58183 · bd58183
1 parent 83d6573
commit bd58183
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 1 deletion.
diff --git 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
@@ -566,10 +566,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) {
@@ -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<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;
@@ -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
@@ -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();