ELF loader strtab fix and tests (#6011)

* ELF loader strtab fix and tests

* ruff

* typos

* only one test

test/test_elf.py

@@ -0,0 +1,19 @@
+import unittest, subprocess, platform
+from tinygrad.runtime.support.elf import elf_loader
+
+class TestElfLoader(unittest.TestCase):
+  def test_load_clang_jit_strtab(self):
+    src = '''
+      void relocation(int); // will be a jump to relocation (needed for .rela.text to exist)
+      void test(int x) {
+        relocation(x+1);
+      }
+    '''
+    args = ('-x', 'c', '-c', '-target', f'{platform.machine()}-none-unknown-elf', '-march=native', '-fPIC', '-O2', '-ffreestanding', '-nostdlib')
+    obj = subprocess.check_output(('clang',) + args + ('-', '-o', '-'), input=src.encode('utf-8'))
+    _, sections, _ = elf_loader(obj)
+    section_names = [sh.name for sh in sections]
+    assert '.text' in section_names and '.rela.text' in section_names, str(section_names)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad/runtime/support/elf.py

@@ -7,12 +7,12 @@ import tinygrad.runtime.autogen.libc as libc
 class ElfSection: name:str; header:libc.Elf64_Shdr; content:bytes # noqa: E702
 
 def elf_loader(blob:bytes, force_section_align:int=1) -> Tuple[memoryview, List[ElfSection], Any]:
-  def _elf_parse_names(tabs): return {sum(len(w) + 1 for w in tabs.split(b'\0')[:i]): w.decode('utf-8') for i, w in enumerate(tabs.split(b'\0'))}
+  def _strtab(blob: bytes, idx: int) -> str: return blob[idx:blob.find(b'\x00', idx)].decode('utf-8')
 
   header = libc.Elf64_Ehdr.from_buffer_copy(blob)
   section_headers = (libc.Elf64_Shdr * header.e_shnum).from_buffer_copy(blob[header.e_shoff:])
-  section_names = _elf_parse_names(blob[(shstrst:=section_headers[header.e_shstrndx].sh_offset):shstrst+section_headers[header.e_shstrndx].sh_size])
-  sections = [ElfSection(section_names[sh.sh_name], sh, blob[sh.sh_offset:sh.sh_offset+sh.sh_size]) for sh in section_headers]
+  sh_strtab = blob[(shstrst:=section_headers[header.e_shstrndx].sh_offset):shstrst+section_headers[header.e_shstrndx].sh_size]
+  sections = [ElfSection(_strtab(sh_strtab, sh.sh_name), sh, blob[sh.sh_offset:sh.sh_offset+sh.sh_size]) for sh in section_headers]
 
   def _to_carray(sh, ctype): return (ctype * (sh.header.sh_size // sh.header.sh_entsize)).from_buffer_copy(sh.content)
   rel = [(sh, sh.name[4:], _to_carray(sh, libc.Elf64_Rel)) for sh in sections if sh.header.sh_type == libc.SHT_REL]