Funnel refuse_new_connections to Godot ENet.

2 files changed, 12 insertions, 0 deletions

diff --git a/thirdparty/enet/enet/enet.h b/thirdparty/enet/enet/enet.h
index 3900353c34..d3b3b54b9d 100644
--- a/thirdparty/enet/enet/enet.h
+++ b/thirdparty/enet/enet/enet.h
@@ -587,6 +587,7 @@ extern   void       enet_host_bandwidth_throttle (ENetHost *);
 extern  enet_uint32 enet_host_random_seed (void);
 ENET_API void enet_host_dtls_server_setup (ENetHost *, void *, void *);
 ENET_API void enet_host_dtls_client_setup (ENetHost *, void *, uint8_t, const char *);
+ENET_API void enet_host_refuse_new_connections (ENetHost *, int);
 
 ENET_API int                 enet_peer_send (ENetPeer *, enet_uint8, ENetPacket *);
 ENET_API ENetPacket *        enet_peer_receive (ENetPeer *, enet_uint8 * channelID);
diff --git a/thirdparty/enet/godot.cpp b/thirdparty/enet/godot.cpp
index 9fefa53e77..36b5131f80 100644
--- a/thirdparty/enet/godot.cpp
+++ b/thirdparty/enet/godot.cpp
@@ -51,6 +51,7 @@ public:
 	virtual Error recvfrom(uint8_t *p_buffer, int p_len, int &r_read, IP_Address &r_ip, uint16_t &r_port) = 0;
 	virtual int set_option(ENetSocketOption p_option, int p_value) = 0;
 	virtual void close() = 0;
+	virtual void set_refuse_new_connections(bool p_refuse) { /* Only used by dtls server */ }
 	virtual ~ENetGodotSocket(){};
 };
 
@@ -250,6 +251,10 @@ public:
 		close();
 	}
 
+	void set_refuse_new_connections(bool p_refuse) {
+		udp_server->set_max_pending_connections(p_refuse ? 0 : 16);
+	}
+
 	Error bind(IP_Address p_ip, uint16_t p_port) {
 		return udp_server->listen(p_port, p_ip);
 	}
@@ -269,6 +274,7 @@ public:
 	}
 
 	Error recvfrom(uint8_t *p_buffer, int p_len, int &r_read, IP_Address &r_ip, uint16_t &r_port) {
+		udp_server->poll();
 		// TODO limits? Maybe we can better enforce allowed connections!
 		if (udp_server->is_connection_available()) {
 			Ref<PacketPeerUDP> udp = udp_server->take_connection();
@@ -409,6 +415,11 @@ void enet_host_dtls_client_setup(ENetHost *host, void *p_cert, uint8_t p_verify,
 	memdelete(sock);
 }
 
+void enet_host_refuse_new_connections(ENetHost *host, int p_refuse) {
+	ERR_FAIL_COND(!host->socket);
+	((ENetGodotSocket *)host->socket)->set_refuse_new_connections(p_refuse);
+}
+
 int enet_socket_bind(ENetSocket socket, const ENetAddress *address) {
 
 	IP_Address ip;