Handle malformed message during dao fork check

Author: pipermerriam

p2p/exceptions.py

@@ -42,6 +42,7 @@ class MalformedMessage(BaseP2PError):
     """
     Raised when a p2p command is received with a malformed message
     """
+    pass
 
 
 class UnknownProtocolCommand(BaseP2PError):

p2p/peer.py

@@ -355,10 +355,18 @@ async def read_msg(self) -> Tuple[protocol.Command, protocol._DecodedMsgType]:
         # too much time is being spent on this again, we need to consider running this in a
         # ProcessPoolExecutor(). Need to make sure we don't use all CPUs in the machine for that,
         # though, otherwise asyncio's event loop can't run and we can't keep up with other peers.
-        decoded_msg = cast(Dict[str, Any], cmd.decode(msg))
-        self.logger.trace("Successfully decoded %s msg: %s", cmd, decoded_msg)
-        self.received_msgs[cmd] += 1
-        return cmd, decoded_msg
+        try:
+            decoded_msg = cast(Dict[str, Any], cmd.decode(msg))
+        except MalformedMessage as err:
+            self.logger.debug(
+                "Malformed message from peer %s: CMD:%s Error: %r",
+                self, type(cmd).__name__, err,
+            )
+            raise
+        else:
+            self.logger.trace("Successfully decoded %s msg: %s", cmd, decoded_msg)
+            self.received_msgs[cmd] += 1
+            return cmd, decoded_msg
 
     def handle_p2p_msg(self, cmd: protocol.Command, msg: protocol._DecodedMsgType) -> None:
         """Handle the base protocol (P2P) messages."""
@@ -837,10 +845,14 @@ async def start_peer(self, peer: BasePeer) -> None:
             # check, we do it here because we want to perform it for incoming peer connections as
             # well.
             msgs = await self.ensure_same_side_on_dao_fork(peer)
-        except DAOForkCheckFailure as e:
-            self.logger.debug("DAO fork check with %s failed: %s", peer, e)
+        except DAOForkCheckFailure as err:
+            self.logger.debug("DAO fork check with %s failed: %s", peer, err)
             await peer.disconnect(DisconnectReason.useless_peer)
             return
+        except MalformedMessage as err:
+            self.logger.debug("DAO fork check with %s failed: %s", peer, err)
+            await peer.disconnect(DisconnectReason.bad_protocol)
+            return
         asyncio.ensure_future(peer.run(finished_callback=self._peer_finished))
         self._add_peer(peer, msgs)
 
@@ -962,9 +974,9 @@ async def ensure_same_side_on_dao_fork(
                     else:
                         msgs.append((cmd, msg))
                         continue
-            except (TimeoutError, PeerConnectionLost) as e:
+            except (TimeoutError, PeerConnectionLost) as err:
                 raise DAOForkCheckFailure(
-                    "Timed out waiting for DAO fork header from {}: {}".format(peer, e))
+                    "Timed out waiting for DAO fork header from {}: {}".format(peer, err))
 
             try:
                 request.validate_headers(headers)

p2p/protocol.py

@@ -22,7 +22,10 @@
 from eth.constants import NULL_BYTE
 from eth.rlp.headers import BlockHeader
 
-from p2p.exceptions import ValidationError
+from p2p.exceptions import (
+    MalformedMessage,
+    ValidationError,
+)
 from p2p.utils import get_devp2p_cmd_id
 
 
@@ -87,7 +90,13 @@ def decode_payload(self, rlp_data: bytes) -> _DecodedMsgType:
         else:
             decoder = sedes.List(
                 [type_ for _, type_ in self.structure], strict=self.decode_strict)
-        data = rlp.decode(rlp_data, sedes=decoder)
+        try:
+            data = rlp.decode(rlp_data, sedes=decoder)
+        except rlp.DecodingError as err:
+            raise MalformedMessage(
+                "Malformed %s message: %r".format(type(self).__name__, err)
+            ) from err
+
         if isinstance(self.structure, sedes.CountableList):
             return data
         return {
@@ -99,7 +108,7 @@ def decode_payload(self, rlp_data: bytes) -> _DecodedMsgType:
     def decode(self, data: bytes) -> _DecodedMsgType:
         packet_type = get_devp2p_cmd_id(data)
         if packet_type != self.cmd_id:
-            raise ValueError("Wrong packet type: {}".format(packet_type))
+            raise MalformedMessage("Wrong packet type: {}".format(packet_type))
         return self.decode_payload(data[1:])
 
     def encode(self, data: _DecodedMsgType) -> Tuple[bytes, bytes]: