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_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_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_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
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(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)
def test_id_from_public_key(): key_pair = create_new_key_pair() public_key = key_pair.public_key key_bin = public_key.serialize() algo = multihash.Func.sha2_256 mh_digest = multihash.digest(key_bin, algo) expected = ID(mh_digest.encode()) actual = ID.from_pubkey(public_key) assert actual == expected
async def run_handshake(self) -> None: """Raise `HandshakeFailure` when handshake failed.""" msg = make_exchange_message(self.local_private_key.get_public_key()) msg_bytes = msg.SerializeToString() encoded_msg_bytes = encode_fixedint_prefixed(msg_bytes) try: await self.write(encoded_msg_bytes) except RawConnError: raise HandshakeFailure("connection closed") try: remote_msg_bytes = await read_fixedint_prefixed(self.conn) except RawConnError: raise HandshakeFailure("connection closed") remote_msg = plaintext_pb2.Exchange() remote_msg.ParseFromString(remote_msg_bytes) received_peer_id = ID(remote_msg.id) # Verify if the receive `ID` matches the one we originally initialize the session. # We only need to check it when we are the initiator, because only in that condition # we possibly knows the `ID` of the remote. if self.is_initiator and self.remote_peer_id != received_peer_id: raise HandshakeFailure( "remote peer sent unexpected peer ID. " f"expected={self.remote_peer_id} received={received_peer_id}" ) # Verify if the given `pubkey` matches the given `peer_id` try: received_pubkey = deserialize_public_key( remote_msg.pubkey.SerializeToString() ) except ValueError: raise HandshakeFailure( f"unknown `key_type` of remote_msg.pubkey={remote_msg.pubkey}" ) except MissingDeserializerError as error: raise HandshakeFailure(error) peer_id_from_received_pubkey = ID.from_pubkey(received_pubkey) if peer_id_from_received_pubkey != received_peer_id: raise HandshakeFailure( "peer id and pubkey from the remote mismatch: " f"received_peer_id={received_peer_id}, remote_pubkey={received_pubkey}, " f"peer_id_from_received_pubkey={peer_id_from_received_pubkey}" ) # Nothing is wrong. Store the `pubkey` and `peer_id` in the session. self.remote_permanent_pubkey = received_pubkey # Only need to set peer's id when we don't know it before, # i.e. we are not the connection initiator. if not self.is_initiator: self.remote_peer_id = received_peer_id
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 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)
def initialize_default_kademlia_router( ksize: int = 20, alpha: int = 3, id_opt: ID = None, storage: IStorage = None) -> KadmeliaPeerRouter: """ initialize kadmelia router when no kademlia router is passed in :param ksize: The k parameter from the paper :param alpha: The alpha parameter from the paper :param id_opt: optional id for host :param storage: An instance that implements :interface:`~kademlia.storage.IStorage` :return: return a default kademlia instance """ if not id_opt: key_pair = generate_new_rsa_identity() id_opt = generate_peer_id_from(key_pair) node_id = id_opt.to_bytes() # ignore type for Kademlia module server = KademliaServer( # type: ignore ksize=ksize, alpha=alpha, node_id=node_id, storage=storage) return KadmeliaPeerRouter(server)
def test_peer_id_from_pubkey(): pubkey = datatypes.PublicKey( b"n\x85UD\xe9^\xbfo\x05\xd1z\xbd\xe5k\x87Y\xe9\xfa\xb3z:\xf8z\xc5\xd7K\xa6\x00\xbbc\xda4M\x10\x1cO\x88\tl\x82\x7f\xd7\xec6\xd8\xdc\xe2\x9c\xdcG\xa5\xea|\x9e\xc57\xf8G\xbe}\xfa\x10\xe9\x12" # noqa: E501 ) peer_id_expected = ID.from_base58( "QmQiv6sR3qHqhUVgC5qUBVWi8YzM6HknYbu4oQKVAqPCGF") assert peer_id_from_pubkey(pubkey) == peer_id_expected
async def dial_peer_maddr_with_retries(self, maddr: Multiaddr) -> None: """ Dial the peer with given multi-address repeatedly for `DIAL_RETRY_COUNT` times """ try: p2p_id = maddr.value_for_protocol(protocols.P_P2P) except (BinaryParseError, ProtocolLookupError) as error: self.logger.debug("Invalid maddr: %s, error: %s", maddr, error) raise DialPeerError from error peer_id = ID.from_base58(p2p_id) for i in range(DIAL_RETRY_COUNT): try: # exponential backoff... await asyncio.sleep(2**i + random.random()) await self.dial_peer_maddr(maddr, peer_id) return except DialPeerError: self.logger.debug( "Could not dial peer: %s, maddr: %s retrying attempt %d of %d...", peer_id, maddr, i, DIAL_RETRY_COUNT, ) continue raise DialPeerError
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.""" 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 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
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
def signature_validator(msg: rpc_pb2.Message) -> bool: """ Verify the message against the given public key. :param pubkey: the public key which signs the message. :param msg: the message signed. """ # Check if signature is attached if msg.signature == b"": logger.debug("Reject because no signature attached for msg: %s", msg) return False # Validate if message sender matches message signer, # i.e., check if `msg.key` matches `msg.from_id` msg_pubkey = deserialize_public_key(msg.key) if ID.from_pubkey(msg_pubkey) != msg.from_id: logger.debug( "Reject because signing key does not match sender ID for msg: %s", msg ) return False # First, construct the original payload that's signed by 'msg.key' msg_without_key_sig = rpc_pb2.Message( data=msg.data, topicIDs=msg.topicIDs, from_id=msg.from_id, seqno=msg.seqno ) payload = PUBSUB_SIGNING_PREFIX.encode() + msg_without_key_sig.SerializeToString() try: return msg_pubkey.verify(payload, msg.signature) except Exception: return False
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())
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
def make_pubsub_msg(origin_id: ID, topic_ids: Sequence[str], data: bytes, seqno: bytes) -> rpc_pb2.Message: return rpc_pb2.Message( from_id=origin_id.to_bytes(), seqno=seqno, data=data, topicIDs=list(topic_ids), )
def __init__( self, local_key_pair: KeyPair, secure_bytes_provider: Callable[[int], bytes] = default_secure_bytes_provider, ) -> None: self.local_private_key = local_key_pair.private_key self.local_peer = ID.from_pubkey(local_key_pair.public_key) self.secure_bytes_provider = secure_bytes_provider
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
async def dial_peer_maddr(self, maddr: Multiaddr) -> None: """ Parse `maddr`, get the ip:port and PeerID, and call `dial_peer` with the parameters. """ ip = maddr.value_for_protocol(protocols.P_IP4) port = maddr.value_for_protocol(protocols.P_TCP) peer_id = ID.from_base58(maddr.value_for_protocol(protocols.P_P2P)) await self.dial_peer(ip=ip, port=port, peer_id=peer_id)
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_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