Merge pull request #18145 from Faless/get_peer_ip_port

Implement get_peer_[address|port] in ENet/WSServer

11 files changed, 121 insertions, 5 deletions

diff --git a/modules/enet/networked_multiplayer_enet.cpp b/modules/enet/networked_multiplayer_enet.cpp
index 0f62714440..28b19224fb 100644
--- a/modules/enet/networked_multiplayer_enet.cpp
+++ b/modules/enet/networked_multiplayer_enet.cpp
@@ -674,7 +674,35 @@ void NetworkedMultiplayerENet::_setup_compressor() {
 
 void NetworkedMultiplayerENet::enet_compressor_destroy(void *context) {
 
-	//do none
+	// Nothing to do
+}
+
+IP_Address NetworkedMultiplayerENet::get_peer_address(int p_peer_id) const {
+
+	ERR_FAIL_COND_V(!peer_map.has(p_peer_id), IP_Address());
+	ERR_FAIL_COND_V(!is_server() && p_peer_id != 1, IP_Address());
+	ERR_FAIL_COND_V(peer_map[p_peer_id] == NULL, IP_Address());
+
+	IP_Address out;
+#ifdef GODOT_ENET
+	out.set_ipv6((uint8_t *)&(peer_map[p_peer_id]->address.host));
+#else
+	out.set_ipv4((uint8_t *)&(peer_map[p_peer_id]->address.host));
+#endif
+
+	return out;
+}
+
+int NetworkedMultiplayerENet::get_peer_port(int p_peer_id) const {
+
+	ERR_FAIL_COND_V(!peer_map.has(p_peer_id), 0);
+	ERR_FAIL_COND_V(!is_server() && p_peer_id != 1, 0);
+	ERR_FAIL_COND_V(peer_map[p_peer_id] == NULL, 0);
+#ifdef GODOT_ENET
+	return peer_map[p_peer_id]->address.port;
+#else
+	return peer_map[p_peer_id]->address.port;
+#endif
 }
 
 void NetworkedMultiplayerENet::_bind_methods() {
@@ -686,6 +714,8 @@ void NetworkedMultiplayerENet::_bind_methods() {
 	ClassDB::bind_method(D_METHOD("set_compression_mode", "mode"), &NetworkedMultiplayerENet::set_compression_mode);
 	ClassDB::bind_method(D_METHOD("get_compression_mode"), &NetworkedMultiplayerENet::get_compression_mode);
 	ClassDB::bind_method(D_METHOD("set_bind_ip", "ip"), &NetworkedMultiplayerENet::set_bind_ip);
+	ClassDB::bind_method(D_METHOD("get_peer_address"), &NetworkedMultiplayerENet::get_peer_address);
+	ClassDB::bind_method(D_METHOD("get_peer_port"), &NetworkedMultiplayerENet::get_peer_port);
 
 	ADD_PROPERTY(PropertyInfo(Variant::INT, "compression_mode", PROPERTY_HINT_ENUM, "None,Range Coder,FastLZ,ZLib,ZStd"), "set_compression_mode", "get_compression_mode");
 
diff --git a/modules/enet/networked_multiplayer_enet.h b/modules/enet/networked_multiplayer_enet.h
index f93f24fd49..0e8dd67160 100644
--- a/modules/enet/networked_multiplayer_enet.h
+++ b/modules/enet/networked_multiplayer_enet.h
@@ -115,6 +115,9 @@ public:
 
 	virtual int get_packet_peer() const;
 
+	virtual IP_Address get_peer_address(int p_peer_id) const;
+	virtual int get_peer_port(int p_peer_id) const;
+
 	Error create_server(int p_port, int p_max_clients = 32, int p_in_bandwidth = 0, int p_out_bandwidth = 0);
 	Error create_client(const IP_Address &p_ip, int p_port, int p_in_bandwidth = 0, int p_out_bandwidth = 0);
 
diff --git a/modules/websocket/emws_peer.cpp b/modules/websocket/emws_peer.cpp
index d3330d683c..e0b987b4d7 100644
--- a/modules/websocket/emws_peer.cpp
+++ b/modules/websocket/emws_peer.cpp
@@ -148,12 +148,14 @@ void EMWSPeer::close() {
 
 IP_Address EMWSPeer::get_connected_host() const {
 
-	return IP_Address();
+	ERR_EXPLAIN("Not supported in HTML5 export");
+	ERR_FAIL_V(IP_Address());
 };
 
 uint16_t EMWSPeer::get_connected_port() const {
 
-	return 1025;
+	ERR_EXPLAIN("Not supported in HTML5 export");
+	ERR_FAIL_V(0);
 };
 
 EMWSPeer::EMWSPeer() {
diff --git a/modules/websocket/emws_server.cpp b/modules/websocket/emws_server.cpp
index c9ddae0c8c..3eb93e4152 100644
--- a/modules/websocket/emws_server.cpp
+++ b/modules/websocket/emws_server.cpp
@@ -58,6 +58,19 @@ PoolVector<String> EMWSServer::get_protocols() const {
 	return out;
 }
 
+IP_Address EMWSServer::get_peer_address(int p_peer_id) const {
+
+	return IP_Address();
+}
+
+int EMWSServer::get_peer_port(int p_peer_id) const {
+
+	return 0;
+}
+
+void EMWSServer::disconnect_peer(int p_peer_id) {
+}
+
 EMWSServer::EMWSServer() {
 }
 
diff --git a/modules/websocket/emws_server.h b/modules/websocket/emws_server.h
index aa089ea40d..9ec4ce72c8 100644
--- a/modules/websocket/emws_server.h
+++ b/modules/websocket/emws_server.h
@@ -46,6 +46,9 @@ public:
 	bool is_listening() const;
 	bool has_peer(int p_id) const;
 	Ref<WebSocketPeer> get_peer(int p_id) const;
+	IP_Address get_peer_address(int p_peer_id) const;
+	int get_peer_port(int p_peer_id) const;
+	void disconnect_peer(int p_peer_id);
 	virtual void poll();
 	virtual PoolVector<String> get_protocols() const;
 
diff --git a/modules/websocket/lws_peer.cpp b/modules/websocket/lws_peer.cpp
index ba45d7688f..8a064fb5a4 100644
--- a/modules/websocket/lws_peer.cpp
+++ b/modules/websocket/lws_peer.cpp
@@ -32,6 +32,13 @@
 #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
+#if !defined(WINDOWS_ENABLED)
+#include <netinet/in.h>
+#endif
+
+#include "drivers/unix/socket_helpers.h"
+
 void LWSPeer::set_wsi(struct lws *p_wsi) {
 	wsi = p_wsi;
 };
@@ -178,12 +185,40 @@ void LWSPeer::close() {
 
 IP_Address LWSPeer::get_connected_host() const {
 
-	return IP_Address();
+	ERR_FAIL_COND_V(!is_connected_to_host(), IP_Address());
+
+	IP_Address ip;
+	int port = 0;
+
+	socklen_t len;
+	struct sockaddr_storage addr;
+	int fd = lws_get_socket_fd(wsi);
+
+	int ret = getpeername(fd, (struct sockaddr *)&addr, &len);
+	ERR_FAIL_COND_V(ret != 0, IP_Address());
+
+	_set_ip_addr_port(ip, port, &addr);
+
+	return ip;
 };
 
 uint16_t LWSPeer::get_connected_port() const {
 
-	return 1025;
+	ERR_FAIL_COND_V(!is_connected_to_host(), 0);
+
+	IP_Address ip;
+	int port = 0;
+
+	socklen_t len;
+	struct sockaddr_storage addr;
+	int fd = lws_get_socket_fd(wsi);
+
+	int ret = getpeername(fd, (struct sockaddr *)&addr, &len);
+	ERR_FAIL_COND_V(ret != 0, 0);
+
+	_set_ip_addr_port(ip, port, &addr);
+
+	return port;
 };
 
 LWSPeer::LWSPeer() {
diff --git a/modules/websocket/lws_server.cpp b/modules/websocket/lws_server.cpp
index 94fe4231ae..8d13dc7a98 100644
--- a/modules/websocket/lws_server.cpp
+++ b/modules/websocket/lws_server.cpp
@@ -164,6 +164,24 @@ Ref<WebSocketPeer> LWSServer::get_peer(int p_id) const {
 	return _peer_map[p_id];
 }
 
+IP_Address LWSServer::get_peer_address(int p_peer_id) const {
+	ERR_FAIL_COND_V(!has_peer(p_peer_id), IP_Address());
+
+	return _peer_map[p_peer_id]->get_connected_host();
+}
+
+int LWSServer::get_peer_port(int p_peer_id) const {
+	ERR_FAIL_COND_V(!has_peer(p_peer_id), 0);
+
+	return _peer_map[p_peer_id]->get_connected_port();
+}
+
+void LWSServer::disconnect_peer(int p_peer_id) {
+	ERR_FAIL_COND(!has_peer(p_peer_id));
+
+	get_peer(p_peer_id)->close();
+}
+
 LWSServer::LWSServer() {
 	context = NULL;
 	_lws_ref = NULL;
diff --git a/modules/websocket/lws_server.h b/modules/websocket/lws_server.h
index de8f59e5ae..9e3fb9b775 100644
--- a/modules/websocket/lws_server.h
+++ b/modules/websocket/lws_server.h
@@ -52,6 +52,9 @@ public:
 	bool is_listening() const;
 	bool has_peer(int p_id) const;
 	Ref<WebSocketPeer> get_peer(int p_id) const;
+	IP_Address get_peer_address(int p_peer_id) const;
+	int get_peer_port(int p_peer_id) const;
+	void disconnect_peer(int p_peer_id);
 	virtual void poll() { _lws_poll(); }
 
 	LWSServer();
diff --git a/modules/websocket/websocket_peer.cpp b/modules/websocket/websocket_peer.cpp
index 6324047846..61f783e377 100644
--- a/modules/websocket/websocket_peer.cpp
+++ b/modules/websocket/websocket_peer.cpp
@@ -43,6 +43,8 @@ void WebSocketPeer::_bind_methods() {
 	ClassDB::bind_method(D_METHOD("is_connected_to_host"), &WebSocketPeer::is_connected_to_host);
 	ClassDB::bind_method(D_METHOD("was_string_packet"), &WebSocketPeer::was_string_packet);
 	ClassDB::bind_method(D_METHOD("close"), &WebSocketPeer::close);
+	ClassDB::bind_method(D_METHOD("get_connected_host"), &WebSocketPeer::get_connected_host);
+	ClassDB::bind_method(D_METHOD("get_connected_port"), &WebSocketPeer::get_connected_port);
 
 	BIND_ENUM_CONSTANT(WRITE_MODE_TEXT);
 	BIND_ENUM_CONSTANT(WRITE_MODE_BINARY);
diff --git a/modules/websocket/websocket_server.cpp b/modules/websocket/websocket_server.cpp
index 5746f61e10..2693b26e47 100644
--- a/modules/websocket/websocket_server.cpp
+++ b/modules/websocket/websocket_server.cpp
@@ -44,6 +44,9 @@ void WebSocketServer::_bind_methods() {
 	ClassDB::bind_method(D_METHOD("listen", "port", "protocols", "gd_mp_api"), &WebSocketServer::listen, DEFVAL(PoolVector<String>()), DEFVAL(false));
 	ClassDB::bind_method(D_METHOD("stop"), &WebSocketServer::stop);
 	ClassDB::bind_method(D_METHOD("has_peer", "id"), &WebSocketServer::has_peer);
+	ClassDB::bind_method(D_METHOD("get_peer_address"), &WebSocketServer::get_peer_address);
+	ClassDB::bind_method(D_METHOD("get_peer_port"), &WebSocketServer::get_peer_port);
+	ClassDB::bind_method(D_METHOD("disconnect_peer"), &WebSocketServer::disconnect_peer);
 
 	ADD_SIGNAL(MethodInfo("client_disconnected", PropertyInfo(Variant::INT, "id")));
 	ADD_SIGNAL(MethodInfo("client_connected", PropertyInfo(Variant::INT, "id"), PropertyInfo(Variant::STRING, "protocol")));
diff --git a/modules/websocket/websocket_server.h b/modules/websocket/websocket_server.h
index 360ff9e6d4..64935f8a58 100644
--- a/modules/websocket/websocket_server.h
+++ b/modules/websocket/websocket_server.h
@@ -52,6 +52,10 @@ public:
 	virtual bool is_server() const;
 	ConnectionStatus get_connection_status() const;
 
+	virtual IP_Address get_peer_address(int p_peer_id) const = 0;
+	virtual int get_peer_port(int p_peer_id) const = 0;
+	virtual void disconnect_peer(int p_peer_id) = 0;
+
 	void _on_peer_packet(int32_t p_peer_id);
 	void _on_connect(int32_t p_peer_id, String p_protocol);
 	void _on_disconnect(int32_t p_peer_id);