Regression tests for python-bitcoinlib msg_headers trailing tx-count varint (petertodd/python-bitcoinlib#320)

test_msg_headers_wire_format.py

# Regression tests for the on-wire format of the P2P 'headers' message.
#
# These three tests were written alongside the msg_headers trailing-tx-count-
# varint fix (petertodd/python-bitcoinlib#320). They were deliberately held
# back from the PR's diff to keep the change minimal -- only the fix itself
# ships in the PR, not new test scaffolding.
#
# They are kept here so the test logic isn't lost. If Peter asks for an
# in-tree regression test during PR review, drop this class into
# bitcoin/tests/test_messages.py.
#
# To run standalone from this directory (with the parent repo importable):
#   python -m unittest test_msg_headers_wire_format -v
#
# To verify the tests actually catch the bug, run them against the
# pre-fix msg_headers and they should all three fail (chain misalignment,
# round-trip mismatch, unread trailing bytes).

import unittest

from bitcoin.messages import msg_headers

from io import BytesIO


class Test_msg_headers_wire_format(unittest.TestCase):
    """Regression tests for the on-wire format of the P2P 'headers' message.

    Each entry in a 'headers' message is an 80-byte block header followed by
    a 0-byte transaction-count varint (the standard CBlock framing, with no
    transactions). msg_headers.msg_deser must consume that trailing varint,
    or subsequent headers in the same message will be misaligned by one byte
    per skipped varint.
    """

    # Mainnet block 0 header (80 bytes, on-wire byte order).
    BLOCK_0_HEADER_HEX = (
        '01000000'
        '0000000000000000000000000000000000000000000000000000000000000000'
        '3ba3edfd7a7b12b27ac72c3e67768f617fc81bc3888a51323a9fb8aa4b1e5e4a'
        '29ab5f49'
        'ffff001d'
        '1dac2b7c'
    )

    # Mainnet block 1 header. hashPrevBlock equals the hash of block 0.
    BLOCK_1_HEADER_HEX = (
        '01000000'
        '6fe28c0ab6f1b372c1a6a246ae63f74f931e8365e15a089c68d6190000000000'
        '982051fd1e4ba744bbbe680e1fee14677ba1a3c3540bf7b1cdb606e857233e0e'
        '61bc6649'
        'ffff001d'
        '01e36299'
    )

    def _build_wire_body(self, header_hex_list):
        """Build the on-wire body of a 'headers' message: varint(N) + per
        header (80 bytes + 0-byte tx-count varint)."""
        body = bytes([len(header_hex_list)])  # varint, fits in one byte
        for h in header_hex_list:
            body += bytes.fromhex(h)
            body += b'\x00'  # trailing tx-count varint
        return body

    def test_deserialize_chain_continuity(self):
        """Header 1's hashPrevBlock must match header 0's hash after deser."""
        body = self._build_wire_body(
            [self.BLOCK_0_HEADER_HEX, self.BLOCK_1_HEADER_HEX])
        msg = msg_headers.msg_deser(BytesIO(body))
        self.assertEqual(len(msg.headers), 2)
        self.assertEqual(
            msg.headers[1].hashPrevBlock,
            msg.headers[0].GetHash(),
            "Header 1 must chain to header 0; misalignment indicates the "
            "trailing tx-count varint was not consumed during deserialization")

    def test_roundtrip_against_wire_body(self):
        """Serializing a deserialized wire body must reproduce the original."""
        body = self._build_wire_body(
            [self.BLOCK_0_HEADER_HEX, self.BLOCK_1_HEADER_HEX])
        msg = msg_headers.msg_deser(BytesIO(body))
        f = BytesIO()
        msg.msg_ser(f)
        self.assertEqual(f.getvalue(), body)

    def test_deserialize_consumes_exact_bytes(self):
        """Deser must consume the entire body and no more (no trailing bytes)."""
        body = self._build_wire_body(
            [self.BLOCK_0_HEADER_HEX, self.BLOCK_1_HEADER_HEX])
        f = BytesIO(body)
        msg_headers.msg_deser(f)
        self.assertEqual(f.read(), b'',
                         "Deserializer left unread bytes in the stream")


if __name__ == '__main__':
    unittest.main()