Adam Ierymenko 6 years ago
parent
commit
0e627f4f68
1 changed files with 54 additions and 24 deletions
  1. 54 24
      root/root.cpp

+ 54 - 24
root/root.cpp

@@ -62,6 +62,7 @@
 #include <vector>
 #include <iostream>
 #include <unordered_map>
+#include <unordered_set>
 #include <vector>
 #include <atomic>
 #include <mutex>
@@ -80,9 +81,6 @@ struct PeerInfo
 	InetAddress ip4,ip6;
 	int64_t lastReceive;
 
-	std::unordered_map< uint64_t,std::unordered_map< MulticastGroup,int64_t,MulticastGroupHasher > > multicastGroups;
-	Mutex multicastGroups_l;
-
 	AtomicCounter __refCount;
 
 	ZT_ALWAYS_INLINE ~PeerInfo() { Utils::burn(key,sizeof(key)); }
@@ -92,12 +90,13 @@ static Identity self;
 static std::atomic_bool run;
 
 static std::vector< SharedPtr<PeerInfo> > newPeers;
-
+static std::unordered_map< uint64_t,std::unordered_map< MulticastGroup,std::unordered_map< Address,int64_t,AddressHasher >,MulticastGroupHasher > > multicastSubscriptions;
 static std::unordered_map< Identity,SharedPtr<PeerInfo>,IdentityHasher > peersByIdentity;
 static std::unordered_map< Address,std::set< SharedPtr<PeerInfo> >,AddressHasher > peersByVirtAddr;
 static std::unordered_map< InetAddress,std::set< SharedPtr<PeerInfo> >,InetAddressHasher > peersByPhysAddr;
 
 static std::mutex newPeers_l;
+static std::mutex multicastSubscriptions_l;
 static std::mutex peersByIdentity_l;
 static std::mutex peersByVirtAddr_l;
 static std::mutex peersByPhysAddr_l;
@@ -197,37 +196,68 @@ static void handlePacket(const int sock,const InetAddress *const ip,Packet &pkt)
 
 			switch(pkt.verb()) {
 				case Packet::VERB_HELLO: {
-					if (pkt.source() == 0x89e92ceee5) {
-						printf("ME!\n");
-						const uint64_t origId = pkt.packetId();
-						const uint64_t ts = pkt.template at<uint64_t>(ZT_PROTO_VERB_HELLO_IDX_TIMESTAMP);
-						pkt.reset(pkt.source(),self.address(),Packet::VERB_OK);
-						pkt.append((uint8_t)Packet::VERB_HELLO);
-						pkt.append(origId);
-						pkt.append(ts);
-						pkt.append((uint8_t)ZT_PROTO_VERSION);
-						pkt.append((uint8_t)1);
-						pkt.append((uint8_t)9);
-						pkt.append((uint16_t)0);
-						ip->serialize(pkt);
-						pkt.armor(peer->key,true);
-						sendto(sock,pkt.data(),pkt.size(),0,(const struct sockaddr *)ip,(socklen_t)((ip->ss_family == AF_INET) ? sizeof(struct sockaddr_in) : sizeof(struct sockaddr_in6)));
-						printf("%s <- OK(HELLO)" ZT_EOL_S,ip->toString(ipstr));
-					}
+					const uint64_t origId = pkt.packetId();
+					const uint64_t ts = pkt.template at<uint64_t>(ZT_PROTO_VERB_HELLO_IDX_TIMESTAMP);
+					pkt.reset(pkt.source(),self.address(),Packet::VERB_OK);
+					pkt.append((uint8_t)Packet::VERB_HELLO);
+					pkt.append(origId);
+					pkt.append(ts);
+					pkt.append((uint8_t)ZT_PROTO_VERSION);
+					pkt.append((uint8_t)0);
+					pkt.append((uint8_t)0);
+					pkt.append((uint16_t)0);
+					ip->serialize(pkt);
+					pkt.armor(peer->key,true);
+					sendto(sock,pkt.data(),pkt.size(),0,(const struct sockaddr *)ip,(socklen_t)((ip->ss_family == AF_INET) ? sizeof(struct sockaddr_in) : sizeof(struct sockaddr_in6)));
+					//printf("%s <- OK(HELLO)" ZT_EOL_S,ip->toString(ipstr));
 				}	break;
 
 				case Packet::VERB_MULTICAST_LIKE: {
 					printf("LIKE\n");
-					Mutex::Lock l(peer->multicastGroups_l);
+					std::lock_guard<std::mutex> l(multicastSubscriptions_l);
 					for(unsigned int ptr=ZT_PACKET_IDX_PAYLOAD;(ptr+18)<=pkt.size();ptr+=18) {
 						const uint64_t nwid = pkt.template at<uint64_t>(ptr);
 						const MulticastGroup mg(MAC(pkt.field(ptr + 8,6),6),pkt.template at<uint32_t>(ptr + 14));
-						peer->multicastGroups[nwid][mg] = now;
+						multicastSubscriptions[nwid][mg][peer->id.address()] = now;
 						printf("%s subscribes to %s/%.8lx on network %.16llx" ZT_EOL_S,ip->toString(ipstr),mg.mac().toString(tmpstr),(unsigned long)mg.adi(),(unsigned long long)nwid);
 					}
 				}	break;
 
 				case Packet::VERB_MULTICAST_GATHER: {
+					const uint64_t nwid = pkt.template at<uint64_t>(ZT_PROTO_VERB_MULTICAST_GATHER_IDX_NETWORK_ID);
+					const unsigned int flags = pkt[ZT_PROTO_VERB_MULTICAST_GATHER_IDX_FLAGS];
+					const MulticastGroup mg(MAC(pkt.field(ZT_PROTO_VERB_MULTICAST_GATHER_IDX_MAC,6),6),pkt.template at<uint32_t>(ZT_PROTO_VERB_MULTICAST_GATHER_IDX_ADI));
+					unsigned int gatherLimit = pkt.template at<uint32_t>(ZT_PROTO_VERB_MULTICAST_GATHER_IDX_GATHER_LIMIT);
+					if (gatherLimit > 255)
+						gatherLimit = 255;
+
+					const uint64_t origId = pkt.packetId();
+					pkt.reset(pkt.source(),self.address(),Packet::VERB_OK);
+					pkt.append((uint8_t)Packet::VERB_MULTICAST_GATHER);
+					pkt.append(origId);
+					pkt.append(nwid);
+					mg.mac().appendTo(pkt);
+					pkt.append((uint32_t)mg.adi());
+
+					{
+						std::lock_guard<std::mutex> l(multicastSubscriptions_l);
+						auto forNet = multicastSubscriptions.find(nwid);
+						if (forNet != multicastSubscriptions.end()) {
+							auto forGroup = forNet->second.find(mg);
+							if (forGroup != forNet->second.end()) {
+								pkt.append((uint32_t)forGroup->second.size());
+								pkt.append((uint16_t)std::min(std::min((unsigned int)forGroup->second.size(),(unsigned int)65535),gatherLimit));
+								auto g = forGroup->second.begin();
+								unsigned int l = 0;
+								for(;((l<gatherLimit)&&(g!=forGroup->second.end()));++l,++g)
+									g->first.appendTo(pkt);
+								if (l > 0) {
+									sendto(sock,pkt.data(),pkt.size(),0,(const struct sockaddr *)ip,(socklen_t)((ip->ss_family == AF_INET) ? sizeof(struct sockaddr_in) : sizeof(struct sockaddr_in6)));
+									printf("%s gathered %u subscribers to %s/%.8lx on network %.16llx" ZT_EOL_S,ip->toString(ipstr),l,mg.mac().toString(tmpstr),(unsigned long)mg.adi(),(unsigned long long)nwid);
+								}
+							}
+						}
+					}
 				}	break;
 
 				default:
@@ -270,7 +300,7 @@ static void handlePacket(const int sock,const InetAddress *const ip,Packet &pkt)
 
 	for(auto i=toAddrs.begin();i!=toAddrs.end();++i) {
 		printf("%s -> %s for %s" ZT_EOL_S,ip->toString(ipstr),i->toString(ipstr2),pkt.destination().toString(astr));
-		//sendto(sock,pkt.data(),pkt.size(),0,(const struct sockaddr *)&(*i),(socklen_t)((i->ss_family == AF_INET) ? sizeof(struct sockaddr_in) : sizeof(struct sockaddr_in6)));
+		sendto(sock,pkt.data(),pkt.size(),0,(const struct sockaddr *)&(*i),(socklen_t)((i->ss_family == AF_INET) ? sizeof(struct sockaddr_in) : sizeof(struct sockaddr_in6)));
 	}
 }