Multicast support in NetSocket/PacketPeerUDP

6 files changed, 121 insertions, 0 deletions

diff --git a/core/io/net_socket.h b/core/io/net_socket.h
index 94e7ef6f75..3bc1369487 100644
--- a/core/io/net_socket.h
+++ b/core/io/net_socket.h
@@ -74,6 +74,8 @@ public:
 	virtual void set_ipv6_only_enabled(bool p_enabled) = 0;
 	virtual void set_tcp_no_delay_enabled(bool p_enabled) = 0;
 	virtual void set_reuse_address_enabled(bool p_enabled) = 0;
+	virtual Error join_multicast_group(const IP_Address &p_multi_address, String p_if_name) = 0;
+	virtual Error leave_multicast_group(const IP_Address &p_multi_address, String p_if_name) = 0;
 };
 
 #endif // NET_SOCKET_H
diff --git a/core/io/packet_peer_udp.cpp b/core/io/packet_peer_udp.cpp
index 5912b8df94..7e9471c053 100644
--- a/core/io/packet_peer_udp.cpp
+++ b/core/io/packet_peer_udp.cpp
@@ -37,6 +37,27 @@ void PacketPeerUDP::set_blocking_mode(bool p_enable) {
 	blocking = p_enable;
 }
 
+Error PacketPeerUDP::join_multicast_group(IP_Address p_multi_address, String p_if_name) {
+
+	ERR_FAIL_COND_V(!_sock.is_valid(), ERR_UNAVAILABLE);
+	ERR_FAIL_COND_V(!p_multi_address.is_valid(), ERR_INVALID_PARAMETER);
+
+	if (!_sock->is_open()) {
+		IP::Type ip_type = p_multi_address.is_ipv4() ? IP::TYPE_IPV4 : IP::TYPE_IPV6;
+		Error err = _sock->open(NetSocket::TYPE_UDP, ip_type);
+		ERR_FAIL_COND_V(err != OK, err);
+		_sock->set_blocking_enabled(false);
+	}
+	return _sock->join_multicast_group(p_multi_address, p_if_name);
+}
+
+Error PacketPeerUDP::leave_multicast_group(IP_Address p_multi_address, String p_if_name) {
+
+	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);
+}
+
 String PacketPeerUDP::_get_packet_ip() const {
 
 	return get_packet_address();
@@ -237,6 +258,8 @@ void PacketPeerUDP::_bind_methods() {
 	ClassDB::bind_method(D_METHOD("get_packet_ip"), &PacketPeerUDP::_get_packet_ip);
 	ClassDB::bind_method(D_METHOD("get_packet_port"), &PacketPeerUDP::get_packet_port);
 	ClassDB::bind_method(D_METHOD("set_dest_address", "host", "port"), &PacketPeerUDP::_set_dest_address);
+	ClassDB::bind_method(D_METHOD("join_multicast_group", "multicast_address", "interface_name"), &PacketPeerUDP::join_multicast_group);
+	ClassDB::bind_method(D_METHOD("leave_multicast_group", "multicast_address", "interface_name"), &PacketPeerUDP::leave_multicast_group);
 }
 
 PacketPeerUDP::PacketPeerUDP() :
diff --git a/core/io/packet_peer_udp.h b/core/io/packet_peer_udp.h
index 0593137604..068bd5cd5a 100644
--- a/core/io/packet_peer_udp.h
+++ b/core/io/packet_peer_udp.h
@@ -77,6 +77,8 @@ public:
 	Error get_packet(const uint8_t **r_buffer, int &r_buffer_size);
 	int get_available_packet_count() const;
 	int get_max_packet_size() const;
+	Error join_multicast_group(IP_Address p_multi_address, String p_if_name);
+	Error leave_multicast_group(IP_Address p_multi_address, String p_if_name);
 
 	PacketPeerUDP();
 	~PacketPeerUDP();
diff --git a/doc/classes/PacketPeerUDP.xml b/doc/classes/PacketPeerUDP.xml
index 9843c16108..376818fb86 100644
--- a/doc/classes/PacketPeerUDP.xml
+++ b/doc/classes/PacketPeerUDP.xml
@@ -37,6 +37,29 @@
 				Returns whether this [PacketPeerUDP] is listening.
 			</description>
 		</method>
+		<method name="join_multicast_group">
+			<return type="int" enum="Error">
+			</return>
+			<argument index="0" name="multicast_address" type="String">
+			</argument>
+			<argument index="1" name="interface_name" type="String">
+			</argument>
+			<description>
+				Join the multicast group specified by [code]multicast_address[/code] using the interface identified by [code]interface_name[/code].
+				You can join the same multicast group with multiple interfaces. Use [method IP.get_local_interfaces] to know which are available.
+			</description>
+		</method>
+		<method name="leave_multicast_group">
+			<return type="int" enum="Error">
+			</return>
+			<argument index="0" name="multicast_address" type="String">
+			</argument>
+			<argument index="1" name="interface_name" type="String">
+			</argument>
+			<description>
+				Remove the interface identified by [code]interface_name[/code] from the multicast group specified by [code]multicast_address[/code].
+			</description>
+		</method>
 		<method name="listen">
 			<return type="int" enum="Error">
 			</return>
diff --git a/drivers/unix/net_socket_posix.cpp b/drivers/unix/net_socket_posix.cpp
index 75d51c8503..fec101d969 100644
--- a/drivers/unix/net_socket_posix.cpp
+++ b/drivers/unix/net_socket_posix.cpp
@@ -59,6 +59,14 @@
 #define MSG_NOSIGNAL SO_NOSIGPIPE
 #endif
 
+// BSD calls this flag IPV6_JOIN_GROUP
+#if !defined(IPV6_ADD_MEMBERSHIP) && defined(IPV6_JOIN_GROUP)
+#define IPV6_ADD_MEMBERSHIP IPV6_JOIN_GROUP
+#endif
+#if !defined(IPV6_DROP_MEMBERSHIP) && defined(IPV6_LEAVE_GROUP)
+#define IPV6_DROP_MEMBERSHIP IPV6_LEAVE_GROUP
+#endif
+
 // Some custom defines to minimize ifdefs
 #define SOCK_EMPTY -1
 #define SOCK_BUF(x) x
@@ -227,6 +235,58 @@ bool NetSocketPosix::_can_use_ip(const IP_Address p_ip, const bool p_for_bind) c
 	return true;
 }
 
+_FORCE_INLINE_ Error NetSocketPosix::_change_multicast_group(IP_Address p_ip, String p_if_name, bool p_add) {
+
+	ERR_FAIL_COND_V(!is_open(), ERR_UNCONFIGURED);
+	ERR_FAIL_COND_V(!_can_use_ip(p_ip, false), ERR_INVALID_PARAMETER);
+
+	// Need to force level and af_family to IP(v4) when using dual stacking and provided multicast group is IPv4
+	IP::Type type = _ip_type == IP::TYPE_ANY && p_ip.is_ipv4() ? IP::TYPE_IPV4 : _ip_type;
+	// This needs to be the proper level for the multicast group, no matter if the socket is dual stacking.
+	int level = type == IP::TYPE_IPV4 ? IPPROTO_IP : IPPROTO_IPV6;
+	int ret = -1;
+
+	IP_Address if_ip;
+	uint32_t if_v6id = 0;
+	Map<String, IP::Interface_Info> if_info;
+	IP::get_singleton()->get_local_interfaces(&if_info);
+	for (Map<String, IP::Interface_Info>::Element *E = if_info.front(); E; E = E->next()) {
+		IP::Interface_Info &c = E->get();
+		if (c.name != p_if_name)
+			continue;
+
+		if_v6id = (uint32_t)c.index.to_int64();
+		if (type == IP::TYPE_IPV6)
+			break; // IPv6 uses index.
+
+		for (List<IP_Address>::Element *F = c.ip_addresses.front(); F; F = F->next()) {
+			if (!F->get().is_ipv4())
+				continue; // Wrong IP type
+			if_ip = F->get();
+			break;
+		}
+		break;
+	}
+
+	if (level == IPPROTO_IP) {
+		ERR_FAIL_COND_V(!if_ip.is_valid(), ERR_INVALID_PARAMETER);
+		struct ip_mreq greq;
+		int sock_opt = p_add ? IP_ADD_MEMBERSHIP : IP_DROP_MEMBERSHIP;
+		copymem(&greq.imr_multiaddr, p_ip.get_ipv4(), 4);
+		copymem(&greq.imr_interface, if_ip.get_ipv4(), 4);
+		ret = setsockopt(_sock, level, sock_opt, (const char *)&greq, sizeof(greq));
+	} else {
+		struct ipv6_mreq greq;
+		int sock_opt = p_add ? IPV6_ADD_MEMBERSHIP : IPV6_DROP_MEMBERSHIP;
+		copymem(&greq.ipv6mr_multiaddr, p_ip.get_ipv6(), 16);
+		greq.ipv6mr_interface = if_v6id;
+		ret = setsockopt(_sock, level, sock_opt, (const char *)&greq, sizeof(greq));
+	}
+	ERR_FAIL_COND_V(ret != 0, FAILED);
+
+	return OK;
+}
+
 void NetSocketPosix::_set_socket(SOCKET_TYPE p_sock, IP::Type p_ip_type, bool p_is_stream) {
 	_sock = p_sock;
 	_ip_type = p_ip_type;
@@ -625,3 +685,11 @@ Ref<NetSocket> NetSocketPosix::accept(IP_Address &r_ip, uint16_t &r_port) {
 	ns->set_blocking_enabled(false);
 	return Ref<NetSocket>(ns);
 }
+
+Error NetSocketPosix::join_multicast_group(const IP_Address &p_ip, String p_if_name) {
+	return _change_multicast_group(p_ip, p_if_name, true);
+}
+
+Error NetSocketPosix::leave_multicast_group(const IP_Address &p_ip, String p_if_name) {
+	return _change_multicast_group(p_ip, p_if_name, false);
+}
diff --git a/drivers/unix/net_socket_posix.h b/drivers/unix/net_socket_posix.h
index b7fb3fdc94..cf0412311d 100644
--- a/drivers/unix/net_socket_posix.h
+++ b/drivers/unix/net_socket_posix.h
@@ -60,6 +60,7 @@ private:
 
 	NetError _get_socket_error();
 	void _set_socket(SOCKET_TYPE p_sock, IP::Type p_ip_type, bool p_is_stream);
+	_FORCE_INLINE_ Error _change_multicast_group(IP_Address p_ip, String p_if_name, bool p_add);
 
 protected:
 	static NetSocket *_create_func();
@@ -93,6 +94,8 @@ public:
 	virtual void set_tcp_no_delay_enabled(bool p_enabled);
 	virtual void set_reuse_address_enabled(bool p_enabled);
 	virtual void set_reuse_port_enabled(bool p_enabled);
+	virtual Error join_multicast_group(const IP_Address &p_multi_address, String p_if_name);
+	virtual Error leave_multicast_group(const IP_Address &p_multi_address, String p_if_name);
 
 	NetSocketPosix();
 	~NetSocketPosix();