Handle exceptions inside read_message

And remove the need of checking `None` for every read messages.

libp2p/stream_muxer/mplex/mplex.py

@@ -178,21 +178,15 @@ class Mplex(IMuxedConn):
         """
         Read a message off of the secured connection and add it to the corresponding message buffer
         """
-        # TODO Deal with other types of messages using flag (currently _)
 
         while True:
             try:
                 channel_id, flag, message = await self._wait_until_shutting_down_or_closed(
                     self.read_message()
                 )
-            except (
-                MplexUnavailable,
-                ConnectionResetError,
-                IncompleteReadError,
-            ) as error:
+            except MplexUnavailable as error:
                 print(f"!@# handle_incoming: read_message: exception={error}")
                 break
-            if channel_id is not None and flag is not None and message is not None:
             stream_id = StreamID(channel_id=channel_id, is_initiator=bool(flag & 1))
             is_stream_id_seen: bool
             stream: MplexStream
@@ -200,8 +194,6 @@ class Mplex(IMuxedConn):
                 is_stream_id_seen = stream_id in self.streams
                 if is_stream_id_seen:
                     stream = self.streams[stream_id]
-                # Other consequent stream message should wait until the stream get accepted
-                # TODO: Handle more tags, and refactor `HeaderTags`
             if flag == HeaderTags.NewStream.value:
                 if is_stream_id_seen:
                     # `NewStream` for the same id is received twice...
@@ -279,6 +271,8 @@ class Mplex(IMuxedConn):
 
             # Force context switch
             await asyncio.sleep(0)
+        # If we enter here, it means this connection is shutting down.
+        # We should clean the things up.
         await self._cleanup()
 
     async def read_message(self) -> Tuple[int, int, bytes]:
@@ -290,15 +284,19 @@ class Mplex(IMuxedConn):
         # FIXME: No timeout is used in Go implementation.
         # Timeout is set to a relatively small value to alleviate wait time to exit
         #  loop in handle_incoming
-        header = await decode_uvarint_from_stream(self.secured_conn)
-        # TODO: Handle the case of EOF and other exceptions?
         try:
+            header = await decode_uvarint_from_stream(self.secured_conn)
             message = await asyncio.wait_for(
                 read_varint_prefixed_bytes(self.secured_conn), timeout=5
             )
-        except asyncio.TimeoutError:
-            # TODO: Investigate what we should do if time is out.
-            return None, None, None
+        except (ConnectionResetError, IncompleteReadError) as error:
+            raise MplexUnavailable(
+                "failed to read messages correctly from the underlying connection"
+            ) from error
+        except asyncio.TimeoutError as error:
+            raise MplexUnavailable(
+                "failed to read more message body within the timeout"
+            ) from error
 
         flag = header & 0x07
         channel_id = header >> 3