yggdrasil/userspace/dyn-loader/src/relocation.rs

use elf::relocation::Rela;

use crate::{
    object::{DynamicSymbol, ObjectMapping},
    Error, State,
};

pub enum RelValue {
    QWord(i64),
    #[allow(unused)] // unused on x86_64
    QDWord(i64, i64),
}

impl RelValue {
    pub fn write(&self, mapping: &mut ObjectMapping, offset: u64) {
        let addr = mapping.base + offset as usize;
        unsafe {
            match *self {
                Self::QWord(value) => {
                    (addr as *mut i64).write_volatile(value);
                }
                Self::QDWord(word0, word1) => {
                    (addr as *mut i64).write_volatile(word0);
                    (addr as *mut i64).add(1).write_volatile(word1);
                }
            }
        }
    }
}

pub trait RelocationExt {
    fn resolve<'a, F: Fn(u32) -> Result<&'a DynamicSymbol, Error>>(
        &'a self,
        state: &mut State,
        tls_module_id: Option<usize>,
        image_symbol: F,
        image_base: usize,
        vma_start: usize,
    ) -> Result<Option<RelValue>, Error>;
}

#[cfg(any(target_arch = "x86_64", rust_analyzer))]
impl RelocationExt for Rela {
    fn resolve<'a, F: Fn(u32) -> Result<&'a DynamicSymbol, Error>>(
        &'a self,
        state: &mut State,
        tls_module_id: Option<usize>,
        image_symbol: F,
        image_base: usize,
        vma_start: usize,
    ) -> Result<Option<RelValue>, Error> {
        let image_base = image_base as i64;
        let symbol = image_symbol(self.r_sym)?;

        let base_value = match self.r_type {
            elf::abi::R_X86_64_JUMP_SLOT => {
                // TODO lazy binding could be implemented here
                state.lookup_resolved_symbol(&symbol.name)
            }
            elf::abi::R_X86_64_GLOB_DAT => state.lookup_resolved_glob_dat(&symbol.name),
            elf::abi::R_X86_64_64 => state.lookup_resolved_symbol(&symbol.name),
            // Retrieve raw symbol value, as this might be an offset into local TLS struct
            elf::abi::R_X86_64_DTPMOD64 | elf::abi::R_X86_64_DTPOFF64 => {
                symbol.raw.st_value as usize
            }
            _ => {
                if let Some(value) = symbol.value {
                    value
                } else if symbol.raw.st_symtype() == elf::abi::STT_NOTYPE {
                    0
                } else {
                    return Err(Error::CannotRelocate);
                }
            }
        } as i64;

        match self.r_type {
            // Direct 64 bit
            elf::abi::R_X86_64_64 => Ok(Some(RelValue::QWord(base_value + self.r_addend))),
            elf::abi::R_X86_64_JUMP_SLOT => Ok(Some(RelValue::QWord(base_value))),
            elf::abi::R_X86_64_COPY => todo!("{:?}: R_X86_64_COPY", symbol.name),
            // GLOB_DAT
            elf::abi::R_X86_64_GLOB_DAT => Ok(Some(RelValue::QWord(base_value))),
            // Adjust by image base
            elf::abi::R_X86_64_RELATIVE => Ok(Some(RelValue::QWord(image_base + self.r_addend))),

            // ID of module containing this symbol (if not present assume local symbol)
            elf::abi::R_X86_64_DTPMOD64 => {
                // TODO I'm not sure if per-object TLS is really needed:
                //      all TLS segments could just be collected into a single one, simplifying
                //      __tls_get_addr()
                if let Some((module_id, _)) = state.lookup_resolved_tls_symbol(&symbol.name) {
                    Ok(Some(RelValue::QWord(module_id as _)))
                } else {
                    Ok(Some(RelValue::QWord(tls_module_id.unwrap() as _)))
                }
            }
            // Offset in module's TLS block
            elf::abi::R_X86_64_DTPOFF64 => {
                if let Some((_, value)) = state.lookup_resolved_tls_symbol(&symbol.name) {
                    Ok(Some(RelValue::QWord(self.r_addend + value as i64)))
                } else {
                    Ok(Some(RelValue::QWord(self.r_addend + base_value)))
                }
            }

            ty => Err(Error::UnsupportedRelocation(ty)),
        }
    }
}

#[cfg(any(target_arch = "aarch64", rust_analyzer))]
impl RelocationExt for Rela {
    fn resolve<'a, F: Fn(u32) -> Result<&'a DynamicSymbol, Error>>(
        &'a self,
        state: &mut State,
        tls_module_id: Option<usize>,
        image_symbol: F,
        image_base: usize,
        vma_start: usize,
    ) -> Result<Option<RelValue>, Error> {
        let image_base = image_base as i64;
        let symbol = image_symbol(self.r_sym)?;

        let base_value = match self.r_type {
            elf::abi::R_AARCH64_ABS64 => state.lookup_resolved_symbol(&symbol.name),
            // TODO lazy binding?
            elf::abi::R_AARCH64_JUMP_SLOT => state.lookup_resolved_symbol(&symbol.name),
            elf::abi::R_AARCH64_GLOB_DAT => state.lookup_resolved_glob_dat(&symbol.name),
            elf::abi::R_AARCH64_TLSDESC => symbol.value.unwrap_or(symbol.raw.st_value as usize),
            _ => {
                if let Some(value) = symbol.value {
                    value
                } else if symbol.raw.st_symtype() == elf::abi::STT_NOTYPE {
                    0
                } else {
                    return Err(Error::CannotRelocate);
                }
            }
        } as i64;

        match self.r_type {
            // Adjust by image base
            elf::abi::R_AARCH64_RELATIVE => Ok(Some(RelValue::QWord(image_base + self.r_addend))),
            elf::abi::R_AARCH64_JUMP_SLOT => Ok(Some(RelValue::QWord(base_value))),
            // GLOB_DAT
            elf::abi::R_AARCH64_GLOB_DAT => Ok(Some(RelValue::QWord(base_value))),
            // Direct 64 bit
            elf::abi::R_AARCH64_ABS64 => Ok(Some(RelValue::QWord(base_value + self.r_addend))),
            elf::abi::R_AARCH64_TLSDESC => {
                let word0 = crate::builtins::__dl_tlsdesc_static as usize;

                let reloc_offset =
                    unsafe { *((image_base as usize + self.r_offset as usize) as *const i64) };

                let word1 = if let Some((module_id, offset)) =
                    state.lookup_resolved_tls_symbol(&symbol.name)
                {
                    // FIXME not really sure about this code
                    let tls_base = state.tls_table[module_id].data.base;

                    base_value + reloc_offset + (tls_base + offset) as i64
                } else {
                    let local_tls = state.tls_table[tls_module_id.unwrap() - 1].data.base;

                    local_tls as i64 + reloc_offset + self.r_addend
                };

                Ok(Some(RelValue::QDWord(word0 as _, word1)))
            }
            _ => Err(Error::UnsupportedRelocation(self.r_type)),
        }
    }
}