UDPServer handles PacketPeerUDP-client association

UDPServer now uses a single socket which is shared with the
PacketPeerUDP it creates and has a new `poll` function to read incoming
packets on that socket and delivers them to the appropriate peer.
PacketPeerUDP created this way never reads from the socket, but are
allowed to write on it using sendto.

This is needed because Windows (unlike Linux/BSD) does not support
packet routing when multiple sockets are bound on the same address/port.

1 files changed, 93 insertions, 7 deletions

diff --git a/core/io/udp_server.cpp b/core/io/udp_server.cpp
index 1d329daf8b..acd15aadc6 100644
--- a/core/io/udp_server.cpp
+++ b/core/io/udp_server.cpp
@@ -32,10 +32,58 @@
 
 void UDPServer::_bind_methods() {
 	ClassDB::bind_method(D_METHOD("listen", "port", "bind_address"), &UDPServer::listen, DEFVAL("*"));
+	ClassDB::bind_method(D_METHOD("poll"), &UDPServer::poll);
 	ClassDB::bind_method(D_METHOD("is_connection_available"), &UDPServer::is_connection_available);
 	ClassDB::bind_method(D_METHOD("is_listening"), &UDPServer::is_listening);
 	ClassDB::bind_method(D_METHOD("take_connection"), &UDPServer::take_connection);
 	ClassDB::bind_method(D_METHOD("stop"), &UDPServer::stop);
+	ClassDB::bind_method(D_METHOD("set_max_pending_connections", "max_pending_connections"), &UDPServer::set_max_pending_connections);
+	ClassDB::bind_method(D_METHOD("get_max_pending_connections"), &UDPServer::get_max_pending_connections);
+	ADD_PROPERTY(PropertyInfo(Variant::INT, "max_pending_connections", PROPERTY_HINT_RANGE, "0,256,1"), "set_max_pending_connections", "get_max_pending_connections");
+}
+
+Error UDPServer::poll() {
+	ERR_FAIL_COND_V(!_sock.is_valid(), ERR_UNAVAILABLE);
+	if (!_sock->is_open()) {
+		return ERR_UNCONFIGURED;
+	}
+	Error err;
+	int read;
+	IP_Address ip;
+	uint16_t port;
+	while (true) {
+		err = _sock->recvfrom(recv_buffer, sizeof(recv_buffer), read, ip, port);
+		if (err != OK) {
+			if (err == ERR_BUSY) {
+				break;
+			}
+			return FAILED;
+		}
+		Peer p;
+		p.ip = ip;
+		p.port = port;
+		List<Peer>::Element *E = peers.find(p);
+		if (!E) {
+			E = pending.find(p);
+		}
+		if (E) {
+			E->get().peer->store_packet(ip, port, recv_buffer, read);
+		} else {
+			if (pending.size() >= max_pending_connections) {
+				// Drop connection.
+				continue;
+			}
+			// It's a new peer, add it to the pending list.
+			Peer peer;
+			peer.ip = ip;
+			peer.port = port;
+			peer.peer = memnew(PacketPeerUDP);
+			peer.peer->connect_shared_socket(_sock, ip, port, this);
+			peer.peer->store_packet(ip, port, recv_buffer, read);
+			pending.push_back(peer);
+		}
+	}
+	return OK;
 }
 
 Error UDPServer::listen(uint16_t p_port, const IP_Address &p_bind_address) {
@@ -82,8 +130,24 @@ bool UDPServer::is_connection_available() const {
 		return false;
 	}
 
-	Error err = _sock->poll(NetSocket::POLL_TYPE_IN, 0);
-	return (err == OK);
+	return pending.size() > 0;
+}
+
+void UDPServer::set_max_pending_connections(int p_max) {
+	ERR_FAIL_COND_MSG(p_max < 0, "Max pending connections value must be a positive number (0 means refuse new connections).");
+	max_pending_connections = p_max;
+	while (p_max > pending.size()) {
+		List<Peer>::Element *E = pending.back();
+		if (!E) {
+			break;
+		}
+		memdelete(E->get().peer);
+		pending.erase(E);
+	}
+}
+
+int UDPServer::get_max_pending_connections() const {
+	return max_pending_connections;
 }
 
 Ref<PacketPeerUDP> UDPServer::take_connection() {
@@ -92,11 +156,20 @@ Ref<PacketPeerUDP> UDPServer::take_connection() {
 		return conn;
 	}
 
-	conn = Ref<PacketPeerUDP>(memnew(PacketPeerUDP));
-	conn->connect_socket(_sock);
-	_sock = Ref<NetSocket>(NetSocket::create());
-	listen(bind_port, bind_address);
-	return conn;
+	Peer peer = pending[0];
+	pending.pop_front();
+	peers.push_back(peer);
+	return peer.peer;
+}
+
+void UDPServer::remove_peer(IP_Address p_ip, int p_port) {
+	Peer peer;
+	peer.ip = p_ip;
+	peer.port = p_port;
+	List<Peer>::Element *E = peers.find(peer);
+	if (E) {
+		peers.erase(E);
+	}
 }
 
 void UDPServer::stop() {
@@ -105,6 +178,19 @@ void UDPServer::stop() {
 	}
 	bind_port = 0;
 	bind_address = IP_Address();
+	List<Peer>::Element *E = peers.front();
+	while (E) {
+		E->get().peer->disconnect_shared_socket();
+		E = E->next();
+	}
+	E = pending.front();
+	while (E) {
+		E->get().peer->disconnect_shared_socket();
+		memdelete(E->get().peer);
+		E = E->next();
+	}
+	peers.clear();
+	pending.clear();
 }
 
 UDPServer::UDPServer() :