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