""" Unit tests for the rocm_kpack.elf package. Tests the new ELF surgery infrastructure against the test_assets binaries. """ import pytest import subprocess from pathlib import Path from rocm_kpack.elf import ( ElfSurgery, ElfVerifier, ProgramHeaderManager, map_section_to_load, set_pointer, conservative_zero_page, zero_page_section, verify_all, verify_with_readelf, create_load_segment, PT_LOAD, SHF_ALLOC, ) from rocm_kpack.elf.types import ( ElfHeader, ProgramHeader, ProgramHeader as PH, SectionHeader as SH, ELF64_EHDR_SIZE, ELF64_PHDR_SIZE, ELF64_SHDR_SIZE, PF_R, PF_X, SHT_NULL, SHT_STRTAB, PAGE_SIZE, page_align_offset, ) class TestElfSurgery: """Tests for the ElfSurgery class.""" def test_load_binary(self, test_assets_dir: Path): """Test loading a valid ELF binary.""" binary = test_assets_dir / "bundled_binaries/linux/cov5/test_kernel_single.exe" surgery = ElfSurgery.load(binary) assert surgery.ehdr is not None assert surgery.is_pie_or_shared # PIE executable def test_find_section(self, test_assets_dir: Path): """Test finding sections by name.""" binary = test_assets_dir / "bundled_binaries/linux/cov5/test_kernel_single.exe" surgery = ElfSurgery.load(binary) # Should find .hip_fatbin section = surgery.find_section(".hip_fatbin") assert section is not None assert section.name == ".hip_fatbin" assert section.size > 0 # Should not find nonexistent section assert surgery.find_section(".nonexistent") is None def test_find_phdr_containing_vaddr(self, test_assets_dir: Path): """Test finding program header by virtual address.""" binary = test_assets_dir / "bundled_binaries/linux/cov5/test_kernel_single.exe" surgery = ElfSurgery.load(binary) # Find the entry point's segment entry = surgery.ehdr.e_entry result = surgery.find_phdr_containing_vaddr(entry) assert result is not None idx, phdr = result assert phdr.p_type == PT_LOAD def test_iter_sections(self, test_assets_dir: Path): """Test iterating over all sections.""" binary = test_assets_dir / "bundled_binaries/linux/cov5/test_kernel_single.exe" surgery = ElfSurgery.load(binary) sections = list(surgery.iter_sections()) assert len(sections) > 0 # Check we have common sections names = [s.name for s in sections] assert ".text" in names assert ".hip_fatbin" in names class TestElfVerifier: """Tests for the ElfVerifier class.""" def test_verify_valid_binary(self, test_assets_dir: Path): """Test verification of a valid binary passes.""" binary = test_assets_dir / "bundled_binaries/linux/cov5/test_kernel_single.exe" result = ElfVerifier.verify(binary) assert result.passed def test_verify_detects_nobits_collision(self, test_assets_dir: Path): """Test that verification detects NOBITS offset collision.""" # Load a valid binary and manually inject a collision binary = test_assets_dir / "bundled_binaries/linux/cov5/test_kernel_single.exe" surgery = ElfSurgery.load(binary) # Find an existing PT_LOAD segment original_phdr = None for idx, phdr in surgery.iter_program_headers(): if phdr.p_type == PT_LOAD and phdr.p_filesz > 0: original_phdr = phdr break assert original_phdr is not None, "Need a PT_LOAD segment for this test" # Create a NOBITS segment (filesz=0) at the same offset - this is invalid collision_phdr = ProgramHeader( p_type=PT_LOAD, p_flags=original_phdr.p_flags, p_offset=original_phdr.p_offset, # Same offset! p_vaddr=original_phdr.p_vaddr + 0x100000, # Different vaddr p_paddr=original_phdr.p_paddr + 0x100000, p_filesz=0, # NOBITS - no file content p_memsz=0x1000, # But has memory size p_align=original_phdr.p_align, ) # Inject the collision surgery._phdrs.append(collision_phdr) surgery._ehdr.e_phnum += 1 # Now verify should detect the collision verifier = ElfVerifier(surgery) result = verifier.check_no_overlapping_load_segments() assert not result.passed, "Should detect NOBITS collision" assert any("NOBITS collision" in e for e in result.errors) class TestMapSectionToLoad: """Tests for the map_section_to_load operation.""" def test_map_non_alloc_section(self, test_assets_dir: Path, tmp_path: Path): """Test mapping a non-ALLOC section to PT_LOAD.""" # Use host_only binary which doesn't have .hip_fatbin binary = test_assets_dir / "bundled_binaries/linux/cov5/libhost_only.so" surgery = ElfSurgery.load(binary) # Find a non-ALLOC section to test with sections = list(surgery.iter_sections()) non_alloc = [s for s in sections if not s.is_alloc and s.size > 0] assert non_alloc, ( "Test binary must have at least one non-ALLOC section with content. " "Available sections: " + ", ".join(s.name for s in sections) ) section = non_alloc[0] result = map_section_to_load(surgery, section.name) assert result.success, f"map_section_to_load failed: {result.error}" assert result.vaddr > 0 # Verify section now has ALLOC flag updated = surgery.find_section(section.name) assert updated is not None assert updated.is_alloc class TestZeroPage: """Tests for zero-page optimization.""" def test_zero_page_hip_fatbin(self, test_assets_dir: Path, tmp_path: Path): """Test zero-paging .hip_fatbin section.""" input_path = ( test_assets_dir / "bundled_binaries/linux/cov5/test_kernel_single.exe" ) output_path = tmp_path / "zero_paged.exe" result = zero_page_section(input_path, output_path) assert result.success assert result.bytes_saved > 0 assert result.pages_zeroed > 0 assert output_path.stat().st_size < input_path.stat().st_size def test_zero_page_no_nobits_collision(self, test_assets_dir: Path, tmp_path: Path): """Test that zero-page doesn't create NOBITS offset collision.""" input_path = ( test_assets_dir / "bundled_binaries/linux/cov5/test_kernel_single.exe" ) output_path = tmp_path / "zero_paged.exe" result = zero_page_section(input_path, output_path) assert result.success # Verify with our internal verifier verify_result = ElfVerifier.verify(output_path) assert verify_result.passed, f"Errors: {verify_result.errors}" def test_zero_page_readelf_valid(self, test_assets_dir: Path, tmp_path: Path): """Test that zero-paged binary is valid according to readelf.""" input_path = ( test_assets_dir / "bundled_binaries/linux/cov5/test_kernel_single.exe" ) output_path = tmp_path / "zero_paged.exe" result = zero_page_section(input_path, output_path) assert result.success # Run readelf and check for errors verify_result = verify_with_readelf(output_path) assert verify_result.passed, f"Errors: {verify_result.errors}" def test_zero_page_small_section(self, test_assets_dir: Path, tmp_path: Path): """Test that small sections are handled gracefully.""" # Create a patched binary with small .hip_fatbin from elf_test_utils import patch_hip_fatbin_size input_path = ( test_assets_dir / "bundled_binaries/linux/cov5/test_kernel_single.exe" ) patched_path = tmp_path / "small_fatbin.exe" output_path = tmp_path / "zero_paged.exe" # Patch to 3000 bytes (less than one page) patch_hip_fatbin_size(input_path, patched_path, new_size=3000) # Should succeed but not save bytes (no full pages) result = zero_page_section(patched_path, output_path) assert result.success # No pages to zero for section < PAGE_SIZE assert result.pages_zeroed == 0 or result.bytes_saved == 0 def test_zero_page_already_nobits(self, test_assets_dir: Path, tmp_path: Path): """Test that zero-paging an already zero-paged binary is safe.""" input_path = ( test_assets_dir / "bundled_binaries/linux/cov5/test_kernel_single.exe" ) first_output = tmp_path / "first_zero_page.exe" second_output = tmp_path / "second_zero_page.exe" # First zero-page should succeed result1 = zero_page_section(input_path, first_output) assert result1.success assert result1.bytes_saved > 0 # Second zero-page should succeed but do nothing result2 = zero_page_section(first_output, second_output) assert result2.success assert result2.bytes_saved == 0 assert result2.pages_zeroed == 0 assert "already NOBITS" in (result2.error or "") # Output should not be created (or if created, should match input) # Actually, zero_page_section only saves on success with changes # Let's verify the first output is still valid verify_result = ElfVerifier.verify(first_output) assert verify_result.passed, f"Errors: {verify_result.errors}" def _build_synthetic_elf(phdr_capacity: int, phdr_count: int = 2) -> bytearray: """Build a minimal valid ELF binary with controllable spare phdr slots. Creates an ELF with: - ELF header at offset 0 - phdr table at offset 64 with room for `phdr_capacity` entries - `phdr_count` actual program headers (PT_LOAD covering ehdr+phdr, PT_LOAD for content) - A .shstrtab section and section header table - Content placed after the phdr gap so get_min_content_offset() works Such ELF files with spare phdr slots are common in practice — prior ELF transformations (notably patchelf) often enlarge the PHDR table, leaving unused slots between e_phnum and the allocated capacity. This is intentionally gross — it exists to test one specific corner case in _write_in_place() where e_phnum < phdr_capacity. """ assert phdr_count <= phdr_capacity phdr_table_size = phdr_capacity * ELF64_PHDR_SIZE # Content starts right after the phdr region (page-aligned) content_offset = ELF64_EHDR_SIZE + phdr_table_size # Pad to 8-byte alignment for section data content_offset = (content_offset + 7) & ~7 # Section string table: "\0.shstrtab\0" shstrtab_data = b"\x00.shstrtab\x00" shstrtab_offset = content_offset shstrtab_size = len(shstrtab_data) # Some dummy content to act as a loadable segment dummy_content = b"\xcc" * 64 # int3 padding dummy_offset = shstrtab_offset + shstrtab_size dummy_offset = (dummy_offset + 7) & ~7 # align # Section headers go after all content shdr_offset = dummy_offset + len(dummy_content) shdr_offset = (shdr_offset + 7) & ~7 # align # 3 section headers: SHT_NULL, .shstrtab, (we keep it minimal) num_shdrs = 2 total_size = shdr_offset + num_shdrs * ELF64_SHDR_SIZE data = bytearray(total_size) # -- ELF Header -- e_ident = bytearray(16) e_ident[0:4] = b"\x7fELF" e_ident[4] = 2 # ELFCLASS64 e_ident[5] = 1 # ELFDATA2LSB e_ident[6] = 1 # EV_CURRENT ehdr = ElfHeader( e_ident=bytes(e_ident), e_type=3, # ET_DYN (shared/PIE) e_machine=0x3E, # EM_X86_64 e_version=1, e_entry=0, e_phoff=ELF64_EHDR_SIZE, e_shoff=shdr_offset, e_flags=0, e_ehsize=ELF64_EHDR_SIZE, e_phentsize=ELF64_PHDR_SIZE, e_phnum=phdr_count, e_shentsize=ELF64_SHDR_SIZE, e_shnum=num_shdrs, e_shstrndx=1, # .shstrtab is section 1 ) ehdr.write_to(data, 0) # -- Program headers -- # PT_LOAD 0: covers ELF header + phdr table (vaddr 0x0) ph0 = PH( p_type=PT_LOAD, p_flags=PF_R, p_offset=0, p_vaddr=0x0, p_paddr=0x0, p_filesz=content_offset, p_memsz=content_offset, p_align=PAGE_SIZE, ) ph0.write_to(data, ELF64_EHDR_SIZE) # PT_LOAD 1: covers dummy content ph1 = PH( p_type=PT_LOAD, p_flags=PF_R | PF_X, p_offset=dummy_offset, p_vaddr=PAGE_SIZE + dummy_offset, # at next page + same page offset p_paddr=PAGE_SIZE + dummy_offset, p_filesz=len(dummy_content), p_memsz=len(dummy_content), p_align=PAGE_SIZE, ) ph1.write_to(data, ELF64_EHDR_SIZE + ELF64_PHDR_SIZE) # Remaining phdr slots are zero (PT_NULL) — this is valid # -- Section content -- data[shstrtab_offset : shstrtab_offset + shstrtab_size] = shstrtab_data data[dummy_offset : dummy_offset + len(dummy_content)] = dummy_content # -- Section headers -- # Section 0: SHT_NULL (required) sh_null = SH( sh_name=0, sh_type=SHT_NULL, sh_flags=0, sh_addr=0, sh_offset=0, sh_size=0, sh_link=0, sh_info=0, sh_addralign=0, sh_entsize=0, ) sh_null.write_to(data, shdr_offset) # Section 1: .shstrtab sh_strtab = SH( sh_name=1, # offset of ".shstrtab" in shstrtab_data sh_type=SHT_STRTAB, sh_flags=0, sh_addr=0, sh_offset=shstrtab_offset, sh_size=shstrtab_size, sh_link=0, sh_info=0, sh_addralign=1, sh_entsize=0, ) sh_strtab.write_to(data, shdr_offset + ELF64_SHDR_SIZE) return data class TestProgramHeaderManager: """Tests for ProgramHeaderManager.""" def test_add_program_header(self, test_assets_dir: Path): """Test adding a program header.""" binary = test_assets_dir / "bundled_binaries/linux/cov5/test_kernel_single.exe" surgery = ElfSurgery.load(binary) original_count = surgery.ehdr.e_phnum manager = ProgramHeaderManager(surgery) # Allocate vaddr first, then align file offset to match vaddr = manager.allocate_vaddr() file_offset = page_align_offset(len(surgery.data), vaddr) # Add a dummy PT_LOAD new_phdr = create_load_segment( vaddr=vaddr, file_offset=file_offset, size=0x1000, ) manager.add_program_header(new_phdr) assert len(manager.program_headers) == original_count + 1 def test_allocate_vaddr(self, test_assets_dir: Path): """Test virtual address allocation.""" binary = test_assets_dir / "bundled_binaries/linux/cov5/test_kernel_single.exe" surgery = ElfSurgery.load(binary) manager = ProgramHeaderManager(surgery) vaddr = manager.allocate_vaddr() # Should be page-aligned and after all existing segments assert vaddr % 0x1000 == 0 assert vaddr >= manager.get_max_vaddr() def test_apply_in_place_with_added_header(self, tmp_path: Path): """Adding a phdr when spare slots exist must write in place. Regression test: _write_in_place() used to call surgery.update_program_header(i) for the new index before updating e_phnum, causing 'Invalid program header index' on binaries where phdr capacity > phdr count (e.g. librccl.so with 13 headers and 418 available slots). """ # Build a synthetic ELF: 2 phdrs used, room for 8 data = _build_synthetic_elf(phdr_capacity=8, phdr_count=2) surgery = ElfSurgery(data) assert surgery.ehdr.e_phnum == 2 manager = ProgramHeaderManager(surgery) # Verify spare capacity exists (should not need relocation) capacity = manager._get_current_capacity() assert capacity >= 3, f"Expected spare slots, got capacity={capacity}" # Add a new PT_LOAD vaddr = manager.allocate_vaddr() file_offset = page_align_offset(len(surgery.data), vaddr) new_phdr = create_load_segment( vaddr=vaddr, file_offset=file_offset, size=0x1000 ) manager.add_program_header(new_phdr) # This is the line that fails without the fix: # ValueError: Invalid program header index: 2 result = manager.apply() assert not result.relocated, "Should have written in place, not relocated" assert surgery.ehdr.e_phnum == 3 # Round-trip: save, reload, verify the new header persists output = tmp_path / "synthetic.so" surgery.save(output) surgery2 = ElfSurgery.load(output) assert surgery2.ehdr.e_phnum == 3 # Verify the new phdr was actually written phdrs2 = list(surgery2.iter_program_headers()) assert len(phdrs2) == 3 _, written_phdr = phdrs2[2] assert written_phdr.p_type == PT_LOAD assert written_phdr.p_vaddr == vaddr class TestVerifyAll: """Tests for the verify_all function.""" def test_verify_all_valid_binary(self, test_assets_dir: Path, tmp_path: Path): """Test that verify_all passes on valid binary.""" binary = test_assets_dir / "bundled_binaries/linux/cov5/test_kernel_single.exe" result = verify_all(binary, tmp_path) assert result.passed, f"Errors: {result.errors}" class TestAddSection: """Tests for ElfSurgery.add_section().""" def test_add_section_basic(self, tmp_path: Path, test_assets_dir: Path): """Test adding a new section to an ELF binary.""" binary = test_assets_dir / "bundled_binaries/linux/cov5/test_kernel_single.exe" output = tmp_path / "with_section.exe" surgery = ElfSurgery.load(binary) # Add a new section content = b"test section content" result = surgery.add_section( name=".test_section", content=content, ) # Verify result assert result.index > 0 assert result.offset > 0 # Save and reload surgery.save(output) surgery2 = ElfSurgery.load(output) # Find the section section = surgery2.find_section(".test_section") assert section is not None assert surgery2.get_section_content(section) == content def test_add_section_preserves_existing( self, tmp_path: Path, test_assets_dir: Path ): """Test that adding a section preserves existing sections.""" binary = test_assets_dir / "bundled_binaries/linux/cov5/test_kernel_multi.exe" output = tmp_path / "with_section.exe" surgery = ElfSurgery.load(binary) # Get original section count original_section_count = surgery.ehdr.e_shnum # Verify .hip_fatbin exists fatbin = surgery.find_section(".hip_fatbin") assert fatbin is not None original_fatbin_size = fatbin.header.sh_size # Add a new section surgery.add_section(name=".my_section", content=b"hello world") surgery.save(output) # Reload and verify surgery2 = ElfSurgery.load(output) assert surgery2.ehdr.e_shnum == original_section_count + 1 # Verify .hip_fatbin preserved fatbin2 = surgery2.find_section(".hip_fatbin") assert fatbin2 is not None assert fatbin2.header.sh_size == original_fatbin_size # Verify new section exists new_section = surgery2.find_section(".my_section") assert new_section is not None assert surgery2.get_section_content(new_section) == b"hello world" def test_add_section_with_flags(self, tmp_path: Path, test_assets_dir: Path): """Test adding a section with specific flags.""" from rocm_kpack.elf import SHT_NOTE, SHF_ALLOC binary = test_assets_dir / "bundled_binaries/linux/cov5/host_only.exe" output = tmp_path / "with_note.exe" surgery = ElfSurgery.load(binary) surgery.add_section( name=".my_note", content=b"note content here", section_type=SHT_NOTE, flags=SHF_ALLOC, addralign=4, ) surgery.save(output) surgery2 = ElfSurgery.load(output) section = surgery2.find_section(".my_note") assert section is not None assert section.header.sh_type == SHT_NOTE assert section.header.sh_flags & SHF_ALLOC assert section.header.sh_addralign == 4