| import hashlib |
| import os |
| import struct |
| import tempfile |
| import unittest |
| from pathlib import Path |
| from unittest.mock import patch |
|
|
| from lightnest.binary import (BinaryFormatError, BinaryReader, Limits, SLOT_MAGIC, |
| find_slots, index_elf, index_lnw1, inspect_binary, read_small_program) |
|
|
|
|
| def blob(shape=(2,), data=None): |
| name = b'weight' |
| header = b'LNW1' + struct.pack('<10I', 1, 8, 1, 2, 32, 32, 16, 16, 1, 1) |
| prefix = header + struct.pack('<I', 1) + b'a' + struct.pack('<I', len(name)) + name |
| prefix += struct.pack('<I', len(shape)) + struct.pack('<' + 'I' * len(shape), *shape) |
| count = 1 |
| for d in shape: |
| count *= d |
| return prefix, data if data is not None else b'\0' * (4 * count) |
|
|
|
|
| class BinaryTests(unittest.TestCase): |
| def setUp(self): |
| self.tmp = tempfile.TemporaryDirectory() |
| self.path = Path(self.tmp.name) / 'input' |
|
|
| def tearDown(self): |
| self.tmp.cleanup() |
|
|
| def write(self, data): |
| self.path.write_bytes(data) |
|
|
| def test_range_and_chunk_guards(self): |
| self.write(bytes(range(100))) |
| with BinaryReader(self.path, max_read=16) as r: |
| for off, size in ((-1, 1), (100, 1), (0, 17), (1, -1)): |
| with self.assertRaises(BinaryFormatError): |
| r.read_at(off, size) |
| self.assertEqual(b''.join(r.iter_chunks(10, 20, chunk_size=3)), bytes(range(10, 30))) |
| with self.assertRaises(ValueError): |
| list(r.iter_chunks(chunk_size=0)) |
| self.assertEqual(r.sha256(), hashlib.sha256(bytes(range(100))).hexdigest()) |
|
|
| def test_magic_cross_boundary_and_overlaps(self): |
| self.write(b'xxxxxxMAGICaaaaaMAGICaaa') |
| with BinaryReader(self.path, max_read=8) as r: |
| self.assertEqual(list(r.find_magic(b'MAGIC')), [6, 16]) |
| self.assertEqual(list(r.find_magic(b'aaa')), [11, 12, 13, 21]) |
| with self.assertRaises(BinaryFormatError): |
| list(r.find_magic(b'aaa', max_matches=1)) |
| self.assertEqual(list(r.find_magic(b'MAGIC', 7, 13)), []) |
|
|
| def test_sparse_gigabyte_tensor_never_read(self): |
| |
| prefix, _ = blob((1,), b'') |
| prefix = prefix[:-4] + struct.pack('<I', 1 << 30) |
| with self.path.open('wb') as f: |
| f.write(prefix) |
| f.truncate(len(prefix) + (1 << 32)) |
| real = os.pread |
| observed = [] |
| def guarded(fd, n, offset): |
| observed.append((offset, n)) |
| self.assertLessEqual(n, 32) |
| self.assertLess(offset, len(prefix)) |
| return real(fd, n, offset) |
| with BinaryReader(self.path, max_read=32) as r, patch('os.pread', guarded): |
| info = index_lnw1(r) |
| self.assertEqual(info['tensors'][0]['nbytes'], 1 << 32) |
| self.assertEqual(info['end_offset'], self.path.stat().st_size) |
| self.assertLess(sum(n for _, n in observed), 100) |
|
|
| def test_lnw1_embedded_offsets(self): |
| prefix, data = blob() |
| self.write(b'prefix' + prefix + data + b'suffix') |
| with BinaryReader(self.path) as r: |
| info = index_lnw1(r, 6, end=6+len(prefix)+len(data)) |
| self.assertEqual(info['tensors'][0]['shape'], [2]) |
| self.assertEqual(info['tensors'][0]['data_offset'], 6+len(prefix)) |
| self.assertEqual(info['architecture']['heads'], 2) |
|
|
| def test_truncated_every_metadata_boundary(self): |
| prefix, data = blob() |
| for length in range(len(prefix) + len(data)): |
| self.write((prefix + data)[:length]) |
| with BinaryReader(self.path) as r, self.assertRaises(BinaryFormatError): |
| index_lnw1(r) |
|
|
| def test_lnw1_guards(self): |
| prefix, data = blob() |
| cases = [] |
| for offset, value in ((4, 2), (8, 0), (16, 3), (36, 100001), (40, 10001), (44, 4097)): |
| b = bytearray(prefix + data) |
| struct.pack_into('<I', b, offset, value) |
| cases.append(b) |
| cases.append(prefix[:-4] + struct.pack('<I', 0) + data) |
| cases.append(prefix[:48] + b'\xff' + prefix[49:] + data) |
| for value in cases: |
| self.write(value) |
| with BinaryReader(self.path) as r, self.assertRaises(BinaryFormatError): |
| index_lnw1(r) |
| self.write(prefix + data) |
| with BinaryReader(self.path) as r, self.assertRaises(BinaryFormatError): |
| index_lnw1(r, limits=Limits(max_metadata_bytes=20)) |
|
|
| def test_nonregular_and_concurrent_truncate(self): |
| with self.assertRaises(BinaryFormatError): |
| BinaryReader(self.tmp.name) |
| fifo = Path(self.tmp.name) / 'fifo' |
| os.mkfifo(fifo) |
| with self.assertRaises(BinaryFormatError): |
| BinaryReader(fifo) |
| self.write(b'1234') |
| with BinaryReader(self.path) as r: |
| self.write(b'') |
| with self.assertRaises(BinaryFormatError): |
| r.read_at(0, 4) |
|
|
| def test_slots_bounded_and_ambiguous(self): |
| slot = SLOT_MAGIC + struct.pack('<II', 256, 4) + b'RV32\0'.ljust(16, b'\0') + bytes(256) |
| self.write(bytes(65530) + slot + slot) |
| with BinaryReader(self.path) as r: |
| slots = find_slots(r) |
| self.assertEqual([s['offset'] for s in slots], [65530, 65826]) |
| self.assertEqual(slots[0]['data_offset'], 65570) |
| self.write(SLOT_MAGIC + struct.pack('<II', 255, 0) + bytes(272)) |
| with BinaryReader(self.path) as r: |
| self.assertEqual(find_slots(r), []) |
|
|
| def test_unknown_and_explicit_raw(self): |
| self.write(b'\x73\x00\x10\x00') |
| self.assertEqual(inspect_binary(self.path)['format'], 'unknown') |
| self.assertEqual(inspect_binary(self.path, raw_arch='RV32')['architecture_source'], 'explicit hint') |
| self.write(b'abc') |
| with self.assertRaises(BinaryFormatError): |
| inspect_binary(self.path, raw_arch='RV32') |
|
|
| def test_streaming_sparse_hash_read_bounds(self): |
| with self.path.open('wb') as f: |
| f.truncate(2 * 1024 * 1024) |
| real = os.pread |
| sizes = [] |
| def guarded(fd, size, offset): |
| sizes.append(size) |
| self.assertLessEqual(size, 1024) |
| return real(fd, size, offset) |
| with BinaryReader(self.path, max_read=1024) as r, patch('os.pread', guarded): |
| digest = r.sha256() |
| expected = hashlib.sha256() |
| for _ in range(2048): |
| expected.update(bytes(1024)) |
| self.assertEqual(digest, expected.hexdigest()) |
| self.assertEqual(len(sizes), 2048) |
|
|
| def test_elf32_big_endian_and_nobits(self): |
| ident = b'\x7fELF' + bytes([1, 2, 1]) + bytes(9) |
| header = ident + struct.pack('>HHIIIIIHHHHHH', 2, 243, 1, 0, 0, 52, 0, 52, 32, 0, 40, 1, 0) |
| section = struct.pack('>IIIIIIIIII', 0, 8, 3, 4096, 92, 1 << 30, 0, 0, 4, 0) |
| self.write(header + section) |
| with BinaryReader(self.path) as r: |
| info = index_elf(r) |
| self.assertEqual(info['bits'], 32) |
| self.assertEqual(info['byteorder'], 'big') |
| self.assertEqual(info['architecture'], 'riscv') |
| self.assertFalse(info['sections'][0]['file_backed']) |
| self.assertEqual(info['sections'][0]['nbytes'], 1 << 30) |
|
|
| def test_unknown_100gb_default_reads_only_magic(self): |
| with self.path.open('wb') as f: |
| f.truncate(100 * 1024**3) |
| real = os.pread |
| reads = [] |
| def guarded(fd, n, offset): |
| reads.append((offset, n)) |
| self.assertLessEqual(offset+n, 4) |
| return real(fd, n, offset) |
| with patch('os.pread', guarded): |
| info = inspect_binary(self.path) |
| self.assertEqual(reads, [(0, 4)]) |
| self.assertEqual(info['weight_search']['scope'], 'not searched') |
| self.assertFalse(info['slot_search']['whole_file']) |
|
|
| def test_elf_huge_embedded_payload_prefix_budget(self): |
| prefix, _ = blob((1,), b'') |
| prefix = prefix[:-4] + struct.pack('<I', 1 << 30) |
| payload_size = len(prefix) + (1 << 32) |
| ident = b'\x7fELF' + bytes([2, 1, 1]) + bytes(9) |
| header = ident + struct.pack('<HHIQQQIHHHHHH', 2, 62, 1, 0, 0, 64, 0, 64, 56, 0, 64, 3, 1) |
| names = b'\0.shstrtab\0.rodata\0' |
| def section(name, kind, flags, offset, size): |
| return struct.pack('<IIQQQQIIQQ', name, kind, flags, 0, offset, size, 0, 0, 1, 0) |
| with self.path.open('wb') as f: |
| f.write(header + section(0, 0, 0, 0, 0) + section(1, 3, 0, 256, len(names)) |
| + section(11, 1, 2, 4096, payload_size) + names) |
| f.seek(4096) |
| f.write(prefix) |
| f.truncate(4096+payload_size) |
| real = os.pread |
| reads = [] |
| def guarded(fd, n, offset): |
| reads.append((offset, n)) |
| self.assertLess(offset, 4096+65536) |
| return real(fd, n, offset) |
| with patch('os.pread', guarded): |
| info = inspect_binary(self.path, scan_bytes=4096) |
| self.assertEqual(len(info['embedded_weights']), 1) |
| self.assertEqual(info['embedded_weights'][0]['tensors'][0]['nbytes'], 1 << 32) |
| self.assertEqual(info['weight_search']['search_bytes_read'], 4096) |
| self.assertFalse(info['weight_search']['whole_file']) |
| self.assertFalse(info['weight_search']['complete_within_selected_ranges']) |
| self.assertLess(sum(n for _, n in reads), 5000) |
| no_scan = inspect_binary(self.path, scan_bytes=0) |
| self.assertEqual(no_scan['embedded_weights'], []) |
| self.assertFalse(no_scan['weight_search']['complete_within_selected_ranges']) |
|
|
| def test_elf_direct_slot_lookup(self): |
| ident = b'\x7fELF' + bytes([2, 1, 1]) + bytes(9) |
| header = ident + struct.pack('<HHIQQQIHHHHHH', 2, 62, 1, 0, 0, 64, 0, 64, 56, 0, 64, 3, 1) |
| names = b'\0.shstrtab\0.lnself\0' |
| def section(name, kind, flags, offset, size): |
| return struct.pack('<IIQQQQIIQQ', name, kind, flags, 0, offset, size, 0, 0, 1, 0) |
| slot = SLOT_MAGIC + struct.pack('<II', 256, 4) + b'RV32\0'.ljust(16, b'\0') + bytes(256) |
| self.write(header + section(0, 0, 0, 0, 0) + section(1, 3, 0, 256, len(names)) |
| + section(11, 1, 3, 512, len(slot)) + names) |
| with self.path.open('r+b') as f: |
| f.seek(512) |
| f.write(slot) |
| info = inspect_binary(self.path, scan_bytes=0) |
| self.assertEqual(info['slots'][0]['offset'], 512) |
| self.assertEqual(info['slots'][0]['arch'], 'RV32') |
| self.assertEqual(info['weight_search']['search_bytes_read'], 0) |
| self.assertFalse(info['slot_search']['uniqueness_verified']) |
|
|
| def test_deep_scan_is_explicit(self): |
| prefix, payload = blob() |
| slot = SLOT_MAGIC + struct.pack('<II', 256, 0) + bytes(272) |
| self.write(b'unknown-prefix' + prefix + payload + slot) |
| fast = inspect_binary(self.path) |
| self.assertEqual(fast['slots'], []) |
| self.assertEqual(fast['embedded_weights'], []) |
| deep = inspect_binary(self.path, deep_scan=True) |
| self.assertEqual(len(deep['slots']), 1) |
| self.assertEqual(len(deep['embedded_weights']), 1) |
| self.assertTrue(deep['slot_search']['whole_file']) |
| self.assertTrue(deep['weight_search']['complete_within_selected_ranges']) |
|
|
| def test_small_program_limit_before_read(self): |
| self.write(b'\x73\x00\x10\x00') |
| self.assertEqual(read_small_program(self.path), self.path.read_bytes()) |
| with self.path.open('wb') as f: |
| f.truncate(1 << 32) |
| with patch('os.pread', side_effect=AssertionError('oversized file payload read')): |
| with self.assertRaisesRegex(ValueError, 'inspect_binary'): |
| read_small_program(self.path) |
| for data in (b'', b'abc', b'\x7fELF', b'LNW1'): |
| self.write(data) |
| with self.assertRaisesRegex(ValueError, 'inspect_binary'): |
| read_small_program(self.path) |
|
|
| def test_minimal_elf_and_guards(self): |
| ident = b'\x7fELF' + bytes([2, 1, 1]) + bytes(9) |
| def header(**kw): |
| vals = dict(kind=2, machine=62, version=1, entry=0, phoff=0, shoff=0, flags=0, |
| ehsize=64, phsize=56, phnum=0, shsize=64, shnum=0, shstr=0) |
| vals.update(kw) |
| return ident + struct.pack('<HHIQQQIHHHHHH', *vals.values()) |
| self.write(header()) |
| with BinaryReader(self.path) as r: |
| self.assertEqual(index_elf(r)['architecture'], 'x86_64') |
| for bad in (header(phnum=1, phoff=64), header(phnum=65535), header(ehsize=1), header(shoff=64), header(phnum=8193)): |
| self.write(bad) |
| with BinaryReader(self.path) as r, self.assertRaises(BinaryFormatError): |
| index_elf(r) |
| segment = struct.pack('<IIQQQQQQ', 1, 5, 0, 0, 0, 120, 120, 4096) |
| self.write(header(phoff=64, phnum=1) + segment) |
| with BinaryReader(self.path) as r: |
| self.assertEqual(index_elf(r)['segments'][0]['nbytes'], 120) |
|
|
|
|
| if __name__ == '__main__': |
| unittest.main() |
|
|
