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.choice(ALPHABETS)
peer_id = ID(random_id_string.encode())
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_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
def test_peer_id(peer_id_string, peer_id_bytes, peer_id):
# test initialized with bytes
assert peer_id.to_bytes() == peer_id_bytes
assert peer_id.to_string() == peer_id_string
# test initialized with string
peer_id_2 = ID.from_base58(peer_id_string)
assert peer_id_2.to_bytes() == peer_id_bytes
assert peer_id_2.to_string() == peer_id_string
# test equal
assert peer_id == peer_id_2
# test not equal
peer_id_3 = ID.from_base58("QmbmfNDEth7Ucvjuxiw3SP3E4PoJzbk7g4Ge6ZDigbCsNp")
assert peer_id != peer_id_3
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_str_more_than_10():
random_id_string = ""
for _ in range(10):
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_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
async def disconnect(self, peer_id: ID) -> None:
disconnect_req = p2pd_pb.DisconnectRequest(peer=peer_id.to_bytes())
req = p2pd_pb.Request(type=p2pd_pb.Request.DISCONNECT,
disconnect=disconnect_req)
stream = await self.daemon_connector.open_connection()
await write_pbmsg(stream, req)
resp = p2pd_pb.Response()
await read_pbmsg_safe(stream, resp)
await stream.close()
raise_if_failed(resp)
async def get_closest_peers(self, key: bytes) -> Tuple[ID, ...]:
"""GET_CLOSEST_PEERS
"""
dht_req = p2pd_pb.DHTRequest(type=p2pd_pb.DHTRequest.GET_CLOSEST_PEERS,
key=key)
resps = await self._do_dht(dht_req)
try:
peer_ids = tuple(ID(dht_resp.value) for dht_resp in resps)
except AttributeError as e:
raise ControlFailure(
f"dht_resp should contains `value`: resps={resps}, e={e}")
return peer_ids
def test_peer_id_interop():
private_key_protobuf_bytes = base64.b64decode(PRIVATE_KEY_PROTOBUF_SERIALIZATION)
private_key_protobuf = pb.PrivateKey()
private_key_protobuf.ParseFromString(private_key_protobuf_bytes)
private_key_data = private_key_protobuf.data
private_key_impl = RSA.import_key(private_key_data)
private_key = RSAPrivateKey(private_key_impl)
public_key = private_key.get_public_key()
peer_id = ID.from_pubkey(public_key)
assert peer_id == EXPECTED_PEER_ID
async def untag_peer(self, peer_id: ID, tag: str) -> None:
"""UNTAG_PEER
"""
connmgr_req = p2pd_pb.ConnManagerRequest(
type=p2pd_pb.ConnManagerRequest.UNTAG_PEER, peer=peer_id.to_bytes(), tag=tag
)
req = p2pd_pb.Request(type=p2pd_pb.Request.CONNMANAGER, connManager=connmgr_req)
stream = await self.daemon_connector.open_connection()
await write_pbmsg(stream, req)
resp = p2pd_pb.Response()
await read_pbmsg_safe(stream, resp)
await stream.close()
raise_if_failed(resp)
async def connect(self, peer_id: ID, maddrs: Iterable[Multiaddr]) -> None:
stream = await self.daemon_connector.open_connection()
maddrs_bytes = [i.to_bytes() for i in maddrs]
connect_req = p2pd_pb.ConnectRequest(peer=peer_id.to_bytes(),
addrs=maddrs_bytes)
req = p2pd_pb.Request(type=p2pd_pb.Request.CONNECT,
connect=connect_req)
await write_pbmsg(stream, req)
resp = p2pd_pb.Response()
await read_pbmsg_safe(stream, resp)
await stream.close()
raise_if_failed(resp)
async def list_peers(self, topic: str) -> Tuple[ID, ...]:
"""PUBSUB LIST_PEERS
"""
pubsub_req = p2pd_pb.PSRequest(type=p2pd_pb.PSRequest.LIST_PEERS,
topic=topic)
req = p2pd_pb.Request(type=p2pd_pb.Request.PUBSUB, pubsub=pubsub_req)
stream = await self.daemon_connector.open_connection()
await write_pbmsg(stream, req)
resp = p2pd_pb.Response()
await read_pbmsg_safe(stream, resp)
await stream.close()
raise_if_failed(resp)
return tuple(
ID(peer_id_bytes) for peer_id_bytes in resp.pubsub.peerIDs)
async def identify(self) -> Tuple[ID, Tuple[Multiaddr, ...]]:
stream = await self.daemon_connector.open_connection()
req = p2pd_pb.Request(type=p2pd_pb.Request.IDENTIFY)
await write_pbmsg(stream, req)
resp = p2pd_pb.Response()
await read_pbmsg_safe(stream, resp)
await stream.close()
raise_if_failed(resp)
peer_id_bytes = resp.identify.id
maddrs_bytes = resp.identify.addrs
maddrs = tuple(Multiaddr(maddr_bytes) for maddr_bytes in maddrs_bytes)
peer_id = ID(peer_id_bytes)
return peer_id, maddrs
async def find_peers_connected_to_peer(
self, peer_id: ID) -> Tuple[PeerInfo, ...]:
"""FIND_PEERS_CONNECTED_TO_PEER
"""
dht_req = p2pd_pb.DHTRequest(
type=p2pd_pb.DHTRequest.FIND_PEERS_CONNECTED_TO_PEER,
peer=peer_id.to_bytes(),
)
resps = await self._do_dht(dht_req)
try:
pinfos = tuple(
PeerInfo.from_pb(dht_resp.peer) for dht_resp in resps)
except AttributeError as e:
raise ControlFailure(
f"dht_resp should contains peer info: resps={resps}, e={e}")
return pinfos # type: ignore
async def get_public_key(self, peer_id: ID) -> crypto_pb.PublicKey:
"""GET_PUBLIC_KEY
"""
dht_req = p2pd_pb.DHTRequest(type=p2pd_pb.DHTRequest.GET_PUBLIC_KEY,
peer=peer_id.to_bytes())
resps = await self._do_dht(dht_req)
if len(resps) != 1:
raise ControlFailure(
f"should only get one response, resps={resps}")
try:
public_key_pb_bytes = resps[0].value
except AttributeError as e:
raise ControlFailure(
f"dht_resp should contains `value`: resps={resps}, e={e}")
public_key_pb = crypto_pb.PublicKey()
public_key_pb.ParseFromString(public_key_pb_bytes)
return public_key_pb
async def find_peer(self, peer_id: ID) -> PeerInfo:
"""FIND_PEER
"""
dht_req = p2pd_pb.DHTRequest(type=p2pd_pb.DHTRequest.FIND_PEER,
peer=peer_id.to_bytes())
resps = await self._do_dht(dht_req)
if len(resps) != 1:
raise ControlFailure(
f"should only get one response from `find_peer`, resps={resps}"
)
dht_resp = resps[0]
try:
pinfo = dht_resp.peer
except AttributeError as e:
raise ControlFailure(
f"dht_resp should contains peer info: dht_resp={dht_resp}, e={e}"
)
return PeerInfo.from_pb(pinfo) # type: ignore
async def stream_open(
self, peer_id: ID, protocols: Sequence[str]
) -> Tuple[StreamInfo, anyio.abc.SocketStream]:
stream = await self.daemon_connector.open_connection()
stream_open_req = p2pd_pb.StreamOpenRequest(peer=peer_id.to_bytes(),
proto=list(protocols))
req = p2pd_pb.Request(type=p2pd_pb.Request.STREAM_OPEN,
streamOpen=stream_open_req)
await write_pbmsg(stream, req)
resp = p2pd_pb.Response()
await read_pbmsg_safe(stream, resp)
raise_if_failed(resp)
pb_stream_info = resp.streamInfo
stream_info = StreamInfo.from_pb(pb_stream_info)
return stream_info, stream
def test_eq_impl_for_bytes():
random_id_string = ""
for _ in range(10):
random_id_string += random.choice(ALPHABETS)
peer_id = ID(random_id_string.encode())
assert peer_id == random_id_string.encode()
def peer_id(peer_id_bytes):
return ID(peer_id_bytes)
def test_eq_false():
peer_id = ID("efgh")
other = ID("abcd")
assert peer_id != other
def peer_id_random():
return ID.from_base58("QmcgpsyWgH8Y8ajJz1Cu72KnS5uo2Aa2LpzU7kinSupNK1")
async def test_client_pubsub_subscribe(p2pcs):
peer_id_0, _ = await p2pcs[0].identify()
peer_id_1, _ = await p2pcs[1].identify()
await connect_safe(p2pcs[0], p2pcs[1])
await connect_safe(p2pcs[1], p2pcs[2])
topic = "topic123"
data = b"data"
stream_0 = await p2pcs[0].pubsub_subscribe(topic)
stream_1 = await p2pcs[1].pubsub_subscribe(topic)
# test case: `get_topics` after subscriptions
assert topic in await p2pcs[0].pubsub_get_topics()
assert topic in await p2pcs[1].pubsub_get_topics()
# wait for mesh built
await anyio.sleep(2)
# test case: `list_topic_peers` after subscriptions
assert peer_id_0 in await p2pcs[1].pubsub_list_peers(topic)
assert peer_id_1 in await p2pcs[0].pubsub_list_peers(topic)
# test case: publish, and both clients receive data
await p2pcs[0].pubsub_publish(topic, data)
pubsub_msg_0 = p2pd_pb.PSMessage()
await read_pbmsg_safe(stream_0, pubsub_msg_0)
assert pubsub_msg_0.data == data
pubsub_msg_1 = p2pd_pb.PSMessage()
await read_pbmsg_safe(stream_1, pubsub_msg_1)
assert pubsub_msg_1.data == data
# test case: publish more data
another_data_0 = b"another_data_0"
another_data_1 = b"another_data_1"
await p2pcs[0].pubsub_publish(topic, another_data_0)
await p2pcs[0].pubsub_publish(topic, another_data_1)
pubsub_msg_1_0 = p2pd_pb.PSMessage()
await read_pbmsg_safe(stream_1, pubsub_msg_1_0)
pubsub_msg_1_1 = p2pd_pb.PSMessage()
await read_pbmsg_safe(stream_1, pubsub_msg_1_1)
assert set([pubsub_msg_1_0.data,
pubsub_msg_1_1.data]) == set([another_data_0, another_data_1])
# test case: subscribe to multiple topics
another_topic = "topic456"
await p2pcs[0].pubsub_subscribe(another_topic)
stream_1_another = await p2pcs[1].pubsub_subscribe(another_topic)
await p2pcs[0].pubsub_publish(another_topic, another_data_0)
pubsub_msg_1_another = p2pd_pb.PSMessage()
await read_pbmsg_safe(stream_1_another, pubsub_msg_1_another)
assert pubsub_msg_1_another.data == another_data_0
# test case: test `from`
assert ID(getattr(pubsub_msg_1_1, "from")) == peer_id_0
# test case: test `from`, when it is sent through 1 hop(p2pcs[1])
stream_2 = await p2pcs[2].pubsub_subscribe(topic)
another_data_2 = b"another_data_2"
await p2pcs[0].pubsub_publish(topic, another_data_2)
pubsub_msg_2_0 = p2pd_pb.PSMessage()
await read_pbmsg_safe(stream_2, pubsub_msg_2_0)
assert ID(getattr(pubsub_msg_2_0, "from")) == peer_id_0
# test case: unsubscribe by closing the stream
await stream_0.close()
await anyio.sleep(0)
assert topic not in await p2pcs[0].pubsub_get_topics()
async def is_peer_removed_from_topic():
return peer_id_0 not in await p2pcs[1].pubsub_list_peers(topic)
await try_until_success(is_peer_removed_from_topic)
def from_pb(cls, peer_info_pb: p2pd_pb2.PeerInfo) -> PeerInfoLibP2P:
peer_id = ID(peer_info_pb.id)
addrs = [Multiaddr(addr) for addr in peer_info_pb.addrs]
return PeerInfoLibP2P(peer_id, addrs)
def from_pb(cls, pb_msg: p2pd_pb2.StreamInfo) -> "StreamInfo":
stream_info = cls(peer_id=ID(pb_msg.peer),
addr=Multiaddr(pb_msg.addr),
proto=pb_msg.proto)
return stream_info
