py-libp2p/libp2p/security/security_multistream.py

107 lines
3.8 KiB
Python
Raw Normal View History

2019-05-02 01:54:19 +08:00
from abc import ABC
2019-04-28 07:42:05 +08:00
from libp2p.protocol_muxer.multiselect_client import MultiselectClient
from libp2p.protocol_muxer.multiselect import Multiselect
2019-08-01 19:12:11 +08:00
from typing import TYPE_CHECKING, NewType, Dict
if TYPE_CHECKING:
from libp2p.network.connection.raw_connection_interface import IRawConnection
from libp2p.peer.id import ID
from .typing import TSecurityDetails
from .secure_conn_interface import ISecureConn
from .secure_transport_interface import ISecureTransport
TProtocol = NewType("TProtocol", str)
2019-05-02 01:54:19 +08:00
2019-04-28 07:42:05 +08:00
"""
Represents a secured connection object, which includes a connection and details about the security
involved in the secured connection
Relevant go repo: https://github.com/libp2p/go-conn-security/blob/master/interface.go
"""
2019-08-01 06:00:12 +08:00
class SecurityMultistream(ABC):
2019-08-01 19:12:11 +08:00
transports: Dict[TProtocol, "ISecureTransport"]
multiselect: "Multiselect"
multiselect_client: "MultiselectClient"
def __init__(self) -> None:
2019-05-02 01:54:19 +08:00
# Map protocol to secure transport
2019-04-28 07:42:05 +08:00
self.transports = {}
# Create multiselect
self.multiselect = Multiselect()
# Create multiselect client
self.multiselect_client = MultiselectClient()
2019-08-01 19:12:11 +08:00
def add_transport(self, protocol: TProtocol, transport: "ISecureTransport") -> None:
2019-04-28 07:42:05 +08:00
# 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)
2019-08-01 19:12:11 +08:00
async def secure_inbound(self, conn: "IRawConnection") -> "ISecureConn":
2019-04-28 07:42:05 +08:00
"""
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
2019-05-01 04:07:26 +08:00
transport = await self.select_transport(conn, False)
2019-04-28 07:42:05 +08:00
# Create secured connection
secure_conn = await transport.secure_inbound(conn)
return secure_conn
2019-08-01 19:12:11 +08:00
async def secure_outbound(
self, conn: "IRawConnection", peer_id: "ID"
) -> "ISecureConn":
2019-04-28 07:42:05 +08:00
"""
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
2019-08-01 19:12:11 +08:00
async def select_transport(
self, conn: "IRawConnection", initiator: bool
) -> "ISecureTransport":
2019-04-28 07:42:05 +08:00
"""
2019-05-02 01:54:19 +08:00
Select a transport that both us and the node on the
2019-04-28 07:42:05 +08:00
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
"""
2019-05-02 01:54:19 +08:00
# 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)
2019-04-28 07:42:05 +08:00
protocol = None
if initiator:
# Select protocol if initiator
2019-08-01 06:00:12 +08:00
protocol = await self.multiselect_client.select_one_of(
list(self.transports.keys()), conn
)
2019-04-28 07:42:05 +08:00
else:
# Select protocol if non-initiator
2019-05-01 04:07:26 +08:00
protocol, _ = await self.multiselect.negotiate(conn)
2019-04-28 07:42:05 +08:00
# Return transport from protocol
return self.transports[protocol]