class MuxerMultistream:
"""
MuxerMultistream is a multistream stream muxed transport multiplexer.
go implementation: github.com/libp2p/go-stream-muxer-multistream/multistream.go
"""
# NOTE: Can be changed to `typing.OrderedDict` since Python 3.7.2.
transports: "OrderedDict[TProtocol, MuxerClassType]"
multiselect: Multiselect
multiselect_client: MultiselectClient
def __init__(
self, muxer_transports_by_protocol: Mapping[TProtocol,
MuxerClassType]) -> None:
self.transports = OrderedDict()
self.multiselect = Multiselect()
self.multiselect_client = MultiselectClient()
for protocol, transport in muxer_transports_by_protocol.items():
self.add_transport(protocol, transport)
def add_transport(self, protocol: TProtocol,
transport: MuxerClassType) -> None:
"""
Add a protocol and its corresponding transport to multistream-select(multiselect).
The order that a protocol is added is exactly the precedence it is negotiated in
multiselect.
:param protocol: the protocol name, which is negotiated in multiselect.
:param transport: the corresponding transportation to the ``protocol``.
"""
# If protocol is already added before, remove it and add it again.
if protocol in self.transports:
del self.transports[protocol]
self.transports[protocol] = transport
self.multiselect.add_handler(protocol, None)
async def select_transport(self, conn: IRawConnection) -> MuxerClassType:
"""
Select a transport that both us and the node on the
other end of conn support and agree on
:param conn: conn to choose a transport over
:return: selected muxer transport
"""
protocol: TProtocol
communicator = MultiselectCommunicator(conn)
if conn.initiator:
protocol = await self.multiselect_client.select_one_of(
tuple(self.transports.keys()), communicator)
else:
protocol, _ = await self.multiselect.negotiate(communicator)
return self.transports[protocol]
async def new_conn(
self,
conn: ISecureConn,
generic_protocol_handler: GenericProtocolHandlerFn,
peer_id: ID,
) -> IMuxedConn:
transport_class = await self.select_transport(conn)
return transport_class(conn, generic_protocol_handler, peer_id)
class SecurityMultistream(ABC):
"""
SSMuxer is a multistream stream security transport multiplexer.
Go implementation: github.com/libp2p/go-conn-security-multistream/ssms.go
"""
# NOTE: Can be changed to `typing.OrderedDict` since Python 3.7.2.
transports: "OrderedDict[TProtocol, ISecureTransport]"
multiselect: Multiselect
multiselect_client: MultiselectClient
def __init__(self,
secure_transports_by_protocol: TSecurityOptions) -> None:
self.transports = OrderedDict()
self.multiselect = Multiselect()
self.multiselect_client = MultiselectClient()
for protocol, transport in secure_transports_by_protocol.items():
self.add_transport(protocol, transport)
def add_transport(self, protocol: TProtocol,
transport: ISecureTransport) -> None:
"""
Add a protocol and its corresponding transport to multistream-select(multiselect).
The order that a protocol is added is exactly the precedence it is negotiated in
multiselect.
:param protocol: the protocol name, which is negotiated in multiselect.
:param transport: the corresponding transportation to the ``protocol``.
"""
# If protocol is already added before, remove it and add it again.
if protocol in self.transports:
del self.transports[protocol]
self.transports[protocol] = transport
# Note: None is added as the handler for the given protocol since
# we only care about selecting the protocol, not any handler function
self.multiselect.add_handler(protocol, None)
async def secure_inbound(self, conn: IRawConnection) -> ISecureConn:
"""
Secure the connection, either locally or by communicating with opposing node via conn,
for an inbound connection (i.e. we are not the initiator)
:return: secure connection object (that implements secure_conn_interface)
"""
transport = await self.select_transport(conn, False)
secure_conn = await transport.secure_inbound(conn)
return secure_conn
async def secure_outbound(self, conn: IRawConnection,
peer_id: ID) -> ISecureConn:
"""
Secure the connection, either locally or by communicating with opposing node via conn,
for an inbound connection (i.e. we are the initiator)
:return: secure connection object (that implements secure_conn_interface)
"""
transport = await self.select_transport(conn, True)
secure_conn = await transport.secure_outbound(conn, peer_id)
return secure_conn
async def select_transport(self, conn: IRawConnection,
is_initiator: bool) -> ISecureTransport:
"""
Select a transport that both us and the node on the
other end of conn support and agree on
:param conn: conn to choose a transport over
:param is_initiator: true if we are the initiator, false otherwise
:return: selected secure transport
"""
protocol: TProtocol
communicator = MultiselectCommunicator(conn)
if is_initiator:
# Select protocol if initiator
protocol = await self.multiselect_client.select_one_of(
list(self.transports.keys()), communicator)
else:
# Select protocol if non-initiator
protocol, _ = await self.multiselect.negotiate(communicator)
# Return transport from protocol
return self.transports[protocol]
class Swarm(INetwork):
# pylint: disable=too-many-instance-attributes, cell-var-from-loop
def __init__(self, peer_id, peerstore, upgrader):
self.self_id = peer_id
self.peerstore = peerstore
self.upgrader = upgrader
self.connections = dict()
self.listeners = dict()
self.stream_handlers = dict()
self.transport = None
# Protocol muxing
self.multiselect = Multiselect()
self.multiselect_client = MultiselectClient()
# Create Notifee array
self.notifees = []
# Create generic protocol handler
self.generic_protocol_handler = create_generic_protocol_handler(self)
def get_peer_id(self):
return self.self_id
def set_stream_handler(self, protocol_id, stream_handler):
"""
:param protocol_id: protocol id used on stream
:param stream_handler: a stream handler instance
:return: true if successful
"""
self.multiselect.add_handler(protocol_id, stream_handler)
return True
async def dial_peer(self, peer_id):
"""
dial_peer try to create a connection to peer_id
:param peer_id: peer if we want to dial
:raises SwarmException: raised when no address if found for peer_id
:return: muxed connection
"""
# Get peer info from peer store
addrs = self.peerstore.addrs(peer_id)
if not addrs:
raise SwarmException("No known addresses to peer")
# TODO: define logic to choose which address to use, or try them all ?
multiaddr = addrs[0]
if peer_id in self.connections:
# If muxed connection already exists for peer_id,
# set muxed connection equal to existing muxed connection
muxed_conn = self.connections[peer_id]
else:
# Dial peer (connection to peer does not yet exist)
# Transport dials peer (gets back a raw conn)
raw_conn = await self.transport.dial(multiaddr, self.self_id)
# Use upgrader to upgrade raw conn to muxed conn
muxed_conn = self.upgrader.upgrade_connection(raw_conn, \
self.generic_protocol_handler, peer_id)
# Store muxed connection in connections
self.connections[peer_id] = muxed_conn
# Call notifiers since event occurred
for notifee in self.notifees:
await notifee.connected(self, muxed_conn)
return muxed_conn
async def new_stream(self, peer_id, protocol_ids):
"""
:param peer_id: peer_id of destination
:param protocol_id: protocol id
:return: net stream instance
"""
# Get peer info from peer store
addrs = self.peerstore.addrs(peer_id)
if not addrs:
raise SwarmException("No known addresses to peer")
multiaddr = addrs[0]
muxed_conn = await self.dial_peer(peer_id)
# Use muxed conn to open stream, which returns
# a muxed stream
# TODO: Remove protocol id from being passed into muxed_conn
muxed_stream = await muxed_conn.open_stream(protocol_ids[0], multiaddr)
# Perform protocol muxing to determine protocol to use
selected_protocol = await self.multiselect_client.select_one_of(
protocol_ids, muxed_stream)
# Create a net stream with the selected protocol
net_stream = NetStream(muxed_stream)
net_stream.set_protocol(selected_protocol)
# Call notifiers since event occurred
for notifee in self.notifees:
await notifee.opened_stream(self, net_stream)
return net_stream
async def listen(self, *args):
"""
:param *args: one or many multiaddrs to start listening on
:return: true if at least one success
For each multiaddr in args
Check if a listener for multiaddr exists already
If listener already exists, continue
Otherwise:
Capture multiaddr in conn handler
Have conn handler delegate to stream handler
Call listener listen with the multiaddr
Map multiaddr to listener
"""
for multiaddr in args:
if str(multiaddr) in self.listeners:
return True
async def conn_handler(reader, writer):
# Read in first message (should be peer_id of initiator) and ack
peer_id = id_b58_decode((await reader.read(1024)).decode())
writer.write("received peer id".encode())
await writer.drain()
# Upgrade reader/write to a net_stream and pass \
# to appropriate stream handler (using multiaddr)
raw_conn = RawConnection(multiaddr.value_for_protocol('ip4'),
multiaddr.value_for_protocol('tcp'),
reader, writer, False)
muxed_conn = self.upgrader.upgrade_connection(raw_conn, \
self.generic_protocol_handler, peer_id)
# 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)
try:
# Success
listener = self.transport.create_listener(conn_handler)
self.listeners[str(multiaddr)] = listener
await listener.listen(multiaddr)
# Call notifiers since event occurred
for notifee in self.notifees:
await notifee.listen(self, multiaddr)
return True
except IOError:
# Failed. Continue looping.
print("Failed to connect to: " + str(multiaddr))
# No multiaddr succeeded
return False
def notify(self, notifee):
"""
:param notifee: object implementing Notifee interface
:return: true if notifee registered successfully, false otherwise
"""
if isinstance(notifee, INotifee):
self.notifees.append(notifee)
return True
return False
def add_transport(self, transport):
# TODO: Support more than one transport
self.transport = transport
class BasicHost(IHost):
"""
BasicHost is a wrapper of a `INetwork` implementation. It performs protocol negotiation
on a stream with multistream-select right after a stream is initialized.
"""
_network: INetwork
_router: KadmeliaPeerRouter
peerstore: IPeerStore
multiselect: Multiselect
multiselect_client: MultiselectClient
def __init__(self,
network: INetwork,
router: KadmeliaPeerRouter = None) -> None:
self._network = network
self._network.set_stream_handler(self._swarm_stream_handler)
self._router = router
self.peerstore = self._network.peerstore
# Protocol muxing
self.multiselect = Multiselect()
self.multiselect_client = MultiselectClient()
def get_id(self) -> ID:
"""
:return: peer_id of host
"""
return self._network.get_peer_id()
def get_network(self) -> INetwork:
"""
:return: network instance of host
"""
return self._network
def get_peerstore(self) -> IPeerStore:
"""
:return: peerstore of the host (same one as in its network instance)
"""
return self.peerstore
def get_mux(self) -> Multiselect:
"""
:return: mux instance of host
"""
return self.multiselect
def get_addrs(self) -> List[multiaddr.Multiaddr]:
"""
:return: all the multiaddr addresses this host is listening to
"""
# TODO: We don't need "/p2p/{peer_id}" postfix actually.
p2p_part = multiaddr.Multiaddr("/p2p/{}".format(
self.get_id().pretty()))
addrs: List[multiaddr.Multiaddr] = []
for transport in self._network.listeners.values():
for addr in transport.get_addrs():
addrs.append(addr.encapsulate(p2p_part))
return addrs
def set_stream_handler(self, protocol_id: TProtocol,
stream_handler: StreamHandlerFn) -> None:
"""
set stream handler for given `protocol_id`
:param protocol_id: protocol id used on stream
:param stream_handler: a stream handler function
"""
self.multiselect.add_handler(protocol_id, stream_handler)
async def new_stream(self, peer_id: ID,
protocol_ids: Sequence[TProtocol]) -> INetStream:
"""
:param peer_id: peer_id that host is connecting
:param protocol_ids: available protocol ids to use for stream
:return: stream: new stream created
"""
net_stream = await self._network.new_stream(peer_id)
# Perform protocol muxing to determine protocol to use
try:
selected_protocol = await self.multiselect_client.select_one_of(
list(protocol_ids), MultiselectCommunicator(net_stream))
except MultiselectClientError as error:
logger.debug("fail to open a stream to peer %s, error=%s", peer_id,
error)
await net_stream.reset()
raise StreamFailure("failt to open a stream to peer %s",
peer_id) from error
net_stream.set_protocol(selected_protocol)
return net_stream
async def connect(self, peer_info: PeerInfo) -> None:
"""
connect ensures there is a connection between this host and the peer with
given `peer_info.peer_id`. connect will absorb the addresses in peer_info into its internal
peerstore. If there is not an active connection, connect will issue a
dial, and block until a connection is opened, or an error is returned.
:param peer_info: peer_info of the peer we want to connect to
:type peer_info: peer.peerinfo.PeerInfo
"""
self.peerstore.add_addrs(peer_info.peer_id, peer_info.addrs, 10)
# there is already a connection to this peer
if peer_info.peer_id in self._network.connections:
return
await self._network.dial_peer(peer_info.peer_id)
async def disconnect(self, peer_id: ID) -> None:
await self._network.close_peer(peer_id)
async def close(self) -> None:
await self._network.close()
# Reference: `BasicHost.newStreamHandler` in Go.
async def _swarm_stream_handler(self, net_stream: INetStream) -> None:
# Perform protocol muxing to determine protocol to use
try:
protocol, handler = await self.multiselect.negotiate(
MultiselectCommunicator(net_stream))
except MultiselectError:
await net_stream.reset()
return
net_stream.set_protocol(protocol)
await handler(net_stream)
class SecurityMultistream(ABC):
def __init__(self):
# Map protocol to secure transport
self.transports = {}
# Create multiselect
self.multiselect = Multiselect()
# Create multiselect client
self.multiselect_client = MultiselectClient()
def add_transport(self, protocol, transport):
# Associate protocol with transport
self.transports[protocol] = transport
# Add protocol and handler to multiselect
# Note: None is added as the handler for the given protocol since
# we only care about selecting the protocol, not any handler function
self.multiselect.add_handler(protocol, None)
async def secure_inbound(self, conn):
"""
Secure the connection, either locally or by communicating with opposing node via conn,
for an inbound connection (i.e. we are not the initiator)
:return: secure connection object (that implements secure_conn_interface)
"""
# Select a secure transport
transport = await self.select_transport(conn, False)
# Create secured connection
secure_conn = await transport.secure_inbound(conn)
return secure_conn
async def secure_outbound(self, conn, peer_id):
"""
Secure the connection, either locally or by communicating with opposing node via conn,
for an inbound connection (i.e. we are the initiator)
:return: secure connection object (that implements secure_conn_interface)
"""
# Select a secure transport
transport = await self.select_transport(conn, True)
# Create secured connection
secure_conn = await transport.secure_outbound(conn, peer_id)
return secure_conn
async def select_transport(self, conn, initiator):
"""
Select a transport that both us and the node on the
other end of conn support and agree on
:param conn: conn to choose a transport over
:param initiator: true if we are the initiator, false otherwise
:return: selected secure transport
"""
# TODO: Is conn acceptable to multiselect/multiselect_client
# instead of stream? In go repo, they pass in a raw conn
# (https://raw.githubusercontent.com/libp2p/go-conn-security-multistream/master/ssms.go)
protocol = None
if initiator:
# Select protocol if initiator
protocol = \
await self.multiselect_client.select_one_of(list(self.transports.keys()), conn)
else:
# Select protocol if non-initiator
protocol, _ = await self.multiselect.negotiate(conn)
# Return transport from protocol
return self.transports[protocol]
class BasicHost(IHost):
"""
BasicHost is a wrapper of a `INetwork` implementation.
It performs protocol negotiation on a stream with multistream-select
right after a stream is initialized.
"""
_network: INetworkService
peerstore: IPeerStore
multiselect: Multiselect
multiselect_client: MultiselectClient
def __init__(
self,
network: INetworkService,
default_protocols: "OrderedDict[TProtocol, StreamHandlerFn]" = None,
) -> None:
self._network = network
self._network.set_stream_handler(self._swarm_stream_handler)
self.peerstore = self._network.peerstore
# Protocol muxing
default_protocols = default_protocols or get_default_protocols(self)
self.multiselect = Multiselect(default_protocols)
self.multiselect_client = MultiselectClient()
def get_id(self) -> ID:
"""
:return: peer_id of host
"""
return self._network.get_peer_id()
def get_public_key(self) -> PublicKey:
return self.peerstore.pubkey(self.get_id())
def get_private_key(self) -> PrivateKey:
return self.peerstore.privkey(self.get_id())
def get_network(self) -> INetworkService:
"""
:return: network instance of host
"""
return self._network
def get_peerstore(self) -> IPeerStore:
"""
:return: peerstore of the host (same one as in its network instance)
"""
return self.peerstore
def get_mux(self) -> Multiselect:
"""
:return: mux instance of host
"""
return self.multiselect
def get_addrs(self) -> List[multiaddr.Multiaddr]:
"""
:return: all the multiaddr addresses this host is listening to
"""
# TODO: We don't need "/p2p/{peer_id}" postfix actually.
p2p_part = multiaddr.Multiaddr(f"/p2p/{self.get_id()!s}")
addrs: List[multiaddr.Multiaddr] = []
for transport in self._network.listeners.values():
for addr in transport.get_addrs():
addrs.append(addr.encapsulate(p2p_part))
return addrs
@asynccontextmanager
async def run(
self, listen_addrs: Sequence[multiaddr.Multiaddr]
) -> AsyncIterator[None]:
"""
run the host instance and listen to ``listen_addrs``.
:param listen_addrs: a sequence of multiaddrs that we want to listen to
"""
network = self.get_network()
async with background_trio_service(network):
await network.listen(*listen_addrs)
yield
def set_stream_handler(
self, protocol_id: TProtocol, stream_handler: StreamHandlerFn
) -> None:
"""
set stream handler for given `protocol_id`
:param protocol_id: protocol id used on stream
:param stream_handler: a stream handler function
"""
self.multiselect.add_handler(protocol_id, stream_handler)
async def new_stream(
self, peer_id: ID, protocol_ids: Sequence[TProtocol]
) -> INetStream:
"""
:param peer_id: peer_id that host is connecting
:param protocol_ids: available protocol ids to use for stream
:return: stream: new stream created
"""
net_stream = await self._network.new_stream(peer_id)
# Perform protocol muxing to determine protocol to use
try:
selected_protocol = await self.multiselect_client.select_one_of(
list(protocol_ids), MultiselectCommunicator(net_stream)
)
except MultiselectClientError as error:
logger.debug("fail to open a stream to peer %s, error=%s", peer_id, error)
await net_stream.reset()
raise StreamFailure(f"failed to open a stream to peer {peer_id}") from error
net_stream.set_protocol(selected_protocol)
return net_stream
async def connect(self, peer_info: PeerInfo) -> None:
"""
connect ensures there is a connection between this host and the peer
with given `peer_info.peer_id`. connect will absorb the addresses in
peer_info into its internal peerstore. If there is not an active
connection, connect will issue a dial, and block until a connection is
opened, or an error is returned.
:param peer_info: peer_info of the peer we want to connect to
:type peer_info: peer.peerinfo.PeerInfo
"""
self.peerstore.add_addrs(peer_info.peer_id, peer_info.addrs, 10)
# there is already a connection to this peer
if peer_info.peer_id in self._network.connections:
return
await self._network.dial_peer(peer_info.peer_id)
async def disconnect(self, peer_id: ID) -> None:
await self._network.close_peer(peer_id)
async def close(self) -> None:
await self._network.close()
# Reference: `BasicHost.newStreamHandler` in Go.
async def _swarm_stream_handler(self, net_stream: INetStream) -> None:
# Perform protocol muxing to determine protocol to use
try:
protocol, handler = await self.multiselect.negotiate(
MultiselectCommunicator(net_stream)
)
except MultiselectError as error:
peer_id = net_stream.muxed_conn.peer_id
logger.debug(
"failed to accept a stream from peer %s, error=%s", peer_id, error
)
await net_stream.reset()
return
net_stream.set_protocol(protocol)
await handler(net_stream)
class Swarm(INetwork):
self_id: ID
peerstore: IPeerStore
upgrader: TransportUpgrader
transport: ITransport
router: IPeerRouting
# TODO: Connection and `peer_id` are 1-1 mapping in our implementation,
# whereas in Go one `peer_id` may point to multiple connections.
connections: Dict[ID, IMuxedConn]
listeners: Dict[str, IListener]
stream_handlers: Dict[INetStream, Callable[[INetStream], None]]
multiselect: Multiselect
multiselect_client: MultiselectClient
notifees: List[INotifee]
def __init__(
self,
peer_id: ID,
peerstore: IPeerStore,
upgrader: TransportUpgrader,
transport: ITransport,
router: IPeerRouting,
):
self.self_id = peer_id
self.peerstore = peerstore
self.upgrader = upgrader
self.transport = transport
self.router = router
self.connections = dict()
self.listeners = dict()
self.stream_handlers = dict()
# Protocol muxing
self.multiselect = Multiselect()
self.multiselect_client = MultiselectClient()
# Create Notifee array
self.notifees = []
# Create generic protocol handler
self.generic_protocol_handler = create_generic_protocol_handler(self)
def get_peer_id(self) -> ID:
return self.self_id
def set_stream_handler(self, protocol_id: TProtocol,
stream_handler: StreamHandlerFn) -> bool:
"""
:param protocol_id: protocol id used on stream
:param stream_handler: a stream handler instance
:return: true if successful
"""
self.multiselect.add_handler(protocol_id, stream_handler)
return True
async def dial_peer(self, peer_id: ID) -> IMuxedConn:
"""
dial_peer try to create a connection to peer_id
:param peer_id: peer if we want to dial
:raises SwarmException: raised when an error occurs
:return: muxed connection
"""
if peer_id in self.connections:
# If muxed connection already exists for peer_id,
# set muxed connection equal to existing muxed connection
return self.connections[peer_id]
try:
# Get peer info from peer store
addrs = self.peerstore.addrs(peer_id)
except PeerStoreError:
raise SwarmException(f"No known addresses to peer {peer_id}")
if not addrs:
raise SwarmException(f"No known addresses to peer {peer_id}")
if not self.router:
multiaddr = addrs[0]
else:
multiaddr = self.router.find_peer(peer_id)
# Dial peer (connection to peer does not yet exist)
# Transport dials peer (gets back a raw conn)
raw_conn = await self.transport.dial(multiaddr, self.self_id)
# Per, https://discuss.libp2p.io/t/multistream-security/130, we first secure
# the conn and then mux the conn
try:
secured_conn = await self.upgrader.upgrade_security(
raw_conn, peer_id, True)
except SecurityUpgradeFailure as error:
# TODO: Add logging to indicate the failure
await raw_conn.close()
raise SwarmException(
f"fail to upgrade the connection to a secured connection from {peer_id}"
) from error
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 connection in connections
self.connections[peer_id] = muxed_conn
# Call notifiers since event occurred
for notifee in self.notifees:
await notifee.connected(self, muxed_conn)
return muxed_conn
async def new_stream(self, peer_id: ID,
protocol_ids: Sequence[TProtocol]) -> NetStream:
"""
:param peer_id: peer_id of destination
:param protocol_id: protocol id
:return: net stream instance
"""
muxed_conn = await self.dial_peer(peer_id)
# Use muxed conn to open stream, which returns a muxed stream
muxed_stream = await muxed_conn.open_stream()
# Perform protocol muxing to determine protocol to use
selected_protocol = await self.multiselect_client.select_one_of(
list(protocol_ids), MultiselectCommunicator(muxed_stream))
# Create a net stream with the selected protocol
net_stream = NetStream(muxed_stream)
net_stream.set_protocol(selected_protocol)
# Call notifiers since event occurred
for notifee in self.notifees:
await notifee.opened_stream(self, net_stream)
return net_stream
async def listen(self, *multiaddrs: Multiaddr) -> bool:
"""
:param multiaddrs: one or many multiaddrs to start listening on
:return: true if at least one success
For each multiaddr
Check if a listener for multiaddr exists already
If listener already exists, continue
Otherwise:
Capture multiaddr in conn handler
Have conn handler delegate to stream handler
Call listener listen with the multiaddr
Map multiaddr to listener
"""
for maddr in multiaddrs:
if str(maddr) in self.listeners:
return True
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)
try:
# Success
listener = self.transport.create_listener(conn_handler)
self.listeners[str(maddr)] = listener
await listener.listen(maddr)
# Call notifiers since event occurred
for notifee in self.notifees:
await notifee.listen(self, maddr)
return True
except IOError:
# Failed. Continue looping.
print("Failed to connect to: " + str(maddr))
# No maddr succeeded
return False
def notify(self, notifee: INotifee) -> bool:
"""
:param notifee: object implementing Notifee interface
:return: true if notifee registered successfully, false otherwise
"""
if isinstance(notifee, INotifee):
self.notifees.append(notifee)
return True
return False
def add_router(self, router: IPeerRouting) -> None:
self.router = router
async def close(self) -> None:
# TODO: Prevent from new listeners and conns being added.
# Reference: https://github.com/libp2p/go-libp2p-swarm/blob/8be680aef8dea0a4497283f2f98470c2aeae6b65/swarm.go#L124-L134 # noqa: E501
# Close listeners
await asyncio.gather(
*[listener.close() for listener in self.listeners.values()])
# Close connections
await asyncio.gather(
*[connection.close() for connection in self.connections.values()])
async def close_peer(self, peer_id: ID) -> None:
if peer_id not in self.connections:
return
connection = self.connections[peer_id]
del self.connections[peer_id]
await connection.close()
