diff --git a/polyfile/nes.py b/polyfile/nes.py index 2983ff2..1e257e1 100644 --- a/polyfile/nes.py +++ b/polyfile/nes.py @@ -18,6 +18,8 @@ def parse_ines_header(header, parent=None): magic = header[:4] if magic != b'NES\x1A': raise InvalidMatch(f"Expected \"NES\\x1A\" magic at the beginning of the file, but found {magic!r}") + if len(header) < 16: + raise InvalidMatch(f"Expected a 16 byte iNES header, but the file only contains {len(header)} bytes") parent = Submatch( name='iNESHeader', match_obj=header, diff --git a/tests/unit/test_nes.py b/tests/unit/test_nes.py new file mode 100644 index 0000000..3613a99 --- /dev/null +++ b/tests/unit/test_nes.py @@ -0,0 +1,36 @@ +from unittest import TestCase + +from polyfile.nes import parse_ines_header +from polyfile.polyfile import InvalidMatch, Match, Matcher + + +VALID_HEADER = b"NES\x1a" + bytes([1, 0, 0, 0]) + bytes(8) + + +class INESHeaderTest(TestCase): + def _root(self) -> Match: + return Match(name="root", match_obj=b"", relative_offset=0, matcher=Matcher()) + + def test_rejects_bad_magic(self): + with self.assertRaises(InvalidMatch): + list(parse_ines_header(b"NOPE" + bytes(12), self._root())) + + def test_rejects_truncated_header(self): + """A header shorter than 16 bytes is not a usable iNES header. + + Lengths 4 to 10 used to raise IndexError, which the parser dispatcher + catches as an unexpected error, logs, and then drops the rest of the + submatches. Lengths 11 to 15 raised nothing at all and produced an + iNESHeader claiming to be 16 bytes with a short UnusedPadding. + """ + for length in range(4, 16): + with self.subTest(length=length): + with self.assertRaises(InvalidMatch): + list(parse_ines_header(VALID_HEADER[:length], self._root())) + + def test_accepts_full_header(self): + names = [m.name for m in parse_ines_header(VALID_HEADER, self._root())] + + self.assertEqual(names[0], "iNESHeader") + for expected in ("iNESMagic", "PRGSize", "CHRSize", "UnusedPadding"): + self.assertIn(expected, names)