From 09ec33a5786cd59cb013dd736250fa58a56203d5 Mon Sep 17 00:00:00 2001
From: uuuvn <83587632+uuuvn@users.noreply.github.com>
Date: Thu, 6 Feb 2025 04:13:44 +0200
Subject: [PATCH] Better errors when relocating against undefined symbol
 (#8902)

---
 test/unit/test_elf.py           | 16 +++++++++++++---
 tinygrad/runtime/support/elf.py |  2 ++
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/test/unit/test_elf.py b/test/unit/test_elf.py
index ed8edce694..f55b11e17e 100644
--- a/test/unit/test_elf.py
+++ b/test/unit/test_elf.py
@@ -1,12 +1,13 @@
 import unittest, subprocess, platform
+from tinygrad.runtime.ops_clang import ClangJITCompiler
 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);
+      int something; // will be a load from a relocation (needed for .rela.text to exist)
+      int test(int x) {
+        return something + x;
       }
     '''
     args = ('-x', 'c', '-c', '-target', f'{platform.machine()}-none-unknown-elf', '-march=native', '-fPIC', '-O2', '-ffreestanding', '-nostdlib')
@@ -14,6 +15,15 @@ class TestElfLoader(unittest.TestCase):
     _, 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)
+  def test_clang_jit_compiler_external_raise(self):
+    src = '''
+      int evil_external_function(int);
+      int test(int x) {
+        return evil_external_function(x+2)*2;
+      }
+    '''
+    with self.assertRaisesRegex(RuntimeError, 'evil_external_function'):
+      ClangJITCompiler().compile(src)
 
 if __name__ == '__main__':
   unittest.main()
diff --git a/tinygrad/runtime/support/elf.py b/tinygrad/runtime/support/elf.py
index ffaeff921f..26495ca3ea 100644
--- a/tinygrad/runtime/support/elf.py
+++ b/tinygrad/runtime/support/elf.py
@@ -32,6 +32,8 @@ def elf_loader(blob:bytes, force_section_align:int=1) -> tuple[memoryview, list[
   for sh, trgt_sh_name, c_rels in rel + rela:
     target_image_off = next(tsh for tsh in sections if tsh.name == trgt_sh_name).header.sh_addr
     rels = [(r.r_offset, symtab[libc.ELF64_R_SYM(r.r_info)], libc.ELF64_R_TYPE(r.r_info), getattr(r, "r_addend", 0)) for r in c_rels]
+    for roff, sym, r_type_, r_addend in rels:
+      if sym.st_shndx == 0: raise RuntimeError(f'Attempting to relocate against an undefined symbol {repr(_strtab(sh_strtab, sym.st_name))}')
     relocs += [(target_image_off + roff, sections[sym.st_shndx].header.sh_addr + sym.st_value, rtype, raddend) for roff, sym, rtype, raddend in rels]
 
   return memoryview(image), sections, relocs