diff --git a/src/rootfilespec/rntuple/envelope.py b/src/rootfilespec/rntuple/envelope.py index eee162c..b5bdcb0 100644 --- a/src/rootfilespec/rntuple/envelope.py +++ b/src/rootfilespec/rntuple/envelope.py @@ -2,6 +2,7 @@ from dataclasses import dataclass, field from typing import Annotated, Generic, TypeVar, cast +import xxhash # type: ignore[import-not-found] from typing_extensions import Self from rootfilespec.bootstrap.compression import decompress @@ -74,12 +75,8 @@ class REnvelope(ROOTSerializable): @classmethod def read(cls, buffer: ReadBuffer) -> tuple[Self, ReadBuffer]: """Reads an REnvelope from the given buffer.""" - #### Save initial buffer position (for checking unknown bytes) - payload_start_pos = buffer.relpos - #### Get the first 64bit integer (lengthType) which contains the length and type of the envelope - # lengthType, buffer = buffer.consume(8) - (lengthType,), buffer = buffer.unpack(" tuple[Self, ReadBuffer]: msg = f"Envelope type {typeID} read does not match passed class {cls.__name__}" raise ValueError(msg) - # Envelope size (uncompressed), encoded in the 48 most significant bits + # Envelope size (uncompressed), encoded in the 48 most significant bits. + # It includes the preamble and the checksum, so it is at least 16 bytes length = lengthType >> 16 - # Ensure that the length of the envelope matches the buffer length - if length - 8 != len(buffer): - msg = f"Length of envelope ({length} minus 8) of type {typeID} does not match buffer length ({len(buffer)})" + if length < 16: + msg = f"Length of envelope ({length}) of type {typeID} is shorter than 16 bytes" + raise ValueError(msg) + if length != len(buffer): + msg = f"Length of envelope ({length}) of type {typeID} does not match buffer length ({len(buffer)})" + raise ValueError(msg) + + #### Split the envelope once: the bytes the checksum covers, then the checksum + # The checksum covers [0, length - 8): the preamble, the payload and any + # unknown trailing bytes (root-io-spec ERRATA 5; "Checksum verification + # ... must include both known and unknown contents") + covered, trailer = buffer[: length - 8], buffer[length - 8 :] + (checksum,), rest = trailer.unpack(" EnvType: Envelopes are compressed, so this decompresses and deserializes. """ + if len(buffer) != self.size: + msg = ( + f"{self.envtype.__name__} at {self.locator}: expected {self.size} " + f"bytes, got {len(buffer)}" + ) + raise ValueError(msg) + #### Decompress the buffer if necessary - if len(buffer) != self.length: + # RNTuple decompression tests equality of the stored size (the locator's) + # and the length: equal means stored raw, smaller compressed, and larger + # is an error (root-io-spec NOTES 2; RNTupleZip.hxx:106-113) + if self.size > self.length: + msg = ( + f"{self.envtype.__name__} at {self.locator}: stored size " + f"{self.size} is larger than its uncompressed length {self.length}" + ) + raise ValueError(msg) + if self.size < self.length: buffer = decompress(buffer, self.length) #### Now read the envelope diff --git a/tests/test_envelope_checksums.py b/tests/test_envelope_checksums.py new file mode 100644 index 0000000..5179881 --- /dev/null +++ b/tests/test_envelope_checksums.py @@ -0,0 +1,161 @@ +from pathlib import Path + +import pytest + +from rootfilespec.bootstrap import BOOTSTRAP_CONTEXT, ROOT3a3aRNTuple +from rootfilespec.reader import open_path +from rootfilespec.rntuple.envelope import REnvelopeLocator +from rootfilespec.rntuple.header import HeaderEnvelope +from rootfilespec.rntuple.RLocator import StandardLocator +from rootfilespec.serializable import BufferContext, ReadBuffer + +DATA = Path(__file__).parent.parent / "reference" / "root-io-spec" / "data" / "rntuple" +pytestmark = pytest.mark.skipif( + not DATA.exists(), reason="reference/root-io-spec not checked out" +) + + +def _anchors(path: Path) -> list[ROOT3a3aRNTuple]: + with open_path(path) as reader: + keylist = reader.keylist() + return [ + reader.fetch(keylist[name]) + for name in keylist + if keylist[name].fClassName == b"ROOT::RNTuple" + ] + + +def _buffer(raw: bytes, offset: int, size: int) -> ReadBuffer: + return ReadBuffer( + memoryview(raw[offset : offset + size]), + 0, + BOOTSTRAP_CONTEXT, + BufferContext(abspos=offset), + ) + + +def _read(raw: bytes, loc): + return loc.read_from(_buffer(raw, loc.offset, loc.size)) + + +@pytest.mark.parametrize("name", sorted(p.name for p in DATA.glob("*.root"))) +def test_every_envelope_verifies(name: str): + """Issue #117: every envelope of every root-io-spec RNTuple fixture is checked + + Including rntuple/compressed.root, whose envelopes are compressed: the + checksum is over the uncompressed envelope. + """ + path = DATA / name + raw = path.read_bytes() + for anchor in _anchors(path): + header = _read(raw, anchor.header_locator) + footer = _read(raw, anchor.footer_locator) + assert footer.headerChecksum == header.checksum + for loc in footer.pagelist_locators: + assert _read(raw, loc).headerChecksum == header.checksum + + +def test_header_checksum_is_the_pinned_value(): + """rntuple/anchor.root's case.toml pins the header envelope's checksum at 500 + and the footer's copy of it at 854""" + path = DATA / "anchor.root" + raw = path.read_bytes() + (anchor,) = _anchors(path) + pinned = bytes([0x2A, 0x13, 0xA2, 0x84, 0xB3, 0x59, 0xD3, 0xDD]) + assert raw[500:508] == raw[854:862] == pinned + header = _read(raw, anchor.header_locator) + footer = _read(raw, anchor.footer_locator) + assert header.checksum == footer.headerChecksum == int.from_bytes(pinned, "little") + + +def test_corrupted_header_envelope_raises(): + """rntuple/anchor.root's header envelope is 268..508 (length 240, checksum at + 500..508). Changing any byte before the checksum must be caught.""" + path = DATA / "anchor.root" + raw = bytearray(path.read_bytes()) + (anchor,) = _anchors(path) + loc = anchor.header_locator + assert (loc.offset, loc.size, loc.length) == (268, 240, 240) + # The "n" of the ntuple's name "ntpl" at 288: flipping its case keeps the + # envelope parseable, so only the checksum can catch it + assert raw[288:292] == b"ntpl" + raw[288] ^= 0x20 + with pytest.raises(ValueError, match="HeaderEnvelope checksum mismatch"): + _read(bytes(raw), loc) + + +def test_corrupted_checksum_raises(): + path = DATA / "anchor.root" + raw = bytearray(path.read_bytes()) + (anchor,) = _anchors(path) + loc = anchor.header_locator + raw[loc.offset + loc.length - 1] ^= 0x01 + with pytest.raises(ValueError, match="HeaderEnvelope checksum mismatch"): + _read(bytes(raw), loc) + + +def test_stored_size_larger_than_length_raises(): + """root-io-spec NOTES 2: RNTuple decompression tests equality; a stored size + larger than the uncompressed length is an error, not a raw payload""" + loc = REnvelopeLocator( + length=16, locator=StandardLocator(size=24, offset=0), envtype=HeaderEnvelope + ) + with pytest.raises(ValueError, match="larger than its uncompressed length"): + loc.read_from(_buffer(bytes(24), 0, 24)) + + +def test_envelope_shorter_than_16_bytes_raises(): + """The length counts the 8-byte preamble and the 8-byte checksum, so it is + at least 16, as ROOT's DeserializeEnvelope requires""" + loc = REnvelopeLocator( + length=8, locator=StandardLocator(size=8, offset=0), envtype=HeaderEnvelope + ) + # The preamble alone: type 1 in the low 16 bits, length 8 in the upper 48 + preamble = (8 << 16 | 0x01).to_bytes(8, "little") + with pytest.raises(ValueError, match="shorter than 16 bytes"): + loc.read_from(_buffer(preamble, 0, 8)) + + +def test_envelope_length_not_matching_the_locator_raises(): + """rntuple/anchor.root's header envelope with the length in its preamble + (bytes 270..276) changed from 240 to 248: the anchor says 240 were stored""" + path = DATA / "anchor.root" + raw = bytearray(path.read_bytes()) + (anchor,) = _anchors(path) + loc = anchor.header_locator + assert raw[268:276] == (240 << 16 | 0x01).to_bytes(8, "little") + raw[268:276] = (248 << 16 | 0x01).to_bytes(8, "little") + with pytest.raises( + ValueError, match=r"Length of envelope \(248\) .* buffer length \(240\)" + ): + _read(bytes(raw), loc) + + +def test_every_corrupted_byte_is_a_checksum_mismatch(): + """The checksum is verified before the payload is parsed, as ROOT does, so + corrupting any byte the checksum covers reports the checksum, not whatever + the parser trips over first (an out-of-range slice, unknown feature flags, + ...).""" + path = DATA / "anchor.root" + raw = path.read_bytes() + (anchor,) = _anchors(path) + loc = anchor.header_locator + # The 8-byte preamble is checked first (type, length); everything after it, + # up to the checksum, is covered only by the checksum + for pos in range(loc.offset + 8, loc.offset + loc.length - 8): + corrupted = bytearray(raw) + corrupted[pos] ^= 0x80 + with pytest.raises(ValueError, match="HeaderEnvelope checksum mismatch"): + _read(bytes(corrupted), loc) + + +@pytest.mark.parametrize("size", [0, 200, 239]) +def test_short_read_raises(size: int): + """Fewer bytes than the locator's size (a truncated file) are reported as + such, not taken for a compressed envelope""" + path = DATA / "anchor.root" + raw = path.read_bytes() + (anchor,) = _anchors(path) + loc = anchor.header_locator + with pytest.raises(ValueError, match=f"expected 240 bytes, got {size}"): + loc.read_from(_buffer(raw, loc.offset, size))