59 lines
1.7 KiB
Python
59 lines
1.7 KiB
Python
from typing import Union
|
|
|
|
from multiaddr import Multiaddr
|
|
import trio
|
|
|
|
from libp2p.host.host_interface import IHost
|
|
from libp2p.peer.id import ID
|
|
from libp2p.peer.peerinfo import PeerInfo
|
|
|
|
from .daemon import Daemon
|
|
|
|
TDaemonOrHost = Union[IHost, Daemon]
|
|
|
|
|
|
def _get_peer_info(node: TDaemonOrHost) -> PeerInfo:
|
|
peer_info: PeerInfo
|
|
if isinstance(node, Daemon):
|
|
peer_info = node.peer_info
|
|
else: # isinstance(node, IHost)
|
|
peer_id = node.get_id()
|
|
maddrs = [
|
|
node.get_addrs()[0].decapsulate(Multiaddr(f"/p2p/{peer_id.to_string()}"))
|
|
]
|
|
peer_info = PeerInfo(peer_id, maddrs)
|
|
return peer_info
|
|
|
|
|
|
async def _is_peer(peer_id: ID, node: TDaemonOrHost) -> bool:
|
|
if isinstance(node, Daemon):
|
|
pinfos = await node.control.list_peers()
|
|
peers = tuple([pinfo.peer_id for pinfo in pinfos])
|
|
return peer_id in peers
|
|
else: # isinstance(node, IHost)
|
|
return peer_id in node.get_network().connections
|
|
|
|
|
|
async def connect(a: TDaemonOrHost, b: TDaemonOrHost) -> None:
|
|
# Type check
|
|
err_msg = (
|
|
f"Type of a={type(a)} or type of b={type(b)} is wrong."
|
|
"Should be either `IHost` or `Daemon`"
|
|
)
|
|
assert all(
|
|
[isinstance(node, IHost) or isinstance(node, Daemon) for node in (a, b)]
|
|
), err_msg
|
|
|
|
b_peer_info = _get_peer_info(b)
|
|
if isinstance(a, Daemon):
|
|
await a.control.connect(b_peer_info.peer_id, b_peer_info.addrs)
|
|
else: # isinstance(b, IHost)
|
|
await a.connect(b_peer_info)
|
|
# Allow additional sleep for both side to establish the connection.
|
|
await trio.sleep(0.1)
|
|
|
|
a_peer_info = _get_peer_info(a)
|
|
|
|
assert await _is_peer(b_peer_info.peer_id, a)
|
|
assert await _is_peer(a_peer_info.peer_id, b)
|