dlopen_rs/core_impl/
register.rs

1use crate::{
2    ElfLibrary, OpenFlags,
3    core_impl::loader::{DylibExt, LoadedDylib},
4};
5use alloc::{
6    borrow::ToOwned,
7    collections::{btree_set::BTreeSet, vec_deque::VecDeque},
8    string::String,
9    sync::Arc,
10    vec::Vec,
11};
12use core::ffi::{c_int, c_void};
13use indexmap::IndexMap;
14use spin::{Lazy, RwLock};
15
16#[macro_export]
17macro_rules! lock_write {
18    ($lock:expr) => {{ $lock.write() }};
19}
20
21#[macro_export]
22macro_rules! lock_read {
23    ($lock:expr) => {{ $lock.read() }};
24}
25
26impl Drop for ElfLibrary {
27    fn drop(&mut self) {
28        let mut removed_libs = Vec::new();
29        {
30            let mut lock = lock_write!(MANAGER);
31            let shortname = self.inner.shortname();
32            let Some(entry) = lock.get(shortname) else {
33                return;
34            };
35
36            if entry.flags.is_nodelete() {
37                return;
38            }
39
40            let ref_count = unsafe { self.inner.core_ref().strong_count() };
41            let has_global = entry.flags.is_global();
42            // Dylib ref in 'all' map + dylib ref in 'deps' list of itself + global ref (if present) + handle's 'inner' ref
43            debug_assert!(self.deps.is_some());
44            let threshold = 3 + has_global as usize;
45
46            log::debug!(
47                "Drop ElfLibrary [{}], ref count: {}, threshold: {}",
48                self.inner.name(),
49                ref_count,
50                threshold
51            );
52
53            if ref_count == threshold {
54                log::info!("Destroying dylib [{}]", self.inner.name());
55                removed_libs.push(self.inner.clone());
56
57                lock.remove(shortname);
58
59                // Check dependencies
60                if let Some(deps) = self.deps.as_ref() {
61                    for dep in deps.iter().skip(1) {
62                        let dep_shortname = dep.shortname();
63                        let Some(dep_lib) = lock.get(dep_shortname) else {
64                            continue;
65                        };
66                        if dep_lib.flags.is_nodelete() {
67                            continue;
68                        }
69                        debug_assert!(
70                            dep_lib.deps.is_some(),
71                            "Dependency [{}] must have its deps set",
72                            dep.name()
73                        );
74                        // Dylib ref in 'all' map + dylib ref in 'deps' list of itself + global ref (if present)
75                        let dep_threshold = 3 + dep_lib.flags.is_global() as usize;
76
77                        if unsafe { dep.core_ref().strong_count() } == dep_threshold {
78                            log::info!("Destroying dylib [{}]", dep.name());
79                            removed_libs.push(dep.clone());
80                            lock.remove(dep_shortname);
81                        }
82                    }
83                }
84            }
85        }
86        for lib in removed_libs {
87            let base = lib.base();
88            let range = base..(base + lib.mapped_len());
89            finalize(base as *mut _, Some(range));
90        }
91    }
92}
93
94/// Represents the current lifecycle state of a dynamic library.
95///
96/// The state uses a compact u8 representation:
97/// - `[0, 254)`: Newly loaded library, value is the index in the current loading batch.
98/// - `254`: Currently undergoing relocation.
99/// - `255`: Successfully relocated and ready for use.
100#[derive(Clone, Copy, Default)]
101pub(crate) struct DylibState(u8);
102
103impl DylibState {
104    const RELOCATED: u8 = 0b11111111;
105    const RELOCATING: u8 = 0b11111110;
106
107    /// Returns true if the library is fully relocated.
108    #[inline]
109    pub(crate) fn is_relocated(&self) -> bool {
110        self.0 == Self::RELOCATED
111    }
112
113    /// If the library is in a "new" state, returns its batch index.
114    #[inline]
115    pub(crate) fn get_new_idx(&self) -> Option<u8> {
116        if self.0 >= Self::RELOCATING {
117            None
118        } else {
119            Some(self.0)
120        }
121    }
122
123    /// Transitions the state to Relocated.
124    #[inline]
125    pub(crate) fn set_relocated(&mut self) -> &mut Self {
126        self.0 = Self::RELOCATED;
127        self
128    }
129
130    /// Transitions the state to Relocating.
131    #[inline]
132    pub(crate) fn set_relocating(&mut self) -> &mut Self {
133        self.0 = Self::RELOCATING;
134        self
135    }
136
137    /// Sets the state to a "new" library with the given batch index.
138    #[allow(unused)]
139    #[inline]
140    pub(crate) fn set_new_idx(&mut self, idx: u8) -> &mut Self {
141        assert!(idx < Self::RELOCATING);
142        self.0 = idx;
143        self
144    }
145}
146
147/// A handle to a dynamic library within the global registry.
148///
149/// This struct wraps a `RelocatedDylib` with additional metadata such as
150/// loading flags, computed dependency scope, and lifecycle state.
151#[derive(Clone)]
152pub(crate) struct GlobalDylib {
153    inner: LoadedDylib,
154    pub(crate) flags: OpenFlags,
155    /// The full dependency scope (Searchlist) used for symbol resolution,
156    /// calculated via BFS starting from this library.
157    pub(crate) deps: Option<Arc<[LoadedDylib]>>,
158    /// Current state in the load/relocate lifecycle.
159    pub(crate) state: DylibState,
160}
161
162unsafe impl Send for GlobalDylib {}
163unsafe impl Sync for GlobalDylib {}
164
165impl GlobalDylib {
166    #[inline]
167    pub(crate) fn get_lib(&self) -> ElfLibrary {
168        debug_assert!(self.deps.is_some());
169        ElfLibrary {
170            inner: self.inner.clone(),
171            deps: self.deps.clone(),
172        }
173    }
174
175    #[inline]
176    pub(crate) fn dylib(&self) -> LoadedDylib {
177        self.inner.clone()
178    }
179
180    #[inline]
181    pub(crate) fn dylib_ref(&self) -> &LoadedDylib {
182        &self.inner
183    }
184
185    #[inline]
186    pub(crate) fn shortname(&self) -> &str {
187        self.inner.shortname()
188    }
189
190    #[inline]
191    pub(crate) fn set_deps(&mut self, deps: Arc<[LoadedDylib]>) {
192        self.deps = Some(deps);
193    }
194}
195
196/// The global manager for all loaded dynamic libraries.
197pub(crate) struct Manager {
198    /// Maps a library's short name (e.g., "libc.so.6") to its full metadata.
199    /// Uses `IndexMap` to preserve loading order for symbol resolution.
200    all: IndexMap<String, GlobalDylib>,
201    /// Libraries available in the global symbol scope (RTLD_GLOBAL).
202    global: IndexMap<String, LoadedDylib>,
203    /// The number of times a new object has been added to the link map.
204    adds: u64,
205    /// The number of times an object has been removed from the link map.
206    subs: u64,
207}
208
209impl Manager {
210    pub(crate) fn add_global(&mut self, name: String, lib: LoadedDylib) {
211        debug_assert!(
212            !self.global.contains_key(&name),
213            "Library [{}] is already in global scope",
214            name
215        );
216        log::trace!("Adding [{}] to global scope", name);
217        self.global.insert(name, lib);
218    }
219
220    pub(crate) fn add(&mut self, name: String, lib: GlobalDylib) {
221        let res = self.all.insert(name.clone(), lib);
222        debug_assert!(res.is_none(), "Library [{}] is already registered", name);
223        self.adds += 1;
224        log::trace!("Registered [{}] in all map", name);
225    }
226
227    pub(crate) fn remove(&mut self, shortname: &str) {
228        let lib = self
229            .all
230            .shift_remove(shortname)
231            .expect("Library is not registered");
232        self.subs += 1;
233        let res = self.global.shift_remove(shortname);
234        debug_assert!(
235            lib.flags.is_global() == res.is_some(),
236            "Inconsistent global scope state when removing [{}]",
237            shortname
238        );
239    }
240
241    #[inline]
242    pub(crate) fn get(&self, name: &str) -> Option<&GlobalDylib> {
243        self.all.get(name)
244    }
245
246    #[inline]
247    pub(crate) fn get_mut(&mut self, name: &str) -> Option<&mut GlobalDylib> {
248        self.all.get_mut(name)
249    }
250
251    #[inline]
252    pub(crate) fn contains(&self, name: &str) -> bool {
253        self.all.contains_key(name)
254    }
255
256    #[inline]
257    pub(crate) fn all_values(&self) -> indexmap::map::Values<'_, String, GlobalDylib> {
258        self.all.values()
259    }
260
261    #[inline]
262    pub(crate) fn global_values(&self) -> indexmap::map::Values<'_, String, LoadedDylib> {
263        self.global.values()
264    }
265
266    #[inline]
267    pub(crate) fn all_iter(&self) -> indexmap::map::Iter<'_, String, GlobalDylib> {
268        self.all.iter()
269    }
270
271    pub(crate) fn adds(&self) -> u64 {
272        self.adds
273    }
274
275    pub(crate) fn subs(&self) -> u64 {
276        self.subs
277    }
278
279    #[inline]
280    pub(crate) fn keys(&self) -> indexmap::map::Keys<'_, String, GlobalDylib> {
281        self.all.keys()
282    }
283
284    #[inline]
285    pub(crate) fn get_index(&self, index: usize) -> Option<(&String, &GlobalDylib)> {
286        self.all.get_index(index)
287    }
288
289    pub(crate) fn promote(&mut self, shortname: &str, flags: OpenFlags) {
290        let promotable = flags.promotable();
291        let entry = self.get_mut(shortname).expect("Library must be registered");
292        if !entry.flags.contains(promotable) {
293            entry.flags |= promotable;
294            if flags.is_global() {
295                let core = entry.inner.clone();
296                self.add_global(shortname.to_owned(), core);
297            }
298        }
299    }
300}
301
302/// The global static instance of the library manager, protected by a readers-writer lock.
303pub(crate) static MANAGER: Lazy<RwLock<Manager>> = Lazy::new(|| {
304    RwLock::new(Manager {
305        all: IndexMap::new(),
306        global: IndexMap::new(),
307        adds: 0,
308        subs: 0,
309    })
310});
311
312/// Registers a library in the global manager.
313///
314/// If the library has `RTLD_GLOBAL` set, it's also added to the global search scope.
315pub(crate) fn register(
316    lib: LoadedDylib,
317    flags: OpenFlags,
318    manager: &mut Manager,
319    state: DylibState,
320) {
321    let name = lib.name();
322    let is_main = name.is_empty();
323    let shortname = lib.shortname().to_owned();
324
325    let mut flags = flags;
326    if name.contains("libc")
327        || name.contains("libpthread")
328        || name.contains("libdl")
329        || name.contains("libgcc_s")
330        || name.contains("ld-linux")
331        || name.contains("ld-musl")
332    {
333        flags |= OpenFlags::RTLD_NODELETE;
334    }
335
336    log::debug!(
337        "Registering library: [{}] (full path: [{}]) flags: [{:?}]",
338        shortname,
339        name,
340        flags
341    );
342
343    manager.add(
344        shortname.clone(),
345        GlobalDylib {
346            state,
347            inner: lib.clone(),
348            flags,
349            deps: None,
350        },
351    );
352    if flags.is_global() || is_main {
353        manager.add_global(shortname, lib);
354    }
355}
356
357/// Finds a symbol in the global search scope.
358///
359/// Iterates through all libraries registered with `RTLD_GLOBAL` in the order they were loaded.
360pub(crate) unsafe fn global_find<'a, T>(name: &str) -> Option<crate::Symbol<'a, T>> {
361    lock_read!(MANAGER).global_values().find_map(|lib| unsafe {
362        lib.get::<T>(name).map(|sym| {
363            log::trace!(
364                "Lazy Binding: find symbol [{}] from [{}] in global scope ",
365                name,
366                lib.name()
367            );
368            core::mem::transmute(sym)
369        })
370    })
371}
372
373/// Finds the next occurrence of a symbol after the specified address.
374pub(crate) unsafe fn next_find<'a, T>(addr: usize, name: &str) -> Option<crate::Symbol<'a, T>> {
375    let lock = lock_read!(MANAGER);
376    // Find the library containing the address
377    let (idx, _) = lock.all_iter().enumerate().find(|(_, (_, v))| {
378        let start = v.inner.base();
379        let end = start + v.inner.mapped_len();
380        (start..end).contains(&addr)
381    })?;
382
383    // Search in all subsequent libraries
384    lock.all_values().skip(idx + 1).find_map(|lib| unsafe {
385        lib.inner.get::<T>(name).map(|sym| {
386            log::trace!(
387                "dlsym: find symbol [{}] from [{}] via RTLD_NEXT",
388                name,
389                lib.inner.name()
390            );
391            core::mem::transmute(sym)
392        })
393    })
394}
395
396pub(crate) fn addr2dso(addr: usize) -> Option<ElfLibrary> {
397    log::trace!("addr2dso: addr [{:#x}]", addr);
398    // Use the manager directly to avoid potential cloning if not needed,
399    // but here we return ElfLibrary which is a wrapper.
400    crate::lock_read!(MANAGER).all_values().find_map(|v| {
401        let start = v.dylib_ref().base();
402        let end = start + v.dylib_ref().mapped_len();
403        if (start..end).contains(&addr) {
404            Some(v.get_lib())
405        } else {
406            None
407        }
408    })
409}
410
411/// Updates the dependency Searchlist for the specified root libraries.
412///
413/// For each root, it performs a Breadth-First Search (BFS) over its dependency tree
414/// to calculate a flat list of all libraries (the Searchlist) used for symbol resolution.
415/// This matches the behavior of the glibc dynamic linker.
416pub(crate) fn update_dependency_scopes<'a>(
417    manager: &mut Manager,
418    roots: impl Iterator<Item = &'a str>,
419) {
420    for root_name in roots {
421        let Some(root_lib) = manager.get(root_name) else {
422            continue;
423        };
424
425        if root_lib.deps.is_some() {
426            continue;
427        }
428
429        let mut scope = Vec::new();
430        let mut visited = BTreeSet::new();
431        let mut queue = VecDeque::new();
432
433        visited.insert(root_name.to_owned());
434        queue.push_back(root_name.to_owned());
435
436        while let Some(current_name) = queue.pop_front() {
437            let Some(lib_entry) = manager.get(&current_name) else {
438                continue;
439            };
440            let dylib = lib_entry.dylib();
441            scope.push(dylib.clone());
442
443            for needed in dylib.needed_libs() {
444                if !manager.contains(needed) {
445                    continue;
446                }
447                if !visited.contains(needed) {
448                    visited.insert(needed.to_owned());
449                    queue.push_back(needed.to_owned());
450                }
451            }
452        }
453        if let Some(entry) = manager.get_mut(root_name) {
454            entry.set_deps(Arc::from(scope));
455        }
456    }
457}
458
459pub(crate) fn register_atexit(
460    dso_handle: *mut c_void,
461    func: unsafe extern "C" fn(*mut c_void),
462    arg: *mut c_void,
463) -> c_int {
464    DESTRUCTORS.write().push(Destructor {
465        dso_handle,
466        func,
467        arg,
468    });
469    0
470}
471
472struct Destructor {
473    dso_handle: *mut c_void,
474    func: unsafe extern "C" fn(*mut c_void),
475    arg: *mut c_void,
476}
477
478unsafe impl Send for Destructor {}
479unsafe impl Sync for Destructor {}
480
481static DESTRUCTORS: Lazy<RwLock<Vec<Destructor>>> = Lazy::new(|| RwLock::new(Vec::new()));
482
483pub(crate) fn finalize(dso_handle: *mut c_void, range: Option<core::ops::Range<usize>>) {
484    let mut to_run = Vec::new();
485    {
486        let mut range = range;
487        if range.is_none() && !dso_handle.is_null() {
488            let manager = MANAGER.read();
489            for v in manager.all_values() {
490                let start = v.dylib_ref().base();
491                let end = start + v.dylib_ref().mapped_len();
492                if (start..end).contains(&(dso_handle as usize)) {
493                    range = Some(start..end);
494                    break;
495                }
496            }
497        }
498
499        let mut all_destructors = DESTRUCTORS.write();
500        let mut i = 0;
501        while i < all_destructors.len() {
502            let matches = match (dso_handle.is_null(), &range) {
503                (true, _) => true, // NULL matches all
504                (false, Some(r)) => r.contains(&(all_destructors[i].dso_handle as usize)),
505                (false, None) => all_destructors[i].dso_handle == dso_handle,
506            };
507
508            if matches {
509                to_run.push(all_destructors.remove(i));
510            } else {
511                i += 1;
512            }
513        }
514    }
515    if !to_run.is_empty() {
516        log::debug!(
517            "Running {} destructors for handle {:p}",
518            to_run.len(),
519            dso_handle
520        );
521    }
522    for destructor in to_run.into_iter().rev() {
523        unsafe { (destructor.func)(destructor.arg) };
524    }
525}
526
527#[unsafe(no_mangle)]
528pub unsafe extern "C" fn __cxa_thread_atexit_impl(
529    func: unsafe extern "C" fn(*mut c_void),
530    arg: *mut c_void,
531    dso_handle: *mut c_void,
532) -> c_int {
533    register_atexit(dso_handle, func, arg)
534}
535
536#[unsafe(no_mangle)]
537pub unsafe extern "C" fn __cxa_atexit(
538    func: unsafe extern "C" fn(*mut c_void),
539    arg: *mut c_void,
540    dso_handle: *mut c_void,
541) -> c_int {
542    register_atexit(dso_handle, func, arg)
543}
544
545#[unsafe(no_mangle)]
546pub unsafe extern "C" fn __cxa_finalize(dso_handle: *mut c_void) {
547    finalize(dso_handle, None);
548}
dlopen_rs/core_impl/register.rs

dlopen_rs/core_impl/
register.rs
