diff --git a/harness/udpstream/frame.py b/harness/udpstream/frame.py index e6bdefa..6de19ae 100644 --- a/harness/udpstream/frame.py +++ b/harness/udpstream/frame.py @@ -103,39 +103,55 @@ def _checked(name: str, value: int, limit: int) -> int: def pack_header(kind: int, frame_no: int, payload_len: int, rows: int, - stretch: int, payload: bytes = b"") -> bytes: + stretch: int, payload: bytes | None = None) -> bytes: """Build the 16-byte header for a frame. The CRC is computed over header bytes 0..11 plus ``payload``. Pass the - payload here to seal a complete frame in one call; with the default empty - payload the header carries the CRC of a header-only frame (use - :func:`seal_header` to add the CRC later -- that is what the chunked sender - does, since it builds the header before splitting the payload). + payload here to seal a complete frame in one call -- and then it must be + exactly ``payload_len`` bytes, or a frame the CRC covers but the length + field denies would go on the wire. Omit ``payload`` to build the header + before the body exists and seal it later with :func:`seal_header`. """ + kind = _checked("kind", kind, 0xFF) + frame_no = _checked("frame_no", frame_no, 0xFFFF) + payload_len = _checked("payload_len", payload_len, 0xFFFF) + rows = _checked("rows", rows, 0xFF) + stretch = _checked("stretch", stretch, 0xFF) + if payload is None: + body = b"" + else: + body = bytes(payload) + if len(body) != payload_len: + raise ValueError("payload is %d bytes but payload_len says %d" + % (len(body), payload_len)) prefix = _PREFIX.pack( MAGIC, VERSION, - _checked("kind", kind, 0xFF), - _checked("frame_no", frame_no, 0xFFFF), - _checked("payload_len", payload_len, 0xFFFF), - _checked("rows", rows, 0xFF), - _checked("stretch", stretch, 0xFF), + kind, + frame_no, + payload_len, + rows, + stretch, 0, # reserved ) assert len(prefix) == CRC_OFFSET - crc = crc16(prefix + bytes(payload)) + crc = crc16(prefix + body) return prefix + _TAIL.pack(crc, 0) def seal_header(header: bytes, payload: bytes = b"") -> bytes: """Return ``header`` with its CRC field recomputed over ``payload``. + ``header`` must be exactly :data:`HEADER_SIZE` bytes: a longer buffer is + rejected rather than silently truncated, so a caller cannot pass + ``header + payload`` here and lose the payload from the CRC. + Byte-for-byte identical to ``pack_header(..., payload=payload)`` for a header that came from :func:`pack_header`; kept separate so a caller that assembles a frame in pieces (header, then payload) can still seal it. """ - if len(header) < HEADER_SIZE: - raise ValueError("header shorter than %d bytes: %d" + if len(header) != HEADER_SIZE: + raise ValueError("header must be exactly %d bytes, got %d" % (HEADER_SIZE, len(header))) crc = crc16(bytes(header[:CRC_OFFSET]) + bytes(payload)) return bytes(header[:CRC_OFFSET]) + _TAIL.pack(crc, 0) diff --git a/harness/udpstream/tests/test_frame.py b/harness/udpstream/tests/test_frame.py index 1406de0..837e70b 100644 --- a/harness/udpstream/tests/test_frame.py +++ b/harness/udpstream/tests/test_frame.py @@ -79,6 +79,34 @@ def test_seal_header_after_the_payload_is_known(): payload=payload) +def test_seal_header_rejects_over_length_input(): + """A buffer longer than a header must not be silently truncated.""" + h = pack_header(KIND_STRIP, 3, 4, 8, 128) + with pytest.raises(ValueError): + seal_header(h + b"\x01\x02\x03\x04") # header + payload, not header + with pytest.raises(ValueError): + seal_header(h + b"\x00" * HEADER_SIZE) + with pytest.raises(ValueError): + seal_header(h[:HEADER_SIZE - 1]) + # Exactly a header is still accepted, and matches the one-shot call. + assert seal_header(h, b"\x01\x02\x03\x04") == pack_header( + KIND_STRIP, 3, 4, 8, 128, payload=b"\x01\x02\x03\x04") + + +def test_pack_header_rejects_payload_length_mismatch(): + """A CRC sealed over N bytes must not be sent under a different payload_len.""" + payload = b"\xAA" * 10 + pack_header(KIND_STRIP, 1, 10, 8, 128, payload=payload) # exact: fine + with pytest.raises(ValueError): + pack_header(KIND_STRIP, 1, 9, 8, 128, payload=payload) + with pytest.raises(ValueError): + pack_header(KIND_STRIP, 1, 11, 8, 128, payload=payload) + with pytest.raises(ValueError): + pack_header(KIND_STRIP, 1, 10, 8, 128, payload=b"") # explicit empty + # Omitting the payload is the header-only path: no length to disagree with. + assert unpack_header(pack_header(KIND_STRIP, 1, 10, 8, 128))["crc_ok"] is True + + # --- the byte layout the C guest will mirror ----------------------------- def test_byte_offsets_and_endianness():