From dc534d1a5493bb08e90d3eb5bb88b006b430fef5 Mon Sep 17 00:00:00 2001
From: radenthefolf <radenthefolf@gmail.com>
Date: Tue, 8 Oct 2024 01:40:18 -0400
Subject: [PATCH] Fixes integration with wslay to prevent accepting packets
 when in_buffer full

---
 modules/websocket/doc_classes/WebSocketPeer.xml |  3 +++
 modules/websocket/packet_buffer.h               |  4 ++++
 modules/websocket/websocket_peer.cpp            | 12 ++++++++++++
 modules/websocket/websocket_peer.h              |  4 ++++
 modules/websocket/wsl_peer.cpp                  |  4 ++++
 5 files changed, 27 insertions(+)

diff --git a/modules/websocket/doc_classes/WebSocketPeer.xml b/modules/websocket/doc_classes/WebSocketPeer.xml
index 238dd305368a..6c42dd109f3d 100644
--- a/modules/websocket/doc_classes/WebSocketPeer.xml
+++ b/modules/websocket/doc_classes/WebSocketPeer.xml
@@ -161,6 +161,9 @@
 		<member name="max_queued_packets" type="int" setter="set_max_queued_packets" getter="get_max_queued_packets" default="2048">
 			The maximum amount of packets that will be allowed in the queues (both inbound and outbound).
 		</member>
+		<member name="min_buffer_free_space" type="int" setter="set_min_buffer_free_space" getter="get_min_buffer_free_space" default="4096">
+			The minimum amount of free space allowed to be in the buffer before reading more messages.
+		</member>
 		<member name="outbound_buffer_size" type="int" setter="set_outbound_buffer_size" getter="get_outbound_buffer_size" default="65535">
 			The size of the input buffer in bytes (roughly the maximum amount of memory that will be allocated for the outbound packets).
 		</member>
diff --git a/modules/websocket/packet_buffer.h b/modules/websocket/packet_buffer.h
index f98ee12ef9a6..b1a7f70eb6d2 100644
--- a/modules/websocket/packet_buffer.h
+++ b/modules/websocket/packet_buffer.h
@@ -104,6 +104,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/websocket_peer.cpp b/modules/websocket/websocket_peer.cpp
index 95a1a238e95a..3a6499cc462a 100644
--- a/modules/websocket/websocket_peer.cpp
+++ b/modules/websocket/websocket_peer.cpp
@@ -70,6 +70,9 @@ void WebSocketPeer::_bind_methods() {
 	ClassDB::bind_method(D_METHOD("set_max_queued_packets", "buffer_size"), &WebSocketPeer::set_max_queued_packets);
 	ClassDB::bind_method(D_METHOD("get_max_queued_packets"), &WebSocketPeer::get_max_queued_packets);
 
+	ClassDB::bind_method(D_METHOD("set_min_buffer_free_space", "min_buffer_free_space"), &WebSocketPeer::set_min_buffer_free_space);
+	ClassDB::bind_method(D_METHOD("get_min_buffer_free_space"), &WebSocketPeer::get_min_buffer_free_space);
+
 	ADD_PROPERTY(PropertyInfo(Variant::PACKED_STRING_ARRAY, "supported_protocols"), "set_supported_protocols", "get_supported_protocols");
 	ADD_PROPERTY(PropertyInfo(Variant::PACKED_STRING_ARRAY, "handshake_headers"), "set_handshake_headers", "get_handshake_headers");
 
@@ -77,6 +80,7 @@ void WebSocketPeer::_bind_methods() {
 	ADD_PROPERTY(PropertyInfo(Variant::INT, "outbound_buffer_size"), "set_outbound_buffer_size", "get_outbound_buffer_size");
 
 	ADD_PROPERTY(PropertyInfo(Variant::INT, "max_queued_packets"), "set_max_queued_packets", "get_max_queued_packets");
+	ADD_PROPERTY(PropertyInfo(Variant::INT, "min_buffer_free_space"), "set_min_buffer_free_space", "get_min_buffer_free_space");
 
 	BIND_ENUM_CONSTANT(WRITE_MODE_TEXT);
 	BIND_ENUM_CONSTANT(WRITE_MODE_BINARY);
@@ -151,3 +155,11 @@ void WebSocketPeer::set_max_queued_packets(int p_max_queued_packets) {
 int WebSocketPeer::get_max_queued_packets() const {
 	return max_queued_packets;
 }
+
+void WebSocketPeer::set_min_buffer_free_space(int p_min_buffer_free_space) {
+	min_buffer_free_space = p_min_buffer_free_space;
+}
+
+int WebSocketPeer::get_min_buffer_free_space() const {
+	return min_buffer_free_space;
+}
diff --git a/modules/websocket/websocket_peer.h b/modules/websocket/websocket_peer.h
index ef0197cf6c96..bfd12cb19e0b 100644
--- a/modules/websocket/websocket_peer.h
+++ b/modules/websocket/websocket_peer.h
@@ -72,6 +72,7 @@ class WebSocketPeer : public PacketPeer {
 	int outbound_buffer_size = DEFAULT_BUFFER_SIZE;
 	int inbound_buffer_size = DEFAULT_BUFFER_SIZE;
 	int max_queued_packets = 2048;
+	int min_buffer_free_space = 4096;
 
 public:
 	static WebSocketPeer *create(bool p_notify_postinitialize = true) {
@@ -117,6 +118,9 @@ class WebSocketPeer : public PacketPeer {
 	void set_max_queued_packets(int p_max_queued_packets);
 	int get_max_queued_packets() const;
 
+	void set_min_buffer_free_space(int p_min_buffer_free_space);
+	int get_min_buffer_free_space() const;
+
 	WebSocketPeer();
 	~WebSocketPeer();
 };
diff --git a/modules/websocket/wsl_peer.cpp b/modules/websocket/wsl_peer.cpp
index 0c0a046805c1..42a974c1826a 100644
--- a/modules/websocket/wsl_peer.cpp
+++ b/modules/websocket/wsl_peer.cpp
@@ -564,6 +564,10 @@ 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 (peer->in_buffer.space_left() < peer->min_buffer_free_space) {
+		return WSLAY_ERR_NOMEM;
+	}
+
 	if (conn.is_null()) {
 		wslay_event_set_error(ctx, WSLAY_ERR_CALLBACK_FAILURE);
 		return -1;