Added set_extra_headers() to WebSocketServer

7 files changed, 25 insertions, 3 deletions

diff --git a/modules/websocket/doc_classes/WebSocketServer.xml b/modules/websocket/doc_classes/WebSocketServer.xml
index ef3279aac4..46b0274de3 100644
--- a/modules/websocket/doc_classes/WebSocketServer.xml
+++ b/modules/websocket/doc_classes/WebSocketServer.xml
@@ -60,6 +60,13 @@
 				If [code]false[/code] is passed instead (default), you must call [PacketPeer] functions ([code]put_packet[/code], [code]get_packet[/code], etc.), on the [WebSocketPeer] returned via [code]get_peer(id)[/code] to communicate with the peer with given [code]id[/code] (e.g. [code]get_peer(id).get_available_packet_count[/code]).
 			</description>
 		</method>
+		<method name="set_extra_headers">
+			<return type="void" />
+			<argument index="0" name="headers" type="PackedStringArray" default="PackedStringArray()" />
+			<description>
+				Sets additional headers to be sent to clients during the HTTP handshake.
+			</description>
+		</method>
 		<method name="stop">
 			<return type="void" />
 			<description>
diff --git a/modules/websocket/emws_server.cpp b/modules/websocket/emws_server.cpp
index 53b4a0207d..2033098cad 100644
--- a/modules/websocket/emws_server.cpp
+++ b/modules/websocket/emws_server.cpp
@@ -33,6 +33,9 @@
 #include "emws_server.h"
 #include "core/os/os.h"
 
+void EMWSServer::set_extra_headers(const Vector<String> &p_headers) {
+}
+
 Error EMWSServer::listen(int p_port, Vector<String> p_protocols, bool gd_mp_api) {
 	return FAILED;
 }
diff --git a/modules/websocket/emws_server.h b/modules/websocket/emws_server.h
index 0d193d423a..ae31d9dbb0 100644
--- a/modules/websocket/emws_server.h
+++ b/modules/websocket/emws_server.h
@@ -42,6 +42,7 @@ class EMWSServer : public WebSocketServer {
 
 public:
 	Error set_buffers(int p_in_buffer, int p_in_packets, int p_out_buffer, int p_out_packets) override;
+	void set_extra_headers(const Vector<String> &p_headers) override;
 	Error listen(int p_port, Vector<String> p_protocols = Vector<String>(), bool gd_mp_api = false) override;
 	void stop() override;
 	bool is_listening() const override;
diff --git a/modules/websocket/websocket_server.cpp b/modules/websocket/websocket_server.cpp
index b3f0140b80..b7851b02c4 100644
--- a/modules/websocket/websocket_server.cpp
+++ b/modules/websocket/websocket_server.cpp
@@ -42,6 +42,7 @@ WebSocketServer::~WebSocketServer() {
 
 void WebSocketServer::_bind_methods() {
 	ClassDB::bind_method(D_METHOD("is_listening"), &WebSocketServer::is_listening);
+	ClassDB::bind_method(D_METHOD("set_extra_headers", "headers"), &WebSocketServer::set_extra_headers, DEFVAL(Vector<String>()));
 	ClassDB::bind_method(D_METHOD("listen", "port", "protocols", "gd_mp_api"), &WebSocketServer::listen, DEFVAL(Vector<String>()), DEFVAL(false));
 	ClassDB::bind_method(D_METHOD("stop"), &WebSocketServer::stop);
 	ClassDB::bind_method(D_METHOD("has_peer", "id"), &WebSocketServer::has_peer);
diff --git a/modules/websocket/websocket_server.h b/modules/websocket/websocket_server.h
index f6f3b80045..7bd80851f5 100644
--- a/modules/websocket/websocket_server.h
+++ b/modules/websocket/websocket_server.h
@@ -51,6 +51,7 @@ protected:
 	uint32_t handshake_timeout = 3000;
 
 public:
+	virtual void set_extra_headers(const Vector<String> &p_headers) = 0;
 	virtual Error listen(int p_port, const Vector<String> p_protocols = Vector<String>(), bool gd_mp_api = false) = 0;
 	virtual void stop() = 0;
 	virtual bool is_listening() const = 0;
diff --git a/modules/websocket/wsl_server.cpp b/modules/websocket/wsl_server.cpp
index 8cd4b78ab3..b58b2e4724 100644
--- a/modules/websocket/wsl_server.cpp
+++ b/modules/websocket/wsl_server.cpp
@@ -96,7 +96,7 @@ bool WSLServer::PendingPeer::_parse_request(const Vector<String> p_protocols, St
 	return true;
 }
 
-Error WSLServer::PendingPeer::do_handshake(const Vector<String> p_protocols, uint64_t p_timeout, String &r_resource_name) {
+Error WSLServer::PendingPeer::do_handshake(const Vector<String> p_protocols, uint64_t p_timeout, String &r_resource_name, const Vector<String> &p_extra_headers) {
 	if (OS::get_singleton()->get_ticks_msec() - time > p_timeout) {
 		print_verbose(vformat("WebSocket handshake timed out after %.3f seconds.", p_timeout * 0.001));
 		return ERR_TIMEOUT;
@@ -141,6 +141,9 @@ Error WSLServer::PendingPeer::do_handshake(const Vector<String> p_protocols, uin
 				if (!protocol.is_empty()) {
 					s += "Sec-WebSocket-Protocol: " + protocol + "\r\n";
 				}
+				for (int i = 0; i < p_extra_headers.size(); i++) {
+					s += p_extra_headers[i] + "\r\n";
+				}
 				s += "\r\n";
 				response = s.utf8();
 				has_request = true;
@@ -167,6 +170,10 @@ Error WSLServer::PendingPeer::do_handshake(const Vector<String> p_protocols, uin
 	return OK;
 }
 
+void WSLServer::set_extra_headers(const Vector<String> &p_headers) {
+	_extra_headers = p_headers;
+}
+
 Error WSLServer::listen(int p_port, const Vector<String> p_protocols, bool gd_mp_api) {
 	ERR_FAIL_COND_V(is_listening(), ERR_ALREADY_IN_USE);
 
@@ -199,7 +206,7 @@ void WSLServer::poll() {
 	for (const Ref<PendingPeer> &E : _pending) {
 		String resource_name;
 		Ref<PendingPeer> ppeer = E;
-		Error err = ppeer->do_handshake(_protocols, handshake_timeout, resource_name);
+		Error err = ppeer->do_handshake(_protocols, handshake_timeout, resource_name, _extra_headers);
 		if (err == ERR_BUSY) {
 			continue;
 		} else if (err != OK) {
diff --git a/modules/websocket/wsl_server.h b/modules/websocket/wsl_server.h
index 6a9dd0dd2f..a920e9c665 100644
--- a/modules/websocket/wsl_server.h
+++ b/modules/websocket/wsl_server.h
@@ -62,7 +62,7 @@ private:
 		CharString response;
 		int response_sent = 0;
 
-		Error do_handshake(const Vector<String> p_protocols, uint64_t p_timeout, String &r_resource_name);
+		Error do_handshake(const Vector<String> p_protocols, uint64_t p_timeout, String &r_resource_name, const Vector<String> &p_extra_headers);
 	};
 
 	int _in_buf_size = DEF_BUF_SHIFT;
@@ -73,9 +73,11 @@ private:
 	List<Ref<PendingPeer>> _pending;
 	Ref<TCPServer> _server;
 	Vector<String> _protocols;
+	Vector<String> _extra_headers;
 
 public:
 	Error set_buffers(int p_in_buffer, int p_in_packets, int p_out_buffer, int p_out_packets) override;
+	void set_extra_headers(const Vector<String> &p_headers) override;
 	Error listen(int p_port, const Vector<String> p_protocols = Vector<String>(), bool gd_mp_api = false) override;
 	void stop() override;
 	bool is_listening() const override;