dlopen_rs/
register.rs

1use crate::{Dylib, OpenFlags};
2use alloc::{borrow::ToOwned, boxed::Box, string::String, sync::Arc};
3use elf_loader::RelocatedDylib;
4use indexmap::IndexMap;
5use spin::{Lazy, RwLock};
6
7impl Drop for Dylib {
8    fn drop(&mut self) {
9        if self.flags.contains(OpenFlags::RTLD_NODELETE)
10            | self.flags.contains(OpenFlags::CUSTOM_NOT_REGISTER)
11        {
12            return;
13        }
14        let ref_count = self.inner.strong_count();
15        // Dylib本身 + 全局
16        let threshold =
17            2 + self.deps.is_some() as usize + self.flags.contains(OpenFlags::RTLD_GLOBAL) as usize;
18        if ref_count == threshold {
19            log::info!("Destroying dylib [{}]", self.inner.shortname());
20            let mut lock = MANAGER.write();
21            lock.all.shift_remove(self.inner.shortname());
22            if self.flags.contains(OpenFlags::RTLD_GLOBAL) {
23                lock.global.shift_remove(self.inner.shortname());
24            }
25            for dep in self.deps.as_ref().unwrap().iter().skip(1) {
26                let dep_threshold = if let Some(lib) = lock.all.get(dep.shortname()) {
27                    if lib.flags.contains(OpenFlags::RTLD_NODELETE) {
28                        continue;
29                    }
30                    2 + lib.deps.is_some() as usize
31                        + lib.flags.contains(OpenFlags::RTLD_GLOBAL) as usize
32                } else {
33                    continue;
34                };
35                if dep.strong_count() == dep_threshold {
36                    log::info!("Destroying dylib [{}]", dep.shortname());
37                    lock.all.shift_remove(dep.shortname());
38                    lock.global.shift_remove(self.inner.shortname());
39                }
40            }
41        }
42    }
43}
44
45#[derive(Clone, Copy)]
46pub(crate) struct DylibState(u8);
47
48impl Default for DylibState {
49    fn default() -> Self {
50        Self(0)
51    }
52}
53
54impl DylibState {
55    const USED_MASK: u8 = 0b10000000;
56    const RELOCATED: u8 = 0b01111111;
57
58    #[inline]
59    pub(crate) fn set_unused(&mut self) -> &mut Self {
60        self.0 &= !Self::USED_MASK;
61        self
62    }
63
64    #[inline]
65    pub(crate) fn set_used(&mut self) -> &mut Self {
66        self.0 |= Self::USED_MASK;
67        self
68    }
69
70    #[inline]
71    pub(crate) fn is_used(&self) -> bool {
72        self.0 & Self::USED_MASK != 0
73    }
74
75    #[inline]
76    pub(crate) fn get_new_idx(&self) -> Option<u8> {
77        let remove_used_bit = self.0 & !Self::USED_MASK;
78        if remove_used_bit == Self::RELOCATED {
79            None
80        } else {
81            Some(remove_used_bit)
82        }
83    }
84
85    #[inline]
86    pub(crate) fn set_relocated(&mut self) -> &mut Self {
87        self.0 |= Self::RELOCATED;
88        self
89    }
90
91    #[allow(unused)]
92    #[inline]
93    pub(crate) fn set_new_idx(&mut self, idx: u8) -> &mut Self {
94        assert!(idx < Self::RELOCATED);
95        self.0 |= idx;
96        self
97    }
98}
99
100#[derive(Clone)]
101pub(crate) struct GlobalDylib {
102    inner: RelocatedDylib<'static>,
103    flags: OpenFlags,
104    deps: Option<Arc<Box<[RelocatedDylib<'static>]>>>,
105    pub(crate) state: DylibState,
106}
107
108unsafe impl Send for GlobalDylib {}
109unsafe impl Sync for GlobalDylib {}
110
111impl GlobalDylib {
112    #[inline]
113    pub(crate) fn get_dylib(&self) -> Dylib {
114        debug_assert!(self.deps.is_some());
115        Dylib {
116            inner: self.inner.clone(),
117            flags: self.flags,
118            deps: self.deps.clone(),
119        }
120    }
121
122    #[inline]
123    pub(crate) fn set_flags(&mut self, flags: OpenFlags) {
124        self.flags = flags;
125    }
126
127    #[inline]
128    pub(crate) fn flags(&self) -> OpenFlags {
129        self.flags
130    }
131
132    #[inline]
133    pub(crate) fn relocated_dylib(&self) -> RelocatedDylib<'static> {
134        self.inner.clone()
135    }
136
137    #[inline]
138    pub(crate) fn relocated_dylib_ref(&self) -> &RelocatedDylib<'static> {
139        &self.inner
140    }
141
142    #[inline]
143    pub(crate) fn shortname(&self) -> &str {
144        self.inner.shortname()
145    }
146
147    #[inline]
148    pub(crate) fn deps(&self) -> Option<&Arc<Box<[RelocatedDylib<'static>]>>> {
149        self.deps.as_ref()
150    }
151}
152
153pub(crate) struct Manager {
154    pub(crate) all: IndexMap<String, GlobalDylib>,
155    pub(crate) global: IndexMap<String, RelocatedDylib<'static>>,
156}
157
158pub(crate) static MANAGER: Lazy<RwLock<Manager>> = Lazy::new(|| {
159    RwLock::new(Manager {
160        all: IndexMap::new(),
161        global: IndexMap::new(),
162    })
163});
164
165pub(crate) fn register(
166    lib: RelocatedDylib<'static>,
167    flags: OpenFlags,
168    deps: Option<Arc<Box<[RelocatedDylib<'static>]>>>,
169    manager: &mut Manager,
170    state: DylibState,
171) {
172    let shortname = lib.shortname().to_owned();
173    log::debug!(
174        "Trying to register a library. Name: [{}] flags:[{:?}]",
175        shortname,
176        flags
177    );
178    manager.all.insert(
179        shortname.to_owned(),
180        GlobalDylib {
181            state,
182            inner: lib.clone(),
183            flags,
184            deps,
185        },
186    );
187    if flags.contains(OpenFlags::RTLD_GLOBAL) {
188        manager.global.insert(shortname.to_owned(), lib);
189    }
190}
191
192#[cfg(feature = "std")]
193pub(crate) fn global_find(name: &str) -> Option<*const ()> {
194    log::debug!("Lazy Binding: [{}]", name);
195    crate::loader::builtin::BUILTIN
196        .get(name)
197        .copied()
198        .or_else(|| {
199            MANAGER.read().global.values().find_map(|lib| unsafe {
200                lib.get::<()>(name).map(|sym| {
201                    log::trace!(
202                        "Lazy Binding: find symbol [{}] from [{}] in global scope ",
203                        name,
204                        lib.name()
205                    );
206                    let val = sym.into_raw();
207                    assert!(lib.base() != val as usize);
208                    //println!("{:x} {:x}", lib.base(), val as usize);
209                    val
210                })
211            })
212        })
213}
dlopen_rs/register.rs

dlopen_rs/
register.rs
