fix dao check

Author: pipermerriam

p2p/peer.py

@@ -7,7 +7,6 @@
 import operator
 import random
 import struct
-import time
 from abc import (
     ABC,
     abstractmethod
@@ -776,8 +775,6 @@ async def ensure_same_side_on_dao_fork(
         wait for that we may receive other messages from the peer, which are returned so that they
         can be re-added to our subscribers' queues when the peer is finally added to the pool.
         """
-        from trinity.protocol.common.commands import BaseBlockHeaders
-        msgs = []
         for start_block, vm_class in self.vm_configuration:
             if not issubclass(vm_class, HomesteadVM):
                 continue
@@ -790,47 +787,40 @@ async def ensure_same_side_on_dao_fork(
             start_block = vm_class.dao_fork_block_number - 1
             # TODO: This can be either an `ETHPeer` or an `LESPeer`.  Will be
             # fixed once full awaitable request API is completed.
-            request = peer.request_block_headers(  # type: ignore
-                start_block,
-                max_headers=2,
-                reverse=False,
-            )
-            start = time.time()
             try:
-                while True:
-                    elapsed = int(time.time() - start)
-                    remaining_timeout = max(0, CHAIN_SPLIT_CHECK_TIMEOUT - elapsed)
-                    cmd, msg = await self.wait(
-                        peer.read_msg(), timeout=remaining_timeout)
-                    if isinstance(cmd, BaseBlockHeaders):
-                        headers = cmd.extract_headers(msg)
-                        break
-                    else:
-                        msgs.append((cmd, msg))
-                        continue
+                class MsgBuffer(PeerSubscriber):
+                    logger = logging.getLogger('p2p.peer.MsgBuffer')
+                    msg_queue_maxsize = 200
+                    subscription_msg_types = {protocol.Command}
+
+                msg_buffer = MsgBuffer()
+
+                with msg_buffer.subscribe_peer(peer):
+                    headers = await peer.handler.get_block_headers(  # type: ignore
+                        start_block,
+                        max_headers=2,
+                        reverse=False,
+                        timeout=CHAIN_SPLIT_CHECK_TIMEOUT,
+                    )
+
+                msgs = [msg_buffer.msg_queue.get_nowait()[1:] for _ in range(msg_buffer.queue_size)]
+
             except (TimeoutError, PeerConnectionLost) as err:
                 raise DAOForkCheckFailure(
-                    "Timed out waiting for DAO fork header from {}: {}".format(peer, err))
+                    "Timed out waiting for DAO fork header from {}: {}".format(peer, err)
+                ) from err
             except MalformedMessage as err:
                 raise DAOForkCheckFailure(
                     "Malformed message while doing DAO fork check with {0}: {1}".format(
                         peer, err,
                     )
                 ) from err
-
-            try:
-                request.validate_headers(headers)
             except ValidationError as err:
                 raise DAOForkCheckFailure(
                     "Invalid header response during DAO fork check: {}".format(err)
-                )
+                ) from err
 
-            if len(headers) != 2:
-                raise DAOForkCheckFailure(
-                    "Peer failed to return all requested headers for DAO fork check"
-                )
-            else:
-                parent, header = headers
+            parent, header = headers
 
             try:
                 vm_class.validate_header(header, parent, check_seal=True)

trinity/protocol/common/handlers.py

@@ -1,21 +1,20 @@
 from abc import abstractmethod
 from typing import (
+    Any,
     Dict,
     Type,
 )
 
 from p2p.peer import BasePeer
-from p2p.protocol import _DecodedMsgType
 from p2p.service import BaseService
 
 from .managers import BaseRequestManager
-from .requests import BaseRequest
 
 
 class BaseRequestResponseHandler(BaseService):
     @property
     @abstractmethod
-    def _managers(self) -> Dict[str, Type[BaseRequestManager[BasePeer, BaseRequest, _DecodedMsgType, _DecodedMsgType]]]:
+    def _managers(self) -> Dict[str, Type[BaseRequestManager[Any, Any, Any, Any]]]:
         pass
 
     def __init__(self, peer: BasePeer) -> None:

trinity/protocol/eth/managers.py

@@ -16,9 +16,7 @@
     BaseRequestManager,
 )
 
-from .commands import (
-    BlockHeaders,
-)
+from .commands import BlockHeaders
 from .requests import HeaderRequest
 
 if TYPE_CHECKING:
@@ -30,7 +28,7 @@ class GetBlockHeadersRequestManager(BaseRequestManager['ETHPeer', HeaderRequest,
 
     _response_msg_type: Type[Command] = BlockHeaders
 
-    async def __call__(self,
+    async def __call__(self,  # type: ignore
                        block_number_or_hash: BlockIdentifier,
                        max_headers: int = None,
                        skip: int = 0,

trinity/protocol/les/managers.py

@@ -35,7 +35,7 @@ class GetBlockHeadersRequestManager(BaseRequestManager['LESPeer', HeaderRequest,
 
     _response_msg_type: Type[Command] = BlockHeaders
 
-    async def __call__(self,
+    async def __call__(self,  # type: ignore
                        block_number_or_hash: BlockIdentifier,
                        max_headers: int = None,
                        skip: int = 0,