fix(profiling): I/O profiling gotchas

profiling/src/io/got_elf64.rs → profiling/src/io/got_elf64/mod.rs

@@ -5,10 +5,14 @@ use crate::profiling::bindings::{
     Elf64_Dyn, Elf64_Rela, Elf64_Sym, Elf64_Xword, DT_JMPREL, DT_NULL, DT_PLTRELSZ, DT_STRTAB,
     DT_SYMTAB, PT_DYNAMIC, PT_LOAD, R_AARCH64_JUMP_SLOT, R_X86_64_JUMP_SLOT,
 };
+use core::ffi::CStr;
+use core::ptr;
 use libc::{c_char, c_int, c_void, dl_phdr_info};
 use log::{error, trace};
-use std::ffi::CStr;
-use std::ptr;
+use std::sync::atomic::{AtomicUsize, Ordering::Relaxed};
+
+#[cfg(test)]
+mod tests;
 
 fn elf64_r_type(info: Elf64_Xword) -> u32 {
     (info & 0xffffffff) as u32
@@ -18,6 +22,73 @@ fn elf64_r_sym(info: Elf64_Xword) -> u32 {
     (info >> 32) as u32
 }
 
+/// Temporarily make a RELRO GOT page writable. Ordinary writable GOT pages
+/// require no permission change. Explicit restoration reports failures; Drop
+/// also restores protection if a caller exits before reaching that step.
+struct GotPageProtection {
+    page: Option<(*mut c_void, usize)>,
+}
+
+impl GotPageProtection {
+    unsafe fn make_writable(info: *mut dl_phdr_info, slot: usize) -> Option<Self> {
+        let mut is_relro = false;
+        for i in 0..(*info).dlpi_phnum {
+            let phdr = &*(*info).dlpi_phdr.add(i as usize);
+            if phdr.p_type == libc::PT_GNU_RELRO {
+                let start = (*info).dlpi_addr as usize + phdr.p_vaddr as usize;
+                let end = start + phdr.p_memsz as usize;
+                if (start..end).contains(&slot) {
+                    is_relro = true;
+                    break;
+                }
+            }
+        }
+        if !is_relro {
+            return Some(Self { page: None });
+        }
+
+        let page_size = libc::sysconf(libc::_SC_PAGESIZE);
+        if page_size <= 0 {
+            error!("Failed to determine page size while changing GOT protection");
+            return None;
+        }
+        let page_size = page_size as usize;
+        let page = (slot & !(page_size - 1)) as *mut c_void;
+        if libc::mprotect(page, page_size, libc::PROT_READ | libc::PROT_WRITE) != 0 {
+            let err = *libc::__errno_location();
+            error!("Failed to make RELRO GOT page writable at {page:p}: {err}");
+            return None;
+        }
+        Some(Self {
+            page: Some((page, page_size)),
+        })
+    }
+
+    fn restore(mut self) -> bool {
+        self.restore_inner()
+    }
+
+    fn restore_inner(&mut self) -> bool {
+        let Some((page, page_size)) = self.page.take() else {
+            return true;
+        };
+        // SAFETY: make_writable obtained this page from a loaded image's GOT.
+        // The caller must keep the image loaded until this guard is restored.
+        if unsafe { libc::mprotect(page, page_size, libc::PROT_READ) } != 0 {
+            let err = unsafe { *libc::__errno_location() };
+            error!("Failed to restore RELRO GOT page protection at {page:p}: {err}");
+            return false;
+        }
+        true
+    }
+}
+
+impl Drop for GotPageProtection {
+    fn drop(&mut self) {
+        self.restore_inner();
+    }
+}
+
 /// Override the GOT entry for symbols specified in `overwrites`.
 ///
 /// See: https://cs4401.walls.ninja/notes/lecture/basics_global_offset_table.html
@@ -35,16 +106,18 @@ unsafe fn override_got_entry(
 ) -> bool {
     let phdr = (*info).dlpi_phdr;
 
-    // Locate the dynamic programm header (`PT_DYNAMIC`)
+    // Locate the dynamic program header (`PT_DYNAMIC`).
     let mut dyn_ptr: *const Elf64_Dyn = ptr::null();
+    let mut dyn_count: usize = 0;
     for i in 0..(*info).dlpi_phnum {
         let phdr_i = phdr.offset(i as isize);
         if (*phdr_i).p_type == PT_DYNAMIC {
             dyn_ptr = ((*info).dlpi_addr as usize + (*phdr_i).p_vaddr as usize) as *const Elf64_Dyn;
+            dyn_count = (*phdr_i).p_memsz as usize / std::mem::size_of::<Elf64_Dyn>();
             break;
         }
     }
-    if dyn_ptr.is_null() {
+    if dyn_ptr.is_null() || dyn_count == 0 {
         trace!("Failed to locate dynamic section");
         return false;
     }
@@ -63,7 +136,7 @@ unsafe fn override_got_entry(
     // - on glibc, addresses are absolutes
     // https://elixir.bootlin.com/glibc/glibc-2.36/source/elf/get-dynamic-info.h#L84
     let mut dyn_iter = dyn_ptr;
-    loop {
+    for _ in 0..dyn_count {
         let d_tag = (*dyn_iter).d_tag as u32;
         if d_tag == DT_NULL {
             break;
@@ -112,77 +185,65 @@ unsafe fn override_got_entry(
 
     let num_relocs = rel_plt_size / std::mem::size_of::<Elf64_Rela>();
 
-    // For each symbol we want to overwrite (from `overwrites`), we scan the relocation entries.
-    // Once the matching symbol name is found, patch its GOT entry to point to our new function.
-    for overwrite in state.overwrites.iter_mut() {
-        for i in 0..num_relocs {
-            let rel = rel_plt.add(i);
-            let r_type = elf64_r_type((*rel).r_info);
+    // Scan relocation entries once and match against symbols we want to overwrite.
+    for i in 0..num_relocs {
+        let rel = rel_plt.add(i);
+        let r_type = elf64_r_type((*rel).r_info);
 
-            // Only handle JUMP_SLOT relocations
-            if r_type != R_AARCH64_JUMP_SLOT && r_type != R_X86_64_JUMP_SLOT {
-                continue;
-            }
+        // Only handle JUMP_SLOT relocations
+        if r_type != R_AARCH64_JUMP_SLOT && r_type != R_X86_64_JUMP_SLOT {
+            continue;
+        }
 
-            // Get the symbol index for this relocation, then the symbol struct
-            let sym_index = elf64_r_sym((*rel).r_info) as usize;
-            let sym = symtab.add(sym_index);
+        // Get the symbol index for this relocation, then the symbol struct
+        let sym_index = elf64_r_sym((*rel).r_info) as usize;
+        let sym = symtab.add(sym_index);
 
-            // Access the symbol name via the string table
-            let name_offset = (*sym).st_name as isize;
-            let name_ptr = strtab.offset(name_offset);
-            let name = CStr::from_ptr(name_ptr).to_str().unwrap_or("");
+        // Access the symbol name via the string table
+        let name_offset = (*sym).st_name as isize;
+        let name_ptr = strtab.offset(name_offset);
+        let name = CStr::from_ptr(name_ptr).to_str().unwrap_or("");
 
+        for overwrite in state.overwrites.iter_mut() {
             if name == overwrite.symbol_name {
                 // Calculate the GOT entry address. Per the ELF spec, `r_offset` for pointer-sized
                 // relocations (such as GOT entries) is guaranteed to be pointer-aligned, see:
                 // https://github.com/ARM-software/abi-aa/blob/main/aaelf64/aaelf64.rst#5733relocation-operations
                 let got_entry =
                     ((*info).dlpi_addr as usize + (*rel).r_offset as usize) as *mut *mut ();
 
-                // Change memory protection so we can write to the GOT entry
-                let page_size = libc::sysconf(libc::_SC_PAGESIZE) as usize;
-                let aligned_addr = (got_entry as usize) & !(page_size - 1);
-                if libc::mprotect(
-                    aligned_addr as *mut c_void,
-                    page_size,
-                    libc::PROT_READ | libc::PROT_WRITE,
-                ) != 0
-                {
-                    let err = *libc::__errno_location();
-                    trace!("mprotect failed: {}", err);
-                    return false;
-                }
-
                 let original = *got_entry;
                 if original == overwrite.new_func {
                     continue;
                 }
 
                 trace!(
-                    "Overriding GOT entry for {} at offset {:?} (abs: {:p}) pointing to {:p} (orig function at {:p})",
+                    "Overriding GOT entry for {} at offset {:?} (abs: {:p}) pointing to {:p}",
                     overwrite.symbol_name,
                     (*rel).r_offset,
                     got_entry,
                     original,
-                    *overwrite.orig_func
                 );
 
-                // This works for musl based linux distros, but not for libc once
-                *overwrite.orig_func = libc::dlsym(libc::RTLD_NEXT, name_ptr) as *mut ();
-                if (*overwrite.orig_func).is_null() {
-                    // libc linux fallback
-                    *overwrite.orig_func = original;
-                }
-                state.restores.push(GotSlotRestore {
+                // Allocate restore bookkeeping before making the page writable.
+                state.restores.reserve(1);
+                let restore = GotSlotRestore {
                     image: (*info).dlpi_addr as usize,
                     image_name: image_name.into(),
                     slot: got_entry as usize,
                     original: original as usize,
                     replacement: overwrite.new_func as usize,
-                });
+                };
+                let Some(protection) = GotPageProtection::make_writable(info, got_entry as usize)
+                else {
+                    return false;
+                };
+                state.restores.push(restore);
                 *got_entry = overwrite.new_func;
-                continue;
+                if !protection.restore() {
+                    return false;
+                }
+                break;
             }
         }
     }
@@ -198,13 +259,28 @@ pub unsafe extern "C" fn callback(
 ) -> c_int {
     let state = &mut *(data as *mut GotHookState);
 
-    // detect myself ...
-    let mut my_info: libc::Dl_info = std::mem::zeroed();
-    if libc::dladdr(callback as *const c_void, &mut my_info) == 0 {
-        error!("Did not find my own `dladdr` and therefore can't hook into the GOT.");
+    // 0 = uninitialized, usize::MAX = failure (too high to be an image base)
+    static MY_BASE_ADDR: AtomicUsize = AtomicUsize::new(0);
+    let mut my_base_addr = MY_BASE_ADDR.load(Relaxed);
+    if my_base_addr == 0 {
+        let mut my_info: libc::Dl_info = unsafe { std::mem::zeroed() };
+        let resolved = if unsafe { libc::dladdr(callback as *const c_void, &mut my_info) } == 0
+            || my_info.dli_fbase.is_null()
+        {
+            error!("Did not find my own `dladdr` and therefore can't hook into the GOT.");
+            usize::MAX
+        } else {
+            my_info.dli_fbase as usize
+        };
+        // Concurrent lookups are harmless; use whichever result was published first.
+        my_base_addr = match MY_BASE_ADDR.compare_exchange(0, resolved, Relaxed, Relaxed) {
+            Ok(_) => resolved,
+            Err(cached) => cached,
+        };
+    }
+    if my_base_addr == usize::MAX {
         return 0;
     }
-    let my_base_addr = my_info.dli_fbase as usize;
     let module_base_addr = (*info).dlpi_addr as usize;
     if module_base_addr == my_base_addr {
         // "this" lib is actually me: skipping GOT hooking for myself
@@ -222,9 +298,9 @@ pub unsafe extern "C" fn callback(
         std::str::from_utf8(image_name).unwrap_or("[Unknown]")
     };
 
-    // I guess if we try to hook into GOT from `linux-vdso` or `ld-linux` our best outcome will be
-    // that nothing happens, but most likely we'll crash and we should avoid that.
-    if name.contains("linux-vdso") || name.contains("ld-linux") {
+    // I guess if we try to hook into GOT from `linux-vdso`, `ld-linux` or `ld-musl` our best
+    // outcome will be that nothing happens, but most likely we'll crash and we should avoid that.
+    if name.contains("linux-vdso") || name.contains("ld-linux") || name.contains("ld-musl") {
         return 0;
     }
 
@@ -286,7 +362,6 @@ unsafe extern "C" fn restore_callback(
     } else {
         CStr::from_ptr((*info).dlpi_name).to_bytes()
     };
-    let page_size = libc::sysconf(libc::_SC_PAGESIZE) as usize;
 
     for restore in restores
         .iter()
@@ -311,22 +386,17 @@ unsafe extern "C" fn restore_callback(
             continue;
         }
 
-        let aligned_addr = restore.slot & !(page_size - 1);
-        if libc::mprotect(
-            aligned_addr as *mut c_void,
-            page_size,
-            libc::PROT_READ | libc::PROT_WRITE,
-        ) != 0
-        {
-            let err = *libc::__errno_location();
-            trace!("mprotect failed while restoring GOT entry at {slot:p}: {err}");
+        let Some(protection) = GotPageProtection::make_writable(info, restore.slot) else {
             *complete = false;
             continue;
-        }
-
-        if restore_slot_if_owned(restore) {
+        };
+        let restored = restore_slot_if_owned(restore);
+        // Restore protection even if the slot is no longer ours.
+        let protected = protection.restore();
+        if restored {
             trace!("Restored GOT entry at {slot:p}");
-        } else {
+        }
+        if !restored || !protected {
             *complete = false;
         }
     }