1 files changed, 46 insertions, 12 deletions

diff --git a/modules/websocket/lws_peer.cpp b/modules/websocket/lws_peer.cpp
index 0989357258..b5c130b308 100644
--- a/modules/websocket/lws_peer.cpp
+++ b/modules/websocket/lws_peer.cpp
@@ -30,6 +30,7 @@
 #ifndef JAVASCRIPT_ENABLED
 
 #include "lws_peer.h"
+
 #include "core/io/ip.h"
 
 // Needed for socket_helpers on Android at least. UNIXes has it, just include if not windows
@@ -38,7 +39,7 @@
 #include <sys/socket.h>
 #endif
 
-#include "drivers/unix/socket_helpers.h"
+#include "drivers/unix/net_socket_posix.h"
 
 void LWSPeer::set_wsi(struct lws *p_wsi) {
 	ERR_FAIL_COND(wsi != NULL);
@@ -60,7 +61,6 @@ Error LWSPeer::read_wsi(void *in, size_t len) {
 
 	ERR_FAIL_COND_V(!is_connected_to_host(), FAILED);
 
-	PeerData *peer_data = (PeerData *)(lws_wsi_user(wsi));
 	uint32_t size = in_size;
 	uint8_t is_string = lws_frame_is_binary(wsi) ? 0 : 1;
 
@@ -88,7 +88,6 @@ Error LWSPeer::write_wsi() {
 
 	ERR_FAIL_COND_V(!is_connected_to_host(), FAILED);
 
-	PeerData *peer_data = (PeerData *)(lws_wsi_user(wsi));
 	PoolVector<uint8_t> tmp;
 	int left = rbw.data_left();
 	uint32_t to_write = 0;
@@ -119,7 +118,6 @@ Error LWSPeer::put_packet(const uint8_t *p_buffer, int p_buffer_size) {
 
 	ERR_FAIL_COND_V(!is_connected_to_host(), FAILED);
 
-	PeerData *peer_data = (PeerData *)lws_wsi_user(wsi);
 	rbw.write((uint8_t *)&p_buffer_size, 4);
 	rbw.write(p_buffer, MIN(p_buffer_size, rbw.space_left()));
 	out_count++;
@@ -132,8 +130,6 @@ Error LWSPeer::get_packet(const uint8_t **r_buffer, int &r_buffer_size) {
 
 	ERR_FAIL_COND_V(!is_connected_to_host(), FAILED);
 
-	PeerData *peer_data = (PeerData *)lws_wsi_user(wsi);
-
 	if (in_count == 0)
 		return ERR_UNAVAILABLE;
 
@@ -178,11 +174,49 @@ bool LWSPeer::is_connected_to_host() const {
 	return wsi != NULL;
 };
 
-void LWSPeer::close() {
+String LWSPeer::get_close_reason(void *in, size_t len, int &r_code) {
+	String s;
+	r_code = 0;
+	if (len < 2) // From docs this should not happen
+		return s;
+
+	const uint8_t *b = (const uint8_t *)in;
+	r_code = b[0] << 8 | b[1];
+
+	if (len > 2) {
+		const char *utf8 = (const char *)&b[2];
+		s.parse_utf8(utf8, len - 2);
+	}
+	return s;
+}
+
+void LWSPeer::send_close_status(struct lws *p_wsi) {
+	if (close_code == -1)
+		return;
+
+	int len = close_reason.size();
+	ERR_FAIL_COND(len > 123); // Maximum allowed reason size in bytes
+
+	lws_close_status code = (lws_close_status)close_code;
+	unsigned char *reason = len > 0 ? (unsigned char *)close_reason.utf8().ptrw() : NULL;
+
+	lws_close_reason(p_wsi, code, reason, len);
+
+	close_code = -1;
+	close_reason = "";
+}
+
+void LWSPeer::close(int p_code, String p_reason) {
 	if (wsi != NULL) {
+		close_code = p_code;
+		close_reason = p_reason;
 		PeerData *data = ((PeerData *)lws_wsi_user(wsi));
 		data->force_close = true;
-		lws_callback_on_writable(wsi); // notify that we want to disconnect
+		data->clean_close = true;
+		lws_callback_on_writable(wsi); // Notify that we want to disconnect
+	} else {
+		close_code = -1;
+		close_reason = "";
 	}
 	wsi = NULL;
 	rbw.resize(0);
@@ -198,7 +232,7 @@ IP_Address LWSPeer::get_connected_host() const {
 	ERR_FAIL_COND_V(!is_connected_to_host(), IP_Address());
 
 	IP_Address ip;
-	int port = 0;
+	uint16_t port = 0;
 
 	struct sockaddr_storage addr;
 	socklen_t len = sizeof(addr);
@@ -209,7 +243,7 @@ IP_Address LWSPeer::get_connected_host() const {
 	int ret = getpeername(fd, (struct sockaddr *)&addr, &len);
 	ERR_FAIL_COND_V(ret != 0, IP_Address());
 
-	_set_ip_addr_port(ip, port, &addr);
+	NetSocketPosix::_set_ip_port(&addr, ip, port);
 
 	return ip;
 };
@@ -219,7 +253,7 @@ uint16_t LWSPeer::get_connected_port() const {
 	ERR_FAIL_COND_V(!is_connected_to_host(), 0);
 
 	IP_Address ip;
-	int port = 0;
+	uint16_t port = 0;
 
 	struct sockaddr_storage addr;
 	socklen_t len = sizeof(addr);
@@ -230,7 +264,7 @@ uint16_t LWSPeer::get_connected_port() const {
 	int ret = getpeername(fd, (struct sockaddr *)&addr, &len);
 	ERR_FAIL_COND_V(ret != 0, 0);
 
-	_set_ip_addr_port(ip, port, &addr);
+	NetSocketPosix::_set_ip_port(&addr, ip, port);
 
 	return port;
 };