1 files changed, 40 insertions, 25 deletions

diff --git a/core/io/packet_peer_udp.cpp b/core/io/packet_peer_udp.cpp
index 862fca96fc..e633a56d54 100644
--- a/core/io/packet_peer_udp.cpp
+++ b/core/io/packet_peer_udp.cpp
@@ -31,12 +31,14 @@
 #include "packet_peer_udp.h"
 
 #include "core/io/ip.h"
+#include "core/io/udp_server.h"
 
 void PacketPeerUDP::set_blocking_mode(bool p_enable) {
 	blocking = p_enable;
 }
 
 void PacketPeerUDP::set_broadcast_enabled(bool p_enabled) {
+	ERR_FAIL_COND(udp_server);
 	broadcast = p_enabled;
 	if (_sock.is_valid() && _sock->is_open()) {
 		_sock->set_broadcasting_enabled(p_enabled);
@@ -44,6 +46,7 @@ void PacketPeerUDP::set_broadcast_enabled(bool p_enabled) {
 }
 
 Error PacketPeerUDP::join_multicast_group(IP_Address p_multi_address, String p_if_name) {
+	ERR_FAIL_COND_V(udp_server, ERR_LOCKED);
 	ERR_FAIL_COND_V(!_sock.is_valid(), ERR_UNAVAILABLE);
 	ERR_FAIL_COND_V(!p_multi_address.is_valid(), ERR_INVALID_PARAMETER);
 
@@ -58,6 +61,7 @@ Error PacketPeerUDP::join_multicast_group(IP_Address p_multi_address, String p_i
 }
 
 Error PacketPeerUDP::leave_multicast_group(IP_Address p_multi_address, String p_if_name) {
+	ERR_FAIL_COND_V(udp_server, ERR_LOCKED);
 	ERR_FAIL_COND_V(!_sock.is_valid(), ERR_UNAVAILABLE);
 	ERR_FAIL_COND_V(!_sock->is_open(), ERR_UNCONFIGURED);
 	return _sock->leave_multicast_group(p_multi_address, p_if_name);
@@ -130,7 +134,7 @@ Error PacketPeerUDP::put_packet(const uint8_t *p_buffer, int p_buffer_size) {
 	}
 
 	do {
-		if (connected) {
+		if (connected && !udp_server) {
 			err = _sock->send(p_buffer, p_buffer_size, sent);
 		} else {
 			err = _sock->sendto(p_buffer, p_buffer_size, sent, peer_addr, peer_port);
@@ -186,26 +190,25 @@ Error PacketPeerUDP::listen(int p_port, const IP_Address &p_bind_address, int p_
 	return OK;
 }
 
-Error PacketPeerUDP::connect_socket(Ref<NetSocket> p_sock) {
-	Error err;
-	int read = 0;
-	uint16_t r_port;
-	IP_Address r_ip;
-
-	err = p_sock->recvfrom(recv_buffer, sizeof(recv_buffer), read, r_ip, r_port, true);
-	ERR_FAIL_COND_V(err != OK, err);
-	err = p_sock->connect_to_host(r_ip, r_port);
-	ERR_FAIL_COND_V(err != OK, err);
+Error PacketPeerUDP::connect_shared_socket(Ref<NetSocket> p_sock, IP_Address p_ip, uint16_t p_port, UDPServer *p_server) {
+	udp_server = p_server;
+	connected = true;
 	_sock = p_sock;
-	peer_addr = r_ip;
-	peer_port = r_port;
+	peer_addr = p_ip;
+	peer_port = p_port;
 	packet_ip = peer_addr;
 	packet_port = peer_port;
-	connected = true;
 	return OK;
 }
 
+void PacketPeerUDP::disconnect_shared_socket() {
+	udp_server = nullptr;
+	_sock = Ref<NetSocket>(NetSocket::create());
+	close();
+}
+
 Error PacketPeerUDP::connect_to_host(const IP_Address &p_host, int p_port) {
+	ERR_FAIL_COND_V(udp_server, ERR_LOCKED);
 	ERR_FAIL_COND_V(!_sock.is_valid(), ERR_UNAVAILABLE);
 	ERR_FAIL_COND_V(!p_host.is_valid(), ERR_INVALID_PARAMETER);
 
@@ -243,7 +246,11 @@ bool PacketPeerUDP::is_connected_to_host() const {
 }
 
 void PacketPeerUDP::close() {
-	if (_sock.is_valid()) {
+	if (udp_server) {
+		udp_server->remove_peer(peer_addr, peer_port);
+		udp_server = nullptr;
+		_sock = Ref<NetSocket>(NetSocket::create());
+	} else if (_sock.is_valid()) {
 		_sock->close();
 	}
 	rb.resize(16);
@@ -262,6 +269,9 @@ Error PacketPeerUDP::_poll() {
 	if (!_sock->is_open()) {
 		return FAILED;
 	}
+	if (udp_server) {
+		return OK; // Handled by UDPServer.
+	}
 
 	Error err;
 	int read;
@@ -284,24 +294,29 @@ Error PacketPeerUDP::_poll() {
 			return FAILED;
 		}
 
-		if (rb.space_left() < read + 24) {
+		err = store_packet(ip, port, recv_buffer, read);
 #ifdef TOOLS_ENABLED
+		if (err != OK) {
 			WARN_PRINT("Buffer full, dropping packets!");
-#endif
-			continue;
 		}
-
-		uint32_t port32 = port;
-		rb.write(ip.get_ipv6(), 16);
-		rb.write((uint8_t *)&port32, 4);
-		rb.write((uint8_t *)&read, 4);
-		rb.write(recv_buffer, read);
-		++queue_count;
+#endif
 	}
 
 	return OK;
 }
 
+Error PacketPeerUDP::store_packet(IP_Address p_ip, uint32_t p_port, uint8_t *p_buf, int p_buf_size) {
+	if (rb.space_left() < p_buf_size + 24) {
+		return ERR_OUT_OF_MEMORY;
+	}
+	rb.write(p_ip.get_ipv6(), 16);
+	rb.write((uint8_t *)&p_port, 4);
+	rb.write((uint8_t *)&p_buf_size, 4);
+	rb.write(p_buf, p_buf_size);
+	++queue_count;
+	return OK;
+}
+
 bool PacketPeerUDP::is_listening() const {
 	return _sock.is_valid() && _sock->is_open();
 }