StarMirror/py-libp2p
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Handle MultiselectCommunicatorError

Browse Source
This commit is contained in:
NIC619 2019-09-15 16:58:22 +08:00
parent 879f193aa1
commit eaa74c4e26
No known key found for this signature in database
GPG Key ID: 570C35F5C2D51B17
2 changed files with 18 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
libp2p/protocol_muxer/multiselect.py
View File

@ -2,7 +2,7 @@ from typing import Dict, Tuple
from libp2p.typing import StreamHandlerFn, TProtocol from libp2p.typing import StreamHandlerFn, TProtocol
from .exceptions import MultiselectError from .exceptions import MultiselectCommunicatorError, MultiselectError
from .multiselect_communicator_interface import IMultiselectCommunicator from .multiselect_communicator_interface import IMultiselectCommunicator
from .multiselect_muxer_interface import IMultiselectMuxer from .multiselect_muxer_interface import IMultiselectMuxer
@ -46,7 +46,10 @@ class Multiselect(IMultiselectMuxer):
# Read and respond to commands until a valid protocol ID is sent # Read and respond to commands until a valid protocol ID is sent
while True: while True:
# Read message # Read message
command = await communicator.read() try:
command = await communicator.read()
except MultiselectCommunicatorError as error:
raise MultiselectError(str(error))
# Command is ls or a protocol # Command is ls or a protocol
if command == "ls": if command == "ls":
@ -76,7 +79,10 @@ class Multiselect(IMultiselectMuxer):
await communicator.write(MULTISELECT_PROTOCOL_ID) await communicator.write(MULTISELECT_PROTOCOL_ID)
# Read in the protocol ID from other party # Read in the protocol ID from other party
handshake_contents = await communicator.read() try:
handshake_contents = await communicator.read()
except MultiselectCommunicatorError as error:
raise MultiselectError(str(error))
# Confirm that the protocols are the same # Confirm that the protocols are the same
if not validate_handshake(handshake_contents): if not validate_handshake(handshake_contents):

12
libp2p/protocol_muxer/multiselect_client.py
View File

@ -2,7 +2,7 @@ from typing import Sequence
from libp2p.typing import TProtocol from libp2p.typing import TProtocol
from .exceptions import MultiselectClientError from .exceptions import MultiselectClientError, MultiselectCommunicatorError
from .multiselect_client_interface import IMultiselectClient from .multiselect_client_interface import IMultiselectClient
from .multiselect_communicator_interface import IMultiselectCommunicator from .multiselect_communicator_interface import IMultiselectCommunicator
@ -30,7 +30,10 @@ class MultiselectClient(IMultiselectClient):
await communicator.write(MULTISELECT_PROTOCOL_ID) await communicator.write(MULTISELECT_PROTOCOL_ID)
# Read in the protocol ID from other party # Read in the protocol ID from other party
handshake_contents = await communicator.read() try:
handshake_contents = await communicator.read()
except MultiselectCommunicatorError as error:
raise MultiselectClientError(str(error))
# Confirm that the protocols are the same # Confirm that the protocols are the same
if not validate_handshake(handshake_contents): if not validate_handshake(handshake_contents):
@ -79,7 +82,10 @@ class MultiselectClient(IMultiselectClient):
await communicator.write(protocol) await communicator.write(protocol)
# Get what counterparty says in response # Get what counterparty says in response
response = await communicator.read() try:
response = await communicator.read()
except MultiselectCommunicatorError as error:
raise MultiselectClientError(str(error))
# Return protocol if response is equal to protocol or raise error # Return protocol if response is equal to protocol or raise error
if response == protocol: if response == protocol: