From b9e05057158f96bf03840a6cf223e97015d8b6d9 Mon Sep 17 00:00:00 2001
From: Anhad Singh <andypython@protonmail.com>
Date: Mon, 13 Jan 2025 23:48:03 +1100
Subject: [PATCH] fix(ld.so): less locks

Acquire the global scope lock at the start of relocate.

Signed-off-by: Anhad Singh <andypython@protonmail.com>
---
 src/ld_so/dso.rs    | 40 +++++++++++++++++++++-------------------
 src/ld_so/linker.rs | 10 ++--------
 2 files changed, 23 insertions(+), 27 deletions(-)

diff --git a/src/ld_so/dso.rs b/src/ld_so/dso.rs
index bd1630d01a..9be19219c8 100644
--- a/src/ld_so/dso.rs
+++ b/src/ld_so/dso.rs
@@ -23,6 +23,7 @@ use crate::{
 use alloc::{
     collections::BTreeMap,
     string::{String, ToString},
+    sync::Arc,
     vec::Vec,
 };
 use core::{
@@ -138,6 +139,13 @@ impl<'data> Dynamic<'data> {
         let name = sym.name(NativeEndian, self.dynstrtab).ok()?;
         Some(core::str::from_utf8(name).expect("non UTF-8 ELF symbol name"))
     }
+
+    fn static_relocations(&self) -> impl Iterator<Item = Relocation> + '_ {
+        self.rela
+            .iter()
+            .map(Relocation::from)
+            .chain(self.rel.iter().map(Relocation::from))
+    }
 }
 
 unsafe impl Send for Dynamic<'_> {}
@@ -671,16 +679,13 @@ impl DSO {
         ))
     }
 
-    fn static_relocate(&self, reloc: Relocation) -> object::Result<()> {
+    fn static_relocate(&self, global_scope: &Scope, reloc: Relocation) -> object::Result<()> {
         let b = self.mmap.as_ptr() as usize;
 
         let sym = if reloc.sym.0 > 0 {
             let name = self.dynamic.symbol_name(reloc.sym).unwrap();
 
-            GLOBAL_SCOPE
-                .read()
-                .get_sym(name)
-                .or_else(|| self.scope.get_sym(name))
+            resolve_sym(name, &[global_scope, &self.scope])
                 .map(|(sym, _, obj)| (sym, obj.tls_offset))
         } else {
             None
@@ -740,7 +745,7 @@ impl DSO {
         Ok(())
     }
 
-    fn lazy_relocate(&self, resolve: Resolve) -> object::Result<()> {
+    fn lazy_relocate(&self, global_scope: &Scope, resolve: Resolve) -> object::Result<()> {
         let Some(got) = self.got() else {
             assert_eq!(self.dynamic.jmprel, 0);
             return Ok(());
@@ -786,10 +791,7 @@ impl DSO {
                 (elf::R_X86_64_JUMP_SLOT, Resolve::Now) => {
                     let name = self.dynamic.symbol_name(reloc.sym).unwrap();
 
-                    let resolved = GLOBAL_SCOPE
-                        .read()
-                        .get_sym(name)
-                        .or_else(|| self.scope.get_sym(name))
+                    let resolved = resolve_sym(name, &[global_scope, &self.scope])
                         .map(|(sym, _, _)| sym.as_ptr() as usize)
                         .expect("unresolved symbol");
 
@@ -808,17 +810,13 @@ impl DSO {
     }
 
     pub fn relocate(&self, ph: &[ProgramHeader], resolve: Resolve) -> object::Result<()> {
-        self.dynamic
-            .rela
-            .iter()
-            .try_for_each(|reloc| self.static_relocate(reloc.into()))?;
+        let global_scope = GLOBAL_SCOPE.read();
 
         self.dynamic
-            .rel
-            .iter()
-            .try_for_each(|reloc| self.static_relocate(reloc.into()))?;
+            .static_relocations()
+            .try_for_each(|reloc| self.static_relocate(&global_scope, reloc))?;
 
-        self.lazy_relocate(resolve)?;
+        self.lazy_relocate(&global_scope, resolve)?;
 
         // Protect pages
         for ph in ph
@@ -864,7 +862,7 @@ impl Drop for DSO {
     }
 }
 
-pub fn is_pie_enabled(elf: &ElfFile) -> bool {
+fn is_pie_enabled(elf: &ElfFile) -> bool {
     elf.elf_header().e_type.get(elf.endian()) == elf::ET_DYN
 }
 
@@ -877,3 +875,7 @@ fn dirname(path: &str) -> String {
     parts.truncate(parts.len() - 1);
     parts.join("/")
 }
+
+pub fn resolve_sym(name: &str, scopes: &[&Scope]) -> Option<(Symbol, SymbolBinding, Arc<DSO>)> {
+    scopes.iter().find_map(|scope| scope.get_sym(name))
+}
diff --git a/src/ld_so/linker.rs b/src/ld_so/linker.rs
index 61eba61f9a..d44181a652 100644
--- a/src/ld_so/linker.rs
+++ b/src/ld_so/linker.rs
@@ -24,7 +24,7 @@ use crate::{
         fcntl, sys_mman,
         unistd::F_OK,
     },
-    ld_so::dso::SymbolBinding,
+    ld_so::dso::{resolve_sym, SymbolBinding},
     platform::{
         types::{c_int, c_uint, c_void},
         Pal, Sys,
@@ -879,13 +879,6 @@ impl Linker {
 //
 // FIXME(andypython): 32-bit
 extern "C" fn __plt_resolve_inner(obj: *const DSO, relocation_index: c_uint) -> *mut c_void {
-    let resolve_sym = |name: &str, scopes: &[&Scope]| -> Option<Symbol> {
-        scopes
-            .iter()
-            .find_map(|scope| scope.get_sym(name))
-            .map(|(sym, _, _)| sym)
-    };
-
     let obj = unsafe { &*obj };
     let obj_base = obj.mmap.as_ptr() as usize;
     let jmprel = obj.dynamic.jmprel;
@@ -908,6 +901,7 @@ extern "C" fn __plt_resolve_inner(obj: *const DSO, relocation_index: c_uint) ->
     .expect("non utf8 symbol name");
 
     let resolved = resolve_sym(name, &[&GLOBAL_SCOPE.read(), &obj.scope])
+        .map(|(sym, _, _)| sym)
         .unwrap_or_else(|| panic!("symbol '{name}' not found"))
         .as_ptr();