improve validations

Author: msbrogli

hathor/p2p/sync_v2/agent.py

@@ -622,42 +622,42 @@ def handle_get_next_blocks(self, payload: str) -> None:
             self.protocol.send_error_and_close_connection('GET-NEXT-BLOCKS received before previous one finished')
             return
         data = GetNextBlocksPayload.parse_raw(payload)
+        start_block = self._validate_block(data.start_hash)
+        if start_block is None:
+            return
+        end_block = self._validate_block(data.end_hash)
+        if end_block is None:
+            return
         self.send_next_blocks(
-            start_hash=data.start_hash,
+            start_block=start_block,
             end_hash=data.end_hash,
             quantity=data.quantity,
         )
 
-    def send_next_blocks(self, start_hash: bytes, end_hash: bytes, quantity: int) -> None:
-        """ Send a NEXT-BLOCKS message.
-        """
-        self.log.debug('start NEXT-BLOCKS stream')
+    def _validate_block(self, _hash: VertexId) -> Optional[Block]:
+        """Validate block given in the GET-NEXT-BLOCKS and GET-TRANSACTIONS-BFS messages."""
         try:
-            blk = self.tx_storage.get_transaction(start_hash)
+            blk = self.tx_storage.get_transaction(_hash)
         except TransactionDoesNotExist:
-            # In case the tx does not exist we send a NOT-FOUND message
-            self.log.debug('requested start_hash not found', start_hash=start_hash.hex())
-            self.send_message(ProtocolMessages.NOT_FOUND, start_hash.hex())
-            return
+            self.log.debug('requested block not found', blk_id=_hash.hex())
+            self.send_message(ProtocolMessages.NOT_FOUND, _hash.hex())
+            return None
+
         if not isinstance(blk, Block):
-            self.log.debug('request start_hash is not a block', start_hash=start_hash.hex())
-            self.send_message(ProtocolMessages.NOT_FOUND, start_hash.hex())
-            return
-        assert blk.hash is not None
-        # XXX: it is not an error for the other peer to request a voided block, we'll pretend it doesn't exist, butf
-        blk_meta = blk.get_metadata()
-        if blk_meta.voided_by:
-            # In case the tx does not exist we send a NOT-FOUND message
-            self.log.debug('requested start_hash is voided, continue anyway', start_hash=start_hash.hex(),
-                           voided_by=[i.hex() for i in blk_meta.voided_by])
-            # XXX: we want to be able to not send this, but we do because the remote node could get stuck otherwise
-            # (tracked by issue #711)
-            # self.send_message(ProtocolMessages.NOT_FOUND, start_hash.hex())
-            # return
+            self.log.debug('request block is not a block', blk_id=_hash.hex())
+            self.send_message(ProtocolMessages.NOT_FOUND, _hash.hex())
+            return None
+
+        return blk
+
+    def send_next_blocks(self, start_block: Block, end_hash: bytes, quantity: int) -> None:
+        """ Send a NEXT-BLOCKS message.
+        """
+        self.log.debug('start NEXT-BLOCKS stream')
         if self._blk_streaming_server is not None and self._blk_streaming_server.is_running:
             self._blk_streaming_server.stop()
         limit = min(quantity, self.DEFAULT_STREAMING_LIMIT)
-        self._blk_streaming_server = BlockchainStreamingServer(self, blk, end_hash, limit=limit)
+        self._blk_streaming_server = BlockchainStreamingServer(self, start_block, end_hash, limit=limit)
         self._blk_streaming_server.start()
 
     def send_blocks(self, blk: Block) -> None:
@@ -867,44 +867,44 @@ def handle_get_transactions_bfs(self, payload: str) -> None:
             self.log.warn('ignore GET-TRANSACTIONS-BFS, already streaming')
             return
         data = GetTransactionsBFSPayload.parse_raw(payload)
-        # XXX: todo verify this limit while parsing the payload.
+
         if len(data.start_from) > MAX_GET_TRANSACTIONS_BFS_LEN:
             self.log.error('too many transactions in GET-TRANSACTIONS-BFS', state=self.state)
             self.protocol.send_error_and_close_connection('Too many transactions in GET-TRANSACTIONS-BFS')
             return
-        self.send_transactions_bfs(data.start_from, data.first_block_hash, data.last_block_hash)
 
-    def send_transactions_bfs(self,
-                              start_from: list[bytes],
-                              first_block_hash: bytes,
-                              last_block_hash: bytes) -> None:
-        """ Start a transactions BFS stream.
-        """
+        first_block = self._validate_block(data.first_block_hash)
+        if first_block is None:
+            return
+
+        last_block = self._validate_block(data.last_block_hash)
+        if last_block is None:
+            return
+
         start_from_txs = []
-        for start_from_hash in start_from:
+        for start_from_hash in data.start_from:
             try:
                 start_from_txs.append(self.tx_storage.get_transaction(start_from_hash))
             except TransactionDoesNotExist:
                 # In case the tx does not exist we send a NOT-FOUND message
                 self.log.debug('requested start_from_hash not found', start_from_hash=start_from_hash.hex())
                 self.send_message(ProtocolMessages.NOT_FOUND, start_from_hash.hex())
                 return
-        if not self.tx_storage.transaction_exists(first_block_hash):
-            # In case the tx does not exist we send a NOT-FOUND message
-            self.log.debug('requested first_block_hash not found', first_block_hash=first_block_hash.hex())
-            self.send_message(ProtocolMessages.NOT_FOUND, first_block_hash.hex())
-            return
-        if not self.tx_storage.transaction_exists(last_block_hash):
-            # In case the tx does not exist we send a NOT-FOUND message
-            self.log.debug('requested last_block_hash not found', last_block_hash=last_block_hash.hex())
-            self.send_message(ProtocolMessages.NOT_FOUND, last_block_hash.hex())
-            return
+
+        self.send_transactions_bfs(start_from_txs, first_block, last_block)
+
+    def send_transactions_bfs(self,
+                              start_from: list[BaseTransaction],
+                              first_block: Block,
+                              last_block: Block) -> None:
+        """ Start a transactions BFS stream.
+        """
         if self._tx_streaming_server is not None and self._tx_streaming_server.is_running:
             self._tx_streaming_server.stop()
         self._tx_streaming_server = TransactionsStreamingServer(self,
-                                                                start_from_txs,
-                                                                first_block_hash,
-                                                                last_block_hash,
+                                                                start_from,
+                                                                first_block,
+                                                                last_block,
                                                                 limit=self.DEFAULT_STREAMING_LIMIT)
         self._tx_streaming_server.start()
 

hathor/p2p/sync_v2/streamers.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from enum import IntFlag
-from typing import TYPE_CHECKING, Iterable, Iterator, Optional, Union, cast
+from typing import TYPE_CHECKING, Iterable, Iterator, Optional, Union
 
 from structlog import get_logger
 from twisted.internet.interfaces import IConsumer, IDelayedCall, IPushProducer
@@ -201,24 +201,21 @@ class TransactionsStreamingServer(_StreamingServerBase):
     def __init__(self,
                  sync_agent: 'NodeBlockSync',
                  start_from: list[BaseTransaction],
-                 first_block_hash: bytes,
-                 last_block_hash: bytes,
+                 first_block: Block,
+                 last_block: Block,
                  *,
                  limit: int = DEFAULT_STREAMING_LIMIT) -> None:
         # XXX: is limit needed for tx streaming? Or let's always send all txs for
         # a block? Very unlikely we'll reach this limit
         super().__init__(sync_agent, limit=limit)
 
-        self.first_block: Block = cast(Block, self.tx_storage.get_transaction(first_block_hash))
-        self.last_block: Block = cast(Block, self.tx_storage.get_transaction(last_block_hash))
+        self.first_block: Block = first_block
+        self.last_block: Block = last_block
         self.start_from = start_from
 
-        assert isinstance(self.first_block, Block)
-        assert isinstance(self.last_block, Block)
-
         # Validate that all transactions in `start_from` are confirmed by the first block.
         for tx in start_from:
-            assert tx.get_metadata().first_block == first_block_hash
+            assert tx.get_metadata().first_block == self.first_block.hash
 
         self.current_block: Optional[Block] = self.first_block
         self.bfs = BFSOrderWalk(self.tx_storage, is_dag_verifications=True, is_dag_funds=True, is_left_to_right=False)