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() :