From cb4aaa2235e4211d75b9a9b919552f25a875b3b1 Mon Sep 17 00:00:00 2001
From: radenthefolf <radenthefolf@gmail.com>
Date: Sun, 13 Oct 2024 07:21:36 -0400
Subject: [PATCH] Allowed wsl_peer to grab the lenth of incoming messages and
 only accept them if they fit in in_buffer

---
 modules/websocket/packet_buffer.h |  4 ++++
 modules/websocket/wsl_peer.cpp    | 17 ++++++++++++++++-
 modules/websocket/wsl_peer.h      |  3 +++
 3 files changed, 23 insertions(+), 1 deletion(-)

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..8e90b613e7b 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<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;
 	}
+
+	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,16 @@ 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) {
+		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<StreamPeer> conn = peer->connection;
@@ -645,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
diff --git a/modules/websocket/wsl_peer.h b/modules/websocket/wsl_peer.h
index 72214476830..33c9185aeb7 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:
+	int length_needed = 0;
+
 	static void initialize();
 	static void deinitialize();