py-libp2p/libp2p/protocol_muxer/multiselect.py
Piper Merriam ded3792924
Merge pull request #331 from dmuhs/fix/docs-format
Add automatic docstring formatting
2019-10-28 09:39:52 -06:00

103 lines
3.5 KiB
Python

from typing import Dict, Tuple
from libp2p.typing import StreamHandlerFn, TProtocol
from .exceptions import MultiselectCommunicatorError, MultiselectError
from .multiselect_communicator_interface import IMultiselectCommunicator
from .multiselect_muxer_interface import IMultiselectMuxer
MULTISELECT_PROTOCOL_ID = "/multistream/1.0.0"
PROTOCOL_NOT_FOUND_MSG = "na"
class Multiselect(IMultiselectMuxer):
"""Multiselect module that is responsible for responding to a multiselect
client and deciding on a specific protocol and handler pair to use for
communication."""
handlers: Dict[TProtocol, StreamHandlerFn]
def __init__(
self, default_handlers: Dict[TProtocol, StreamHandlerFn] = None
) -> None:
if not default_handlers:
default_handlers = {}
self.handlers = default_handlers
def add_handler(self, protocol: TProtocol, handler: StreamHandlerFn) -> None:
"""
Store the handler with the given protocol.
:param protocol: protocol name
:param handler: handler function
"""
self.handlers[protocol] = handler
async def negotiate(
self, communicator: IMultiselectCommunicator
) -> Tuple[TProtocol, StreamHandlerFn]:
"""
Negotiate performs protocol selection.
:param stream: stream to negotiate on
:return: selected protocol name, handler function
:raise MultiselectError: raised when negotiation failed
"""
await self.handshake(communicator)
while True:
try:
command = await communicator.read()
except MultiselectCommunicatorError as error:
raise MultiselectError(error)
if command == "ls":
# TODO: handle ls command
pass
else:
protocol = TProtocol(command)
if protocol in self.handlers:
try:
await communicator.write(protocol)
except MultiselectCommunicatorError as error:
raise MultiselectError(error)
return protocol, self.handlers[protocol]
try:
await communicator.write(PROTOCOL_NOT_FOUND_MSG)
except MultiselectCommunicatorError as error:
raise MultiselectError(error)
async def handshake(self, communicator: IMultiselectCommunicator) -> None:
"""
Perform handshake to agree on multiselect protocol.
:param communicator: communicator to use
:raise MultiselectError: raised when handshake failed
"""
try:
await communicator.write(MULTISELECT_PROTOCOL_ID)
except MultiselectCommunicatorError as error:
raise MultiselectError(error)
try:
handshake_contents = await communicator.read()
except MultiselectCommunicatorError as error:
raise MultiselectError(error)
if not is_valid_handshake(handshake_contents):
raise MultiselectError(
"multiselect protocol ID mismatch: "
f"received handshake_contents={handshake_contents}"
)
def is_valid_handshake(handshake_contents: str) -> bool:
"""
Determine if handshake is valid and should be confirmed.
:param handshake_contents: contents of handshake message
:return: true if handshake is complete, false otherwise
"""
return handshake_contents == MULTISELECT_PROTOCOL_ID