4 files changed, 32 insertions, 23 deletions

diff --git a/modules/websocket/register_types.cpp b/modules/websocket/register_types.cpp
index 6d63938d4f..f562de111f 100644
--- a/modules/websocket/register_types.cpp
+++ b/modules/websocket/register_types.cpp
@@ -55,25 +55,33 @@ static void _editor_init_callback() {
 }
 #endif
 
-void register_websocket_types() {
+void initialize_websocket_module(ModuleInitializationLevel p_level) {
+	if (p_level == MODULE_INITIALIZATION_LEVEL_SCENE) {
 #ifdef JAVASCRIPT_ENABLED
-	EMWSPeer::make_default();
-	EMWSClient::make_default();
-	EMWSServer::make_default();
+		EMWSPeer::make_default();
+		EMWSClient::make_default();
+		EMWSServer::make_default();
 #else
-	WSLPeer::make_default();
-	WSLClient::make_default();
-	WSLServer::make_default();
+		WSLPeer::make_default();
+		WSLClient::make_default();
+		WSLServer::make_default();
 #endif
 
-	GDREGISTER_ABSTRACT_CLASS(WebSocketMultiplayerPeer);
-	ClassDB::register_custom_instance_class<WebSocketServer>();
-	ClassDB::register_custom_instance_class<WebSocketClient>();
-	ClassDB::register_custom_instance_class<WebSocketPeer>();
+		GDREGISTER_ABSTRACT_CLASS(WebSocketMultiplayerPeer);
+		ClassDB::register_custom_instance_class<WebSocketServer>();
+		ClassDB::register_custom_instance_class<WebSocketClient>();
+		ClassDB::register_custom_instance_class<WebSocketPeer>();
+	}
 
 #ifdef TOOLS_ENABLED
-	EditorNode::add_init_callback(&_editor_init_callback);
+	if (p_level == MODULE_INITIALIZATION_LEVEL_EDITOR) {
+		EditorNode::add_init_callback(&_editor_init_callback);
+	}
 #endif
 }
 
-void unregister_websocket_types() {}
+void uninitialize_websocket_module(ModuleInitializationLevel p_level) {
+	if (p_level != MODULE_INITIALIZATION_LEVEL_SCENE) {
+		return;
+	}
+}
diff --git a/modules/websocket/register_types.h b/modules/websocket/register_types.h
index 4ab6c0cfd3..dab42d6ed9 100644
--- a/modules/websocket/register_types.h
+++ b/modules/websocket/register_types.h
@@ -31,7 +31,9 @@
 #ifndef WEBSOCKET_REGISTER_TYPES_H
 #define WEBSOCKET_REGISTER_TYPES_H
 
-void register_websocket_types();
-void unregister_websocket_types();
+#include "modules/register_module_types.h"
+
+void initialize_websocket_module(ModuleInitializationLevel p_level);
+void uninitialize_websocket_module(ModuleInitializationLevel p_level);
 
 #endif // WEBSOCKET_REGISTER_TYPES_H
diff --git a/modules/websocket/wsl_client.cpp b/modules/websocket/wsl_client.cpp
index 6f0b726f5c..894ba7766f 100644
--- a/modules/websocket/wsl_client.cpp
+++ b/modules/websocket/wsl_client.cpp
@@ -91,6 +91,7 @@ void WSLClient::_do_handshake() {
 				data->id = 1;
 				_peer->make_context(data, _in_buf_size, _in_pkt_size, _out_buf_size, _out_pkt_size);
 				_peer->set_no_delay(true);
+				_status = CONNECTION_CONNECTED;
 				_on_connect(protocol);
 				break;
 			}
@@ -231,6 +232,7 @@ Error WSLClient::connect_to_host(String p_host, String p_path, uint16_t p_port,
 	}
 	request += "\r\n";
 	_request = request.utf8();
+	_status = CONNECTION_CONNECTING;
 
 	return OK;
 }
@@ -337,21 +339,19 @@ Ref<WebSocketPeer> WSLClient::get_peer(int p_peer_id) const {
 }
 
 MultiplayerPeer::ConnectionStatus WSLClient::get_connection_status() const {
+	// This is surprising, but keeps the current behaviour to allow clean close requests.
+	// TODO Refactor WebSocket and split Client/Server/Multiplayer like done in other peers.
 	if (_peer->is_connected_to_host()) {
 		return CONNECTION_CONNECTED;
 	}
-
-	if (_tcp->get_status() == StreamPeerTCP::STATUS_CONNECTING || _resolver_id != IP::RESOLVER_INVALID_ID) {
-		return CONNECTION_CONNECTING;
-	}
-
-	return CONNECTION_DISCONNECTED;
+	return _status;
 }
 
 void WSLClient::disconnect_from_host(int p_code, String p_reason) {
 	_peer->close(p_code, p_reason);
 	_connection = Ref<StreamPeer>(nullptr);
 	_tcp = Ref<StreamPeerTCP>(memnew(StreamPeerTCP));
+	_status = CONNECTION_DISCONNECTED;
 
 	_key = "";
 	_host = "";
diff --git a/modules/websocket/wsl_client.h b/modules/websocket/wsl_client.h
index 04af83f28c..22d7ffa839 100644
--- a/modules/websocket/wsl_client.h
+++ b/modules/websocket/wsl_client.h
@@ -52,6 +52,7 @@ private:
 	Ref<WSLPeer> _peer;
 	Ref<StreamPeerTCP> _tcp;
 	Ref<StreamPeer> _connection;
+	ConnectionStatus _status = CONNECTION_DISCONNECTED;
 
 	CharString _request;
 	int _requested = 0;
@@ -59,8 +60,6 @@ private:
 	uint8_t _resp_buf[WSL_MAX_HEADER_SIZE];
 	int _resp_pos = 0;
 
-	String _response;
-
 	String _key;
 	String _host;
 	uint16_t _port = 0;