def test_eq_false(): other = ID("efgh") expected = False actual = ID("abcd").__eq__(other) assert actual == expected
async def test_simple_two_nodes(): node_a = await new_node(transport_opt=[str(LISTEN_MADDR)]) node_b = await new_node(transport_opt=[str(LISTEN_MADDR)]) await node_a.get_network().listen(LISTEN_MADDR) await node_b.get_network().listen(LISTEN_MADDR) supported_protocols = [FLOODSUB_PROTOCOL_ID] topic = "my_topic" data = b"some data" floodsub_a = FloodSub(supported_protocols) pubsub_a = Pubsub(node_a, floodsub_a, ID(b"a" * 32)) floodsub_b = FloodSub(supported_protocols) pubsub_b = Pubsub(node_b, floodsub_b, ID(b"b" * 32)) await connect(node_a, node_b) await asyncio.sleep(0.25) sub_b = await pubsub_b.subscribe(topic) # Sleep to let a know of b's subscription await asyncio.sleep(0.25) await pubsub_a.publish(topic, data) res_b = await sub_b.get() # Check that the msg received by node_b is the same # as the message sent by node_a assert ID(res_b.from_id) == node_a.get_id() assert res_b.data == data assert res_b.topicIDs == [topic] # Success, terminate pending tasks. await cleanup()
def test_eq_true(): random_id_string = "" for _ in range(10): random_id_string += random.choice(ALPHABETS) peer_id = ID(random_id_string.encode()) assert peer_id == base58.b58encode(random_id_string).decode() assert peer_id == random_id_string.encode() assert peer_id == ID(random_id_string.encode())
def test_eq_true(): random_id_string = '' for _ in range(10): random_id_string += random.SystemRandom().choice(ALPHABETS) other = ID(random_id_string) expected = True actual = ID(random_id_string).__eq__(other) assert actual == expected
def create_kad_peerinfo(raw_node_id=None, sender_ip=None, sender_port=None): node_id = ID(raw_node_id) if raw_node_id else ID( digest(random.getrandbits(255))) peer_data = None if sender_ip and sender_port: peer_data = PeerData() #pylint: disable=no-value-for-parameter addr = [Multiaddr("/"+ P_IP +"/" + str(sender_ip) + "/"\ + P_UDP + "/" + str(sender_port))] peer_data.add_addrs(addr) return KadPeerInfo(node_id, peer_data)
async def test_lru_cache_two_nodes(monkeypatch): # two nodes with cache_size of 4 # `node_a` send the following messages to node_b message_indices = [1, 1, 2, 1, 3, 1, 4, 1, 5, 1] # `node_b` should only receive the following expected_received_indices = [1, 2, 3, 4, 5, 1] node_a = await new_node(transport_opt=[str(LISTEN_MADDR)]) node_b = await new_node(transport_opt=[str(LISTEN_MADDR)]) await node_a.get_network().listen(LISTEN_MADDR) await node_b.get_network().listen(LISTEN_MADDR) supported_protocols = SUPPORTED_PROTOCOLS topic = "my_topic" # Mock `get_msg_id` to make us easier to manipulate `msg_id` by `data`. def get_msg_id(msg): # Originally it is `(msg.seqno, msg.from_id)` return (msg.data, msg.from_id) import libp2p.pubsub.pubsub monkeypatch.setattr(libp2p.pubsub.pubsub, "get_msg_id", get_msg_id) # Initialize Pubsub with a cache_size of 4 cache_size = 4 floodsub_a = FloodSub(supported_protocols) pubsub_a = Pubsub(node_a, floodsub_a, ID(b"a" * 32), cache_size) floodsub_b = FloodSub(supported_protocols) pubsub_b = Pubsub(node_b, floodsub_b, ID(b"b" * 32), cache_size) await connect(node_a, node_b) await asyncio.sleep(0.25) sub_b = await pubsub_b.subscribe(topic) await asyncio.sleep(0.25) def _make_testing_data(i: int) -> bytes: num_int_bytes = 4 if i >= 2**(num_int_bytes * 8): raise ValueError("integer is too large to be serialized") return b"data" + i.to_bytes(num_int_bytes, "big") for index in message_indices: await pubsub_a.publish(topic, _make_testing_data(index)) await asyncio.sleep(0.25) for index in expected_received_indices: res_b = await sub_b.get() assert res_b.data == _make_testing_data(index) assert sub_b.empty() # Success, terminate pending tasks. await cleanup()
def create_kad_peerinfo(node_id_bytes=None, sender_ip=None, sender_port=None): node_id = (ID(node_id_bytes) if node_id_bytes else ID( digest(random.getrandbits(255)))) addrs: List[Multiaddr] if sender_ip and sender_port: addrs = [ Multiaddr("/" + P_IP + "/" + str(sender_ip) + "/" + P_UDP + "/" + str(sender_port)) ] else: addrs = [] return KadPeerInfo(node_id, addrs)
async def test_set_and_remove_topic_validator(pubsubs_fsub): is_sync_validator_called = False def sync_validator(peer_id, msg): nonlocal is_sync_validator_called is_sync_validator_called = True is_async_validator_called = False async def async_validator(peer_id, msg): nonlocal is_async_validator_called is_async_validator_called = True topic = "TEST_VALIDATOR" assert topic not in pubsubs_fsub[0].topic_validators # Register sync validator pubsubs_fsub[0].set_topic_validator(topic, sync_validator, False) assert topic in pubsubs_fsub[0].topic_validators topic_validator = pubsubs_fsub[0].topic_validators[topic] assert not topic_validator.is_async # Validate with sync validator topic_validator.validator(peer_id=ID(b"peer"), msg="msg") assert is_sync_validator_called assert not is_async_validator_called # Register with async validator pubsubs_fsub[0].set_topic_validator(topic, async_validator, True) is_sync_validator_called = False assert topic in pubsubs_fsub[0].topic_validators topic_validator = pubsubs_fsub[0].topic_validators[topic] assert topic_validator.is_async # Validate with async validator await topic_validator.validator(peer_id=ID(b"peer"), msg="msg") assert is_async_validator_called assert not is_sync_validator_called # Remove validator pubsubs_fsub[0].remove_topic_validator(topic) assert topic not in pubsubs_fsub[0].topic_validators
async def test_simple_four_nodes(): node_a = KademliaServer() await node_a.listen(5801) node_b = KademliaServer() await node_b.listen(5802) node_c = KademliaServer() await node_c.listen(5803) node_d = KademliaServer() await node_d.listen(5804) node_a_value = await node_b.bootstrap([("127.0.0.1", 5801)]) node_a_kad_peerinfo = node_a_value[0] await node_c.bootstrap([("127.0.0.1", 5802)]) await node_d.bootstrap([("127.0.0.1", 5803)]) await node_b.set(node_a_kad_peerinfo.xor_id, peer_info_to_str(node_a_kad_peerinfo)) router = KadmeliaPeerRouter(node_d) returned_info = await router.find_peer( ID(node_a_kad_peerinfo.peer_id_bytes)) assert returned_info == node_a_kad_peerinfo
async def publish(self, msg_forwarder: ID, pubsub_msg: rpc_pb2.Message) -> None: """Invoked to forward a new message that has been validated.""" self.mcache.put(pubsub_msg) peers_gen = self._get_peers_to_send( pubsub_msg.topicIDs, msg_forwarder=msg_forwarder, origin=ID(pubsub_msg.from_id), ) rpc_msg = rpc_pb2.RPC(publish=[pubsub_msg]) logger.debug("publishing message %s", pubsub_msg) for peer_id in peers_gen: if peer_id not in self.pubsub.peers: continue stream = self.pubsub.peers[peer_id] # FIXME: We should add a `WriteMsg` similar to write delimited messages. # Ref: https://github.com/libp2p/go-libp2p-pubsub/blob/master/comm.go#L107 # TODO: Go use `sendRPC`, which possibly piggybacks gossip/control messages. try: await stream.write(encode_varint_prefixed(rpc_msg.SerializeToString())) except StreamClosed: logger.debug("Fail to publish message to %s: stream closed", peer_id) self.pubsub._handle_dead_peer(peer_id)
async def conn_handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: # Upgrade reader/write to a net_stream and pass \ # to appropriate stream handler (using multiaddr) raw_conn = RawConnection(reader, writer, False) # Per, https://discuss.libp2p.io/t/multistream-security/130, we first secure # the conn and then mux the conn try: # FIXME: This dummy `ID(b"")` for the remote peer is useless. secured_conn = await self.upgrader.upgrade_security( raw_conn, ID(b""), False) except SecurityUpgradeFailure as error: # TODO: Add logging to indicate the failure await raw_conn.close() raise SwarmException( "fail to upgrade the connection to a secured connection" ) from error peer_id = secured_conn.get_remote_peer() try: muxed_conn = await self.upgrader.upgrade_connection( secured_conn, self.generic_protocol_handler, peer_id) except MuxerUpgradeFailure as error: # TODO: Add logging to indicate the failure await secured_conn.close() raise SwarmException( f"fail to upgrade the connection to a muxed connection from {peer_id}" ) from error # Store muxed_conn with peer id self.connections[peer_id] = muxed_conn # Call notifiers since event occurred for notifee in self.notifees: await notifee.connected(self, muxed_conn)
async def add_node(node_id: str) -> None: node = await new_node(transport_opt=[str(LISTEN_MADDR)]) await node.get_network().listen(LISTEN_MADDR) node_map[node_id] = node pubsub_router = router_factory(protocols=obj["supported_protocols"]) pubsub = Pubsub(node, pubsub_router, ID(node_id.encode())) pubsub_map[node_id] = pubsub
async def test_node_dial_peer(nodes): # Test: Exception raised when dialing a wrong addr with pytest.raises(ConnectionRefusedError): await nodes[0].dial_peer( nodes[1].listen_ip, get_open_port(), ID("123"), ) # Test: 0 <-> 1 await nodes[0].dial_peer( nodes[1].listen_ip, nodes[1].listen_port, nodes[1].peer_id, ) assert nodes[0].peer_id in nodes[1].host.get_network().connections assert nodes[1].peer_id in nodes[0].host.get_network().connections # Test: Second dial to a connected peer does not open a new connection original_conn = nodes[1].host.get_network().connections[nodes[0].peer_id] await nodes[0].dial_peer( nodes[1].listen_ip, nodes[1].listen_port, nodes[1].peer_id, ) assert nodes[1].host.get_network().connections[ nodes[0].peer_id] is original_conn # Test: 0 <-> 1 <-> 2 await nodes[2].dial_peer( nodes[1].listen_ip, nodes[1].listen_port, nodes[1].peer_id, ) assert nodes[1].peer_id in nodes[2].host.get_network().connections assert nodes[2].peer_id in nodes[1].host.get_network().connections assert len(nodes[1].host.get_network().connections) == 2
def test_init_(): random_id_string = '' for _ in range(10): random_id_string += random.SystemRandom().choice(ALPHABETS) peer_id = ID(random_id_string) #pylint: disable=protected-access assert peer_id._id_str == random_id_string
def test_handle_subscription(pubsubs_fsub): assert len(pubsubs_fsub[0].peer_topics) == 0 sub_msg_0 = rpc_pb2.RPC.SubOpts(subscribe=True, topicid=TESTING_TOPIC) peer_ids = [ID(b"\x12\x20" + i.to_bytes(32, "big")) for i in range(2)] # Test: One peer is subscribed pubsubs_fsub[0].handle_subscription(peer_ids[0], sub_msg_0) assert ( len(pubsubs_fsub[0].peer_topics) == 1 and TESTING_TOPIC in pubsubs_fsub[0].peer_topics ) assert len(pubsubs_fsub[0].peer_topics[TESTING_TOPIC]) == 1 assert peer_ids[0] in pubsubs_fsub[0].peer_topics[TESTING_TOPIC] # Test: Another peer is subscribed pubsubs_fsub[0].handle_subscription(peer_ids[1], sub_msg_0) assert len(pubsubs_fsub[0].peer_topics) == 1 assert len(pubsubs_fsub[0].peer_topics[TESTING_TOPIC]) == 2 assert peer_ids[1] in pubsubs_fsub[0].peer_topics[TESTING_TOPIC] # Test: Subscribe to another topic another_topic = "ANOTHER_TOPIC" sub_msg_1 = rpc_pb2.RPC.SubOpts(subscribe=True, topicid=another_topic) pubsubs_fsub[0].handle_subscription(peer_ids[0], sub_msg_1) assert len(pubsubs_fsub[0].peer_topics) == 2 assert another_topic in pubsubs_fsub[0].peer_topics assert peer_ids[0] in pubsubs_fsub[0].peer_topics[another_topic] # Test: unsubscribe unsub_msg = rpc_pb2.RPC.SubOpts(subscribe=False, topicid=TESTING_TOPIC) pubsubs_fsub[0].handle_subscription(peer_ids[0], unsub_msg) assert peer_ids[0] not in pubsubs_fsub[0].peer_topics[TESTING_TOPIC]
async def publish(self, msg_forwarder: ID, pubsub_msg: rpc_pb2.Message) -> None: """ Invoked to forward a new message that has been validated. This is where the "flooding" part of floodsub happens With flooding, routing is almost trivial: for each incoming message, forward to all known peers in the topic. There is a bit of logic, as the router maintains a timed cache of previous messages, so that seen messages are not further forwarded. It also never forwards a message back to the source or the peer that forwarded the message. :param msg_forwarder: peer ID of the peer who forwards the message to us :param pubsub_msg: pubsub message in protobuf. """ peers_gen = self._get_peers_to_send( pubsub_msg.topicIDs, msg_forwarder=msg_forwarder, origin=ID(pubsub_msg.from_id), ) rpc_msg = rpc_pb2.RPC(publish=[pubsub_msg], ) for peer_id in peers_gen: stream = self.pubsub.peers[str(peer_id)] # FIXME: We should add a `WriteMsg` similar to write delimited messages. # Ref: https://github.com/libp2p/go-libp2p-pubsub/blob/master/comm.go#L107 await stream.write(rpc_msg.SerializeToString())
def test_id_b58_encode(): random_id_string = '' for _ in range(10): random_id_string += random.SystemRandom().choice(ALPHABETS) expected = base58.b58encode(random_id_string).decode() actual = id_b58_encode(ID(random_id_string)) assert actual == expected
def test_id_to_base58(): random_id_string = "" for _ in range(10): random_id_string += random.choice(ALPHABETS) expected = base58.b58encode(random_id_string).decode() actual = ID(random_id_string.encode()).to_base58() assert actual == expected
def test_id_from_base58(): random_id_string = "" for _ in range(10): random_id_string += random.choice(ALPHABETS) expected = ID(base58.b58decode(random_id_string)) actual = ID.from_base58(random_id_string.encode()) assert actual == expected
def test_hash(): random_id_string = '' for _ in range(10): random_id_string += random.SystemRandom().choice(ALPHABETS) expected = hash(random_id_string) actual = ID(random_id_string).__hash__() assert actual == expected
def test_str_less_than_10(): random_id_string = '' for _ in range(5): random_id_string += random.SystemRandom().choice(ALPHABETS) pid = base58.b58encode(random_id_string).decode() expected = "<peer.ID " + pid + ">" actual = ID(random_id_string).__str__() assert actual == expected
def test_pretty(): random_id_string = '' for _ in range(10): random_id_string += random.SystemRandom().choice(ALPHABETS) peer_id = ID(random_id_string) actual = peer_id.pretty() expected = base58.b58encode(random_id_string).decode() assert actual == expected
def test_str_more_than_10(): random_id_string = '' for _ in range(10): random_id_string += random.SystemRandom().choice(ALPHABETS) pid = base58.b58encode(random_id_string).decode() expected = pid actual = ID(random_id_string).__str__() assert actual == expected
def test_str_less_than_10(): random_id_string = "" for _ in range(5): random_id_string += random.choice(ALPHABETS) peer_id = base58.b58encode(random_id_string).decode() expected = peer_id actual = ID(random_id_string.encode()).__str__() assert actual == expected
def test_id_from_public_key(): bits_list = [1024, 1280, 1536, 1536, 2048] key = RSA.generate(random.choice(bits_list)) key_bin = key.exportKey("DER") algo = multihash.Func.sha2_256 mh_digest = multihash.digest(key_bin, algo) expected = ID(mh_digest.encode()) actual = id_from_public_key(key) assert actual == expected
def test_init_(): random_addrs = [random.randint(0, 255) for r in range(4)] random_id_string = "" for _ in range(10): random_id_string += random.SystemRandom().choice(ALPHABETS) peer_id = ID(random_id_string.encode()) peer_info = PeerInfo(peer_id, random_addrs) assert peer_info.peer_id == peer_id assert peer_info.addrs == random_addrs
class PeerFactory(factory.Factory): class Meta: model = Peer _id = factory.Sequence(lambda n: ID(f'peer{n}')) node = factory.SubFactory(NodeFactory) head_fork_version = None finalized_root = ZERO_HASH32 finalized_epoch = 0 head_root = ZERO_HASH32 head_slot = 0
async def test_message_all_peers(pubsubs_fsub, monkeypatch): peer_ids = [ID(b"\x12\x20" + i.to_bytes(32, "big")) for i in range(10)] mock_peers = {peer_id: FakeNetStream() for peer_id in peer_ids} monkeypatch.setattr(pubsubs_fsub[0], "peers", mock_peers) empty_rpc = rpc_pb2.RPC() empty_rpc_bytes = empty_rpc.SerializeToString() empty_rpc_bytes_len_prefixed = encode_varint_prefixed(empty_rpc_bytes) await pubsubs_fsub[0].message_all_peers(empty_rpc_bytes) for stream in mock_peers.values(): assert (await stream.read()) == empty_rpc_bytes_len_prefixed
def test_str_more_than_10(): random_id_string = '' for _ in range(10): random_id_string += random.SystemRandom().choice(ALPHABETS) pid = base58.b58encode(random_id_string).decode() part_1, part_2 = pid[:2], pid[len(pid) - 6:] expected = "<peer.ID " + part_1 + "*" + part_2 + ">" actual = ID(random_id_string).__str__() assert actual == expected
async def test_get_msg_validators(pubsubs_fsub): times_sync_validator_called = 0 def sync_validator(peer_id, msg): nonlocal times_sync_validator_called times_sync_validator_called += 1 times_async_validator_called = 0 async def async_validator(peer_id, msg): nonlocal times_async_validator_called times_async_validator_called += 1 topic_1 = "TEST_VALIDATOR_1" topic_2 = "TEST_VALIDATOR_2" topic_3 = "TEST_VALIDATOR_3" # Register sync validator for topic 1 and 2 pubsubs_fsub[0].set_topic_validator(topic_1, sync_validator, False) pubsubs_fsub[0].set_topic_validator(topic_2, sync_validator, False) # Register async validator for topic 3 pubsubs_fsub[0].set_topic_validator(topic_3, async_validator, True) msg = make_pubsub_msg( origin_id=pubsubs_fsub[0].my_id, topic_ids=[topic_1, topic_2, topic_3], data=b"1234", seqno=b"\x00" * 8, ) topic_validators = pubsubs_fsub[0].get_msg_validators(msg) for topic_validator in topic_validators: if topic_validator.is_async: await topic_validator.validator(peer_id=ID(b"peer"), msg="msg") else: topic_validator.validator(peer_id=ID(b"peer"), msg="msg") assert times_sync_validator_called == 2 assert times_async_validator_called == 1