elf_loader 0.15.1

use super::{DynLifecycleHandler, LoadHook, LoadHookContext, LoaderInner};
use crate::{
    MmapError, ParsePhdrError, Result,
    elf::{ElfDyn, ElfHeader, ElfLayout, ElfPhdr, ElfPhdrs, ElfProgramType, NativeElfLayout},
    image::RawDynamic,
    input::{ElfReader, PathBuf},
    os::{Mmap, PageSize},
    segment::{ELFRelro, ElfSegments, SegmentBuilder, program::ProgramSegments},
    tls::{TlsInfo, TlsResolver},
};
use alloc::{boxed::Box, vec::Vec};
use core::{ffi::c_char, marker::PhantomData, ptr::NonNull};

/// Builder for creating relocated ELF objects
///
/// This structure is used internally during the loading process to collect
/// and organize the various components of a relocated ELF file before
/// building the final loaded image.
pub(crate) struct ImageBuilder<'hook, H, M, Tls, D = (), L: ElfLayout = NativeElfLayout>
where
    H: LoadHook<L>,
    M: Mmap,
    Tls: TlsResolver,
{
    /// Hook function for processing program headers (always present)
    hook: &'hook mut H,

    /// Mapped program headers
    phdr_mmap: Option<&'static [ElfPhdr<L>]>,

    /// Loader source path or caller-provided source identifier.
    pub(crate) path: PathBuf,

    /// ELF header
    pub(crate) ehdr: ElfHeader<L>,

    /// GNU_RELRO segment information
    pub(crate) relro: Option<ELFRelro>,

    /// Pointer to the dynamic section
    pub(crate) dynamic_ptr: Option<NonNull<ElfDyn<L>>>,

    /// TLS information
    pub(crate) tls_info: Option<TlsInfo>,

    /// Whether to use static TLS
    pub(crate) static_tls: bool,

    /// Page size used for segment layout.
    page_size: usize,

    /// User-defined data
    pub(crate) user_data: D,

    /// Memory segments
    pub(crate) segments: ElfSegments,

    /// Initialization function handler
    pub(crate) init_fn: DynLifecycleHandler,

    /// Finalization function handler
    pub(crate) fini_fn: DynLifecycleHandler,

    /// Pointer to the interpreter path (PT_INTERP)
    pub(crate) interp: Option<NonNull<c_char>>,

    /// Pointer to the .eh_frame_hdr section (PT_GNU_EH_FRAME)
    pub(crate) eh_frame_hdr: Option<NonNull<u8>>,

    /// Phantom data to maintain Mmap type information
    _marker: PhantomData<(M, Tls, L)>,
}

pub(crate) struct ScanBuilder<L: ElfLayout = NativeElfLayout> {
    pub(crate) path: PathBuf,
    pub(crate) ehdr: ElfHeader<L>,
    pub(crate) phdrs: Box<[ElfPhdr<L>]>,
    pub(crate) reader: Box<dyn ElfReader + 'static>,
}

impl<L: ElfLayout> ScanBuilder<L> {
    #[inline]
    pub(crate) fn new(
        path: PathBuf,
        ehdr: ElfHeader<L>,
        phdrs: Box<[ElfPhdr<L>]>,
        reader: Box<dyn ElfReader + 'static>,
    ) -> Self {
        Self {
            path,
            ehdr,
            phdrs,
            reader,
        }
    }
}

impl<'hook, H, M, Tls, D, L> ImageBuilder<'hook, H, M, Tls, D, L>
where
    H: LoadHook<L>,
    Tls: TlsResolver,
    M: Mmap,
    L: ElfLayout,
{
    /// Create a new [`ImageBuilder`].
    ///
    /// # Arguments
    /// * `hook` - Hook function for processing program headers
    /// * `segments` - Memory segments of the ELF file
    /// * `path` - Loader source path or caller-provided source identifier
    /// * `ehdr` - ELF header
    /// * `init_fn` - Initialization function handler
    /// * `fini_fn` - Finalization function handler
    ///
    pub(crate) fn new(
        hook: &'hook mut H,
        segments: ElfSegments,
        path: PathBuf,
        ehdr: ElfHeader<L>,
        init_fn: DynLifecycleHandler,
        fini_fn: DynLifecycleHandler,
        static_tls: bool,
        page_size: usize,
        user_data: D,
    ) -> Self {
        Self {
            hook,
            phdr_mmap: None,
            path,
            ehdr,
            relro: None,
            dynamic_ptr: None,
            tls_info: None,
            static_tls,
            page_size,
            segments,
            user_data,
            init_fn,
            fini_fn,
            interp: None,
            eh_frame_hdr: None,
            _marker: PhantomData,
        }
    }

    /// Parse a program header and extract relevant information.
    pub(crate) fn parse_phdr(&mut self, phdr: &ElfPhdr<L>) -> Result<()> {
        let ctx = LoadHookContext::new(self.path.as_path(), phdr, &self.segments);
        self.hook.call(&ctx)?;

        match phdr.program_type() {
            ElfProgramType::DYNAMIC => {
                self.dynamic_ptr = Some(
                    NonNull::new(self.segments.get_mut_ptr(phdr.p_vaddr())).ok_or(
                        ParsePhdrError::malformed("PT_DYNAMIC is outside mapped segments"),
                    )?,
                )
            }
            ElfProgramType::GNU_RELRO => {
                self.relro = Some(ELFRelro::new::<M, L>(
                    phdr,
                    self.segments.base_addr(),
                    self.page_size,
                ))
            }
            ElfProgramType::PHDR => {
                self.phdr_mmap = Some(
                    self.segments
                        .get_slice::<ElfPhdr<L>>(phdr.p_vaddr(), phdr.p_memsz()),
                );
            }
            ElfProgramType::INTERP => {
                self.interp = Some(
                    NonNull::new(self.segments.get_mut_ptr(phdr.p_vaddr())).ok_or(
                        ParsePhdrError::malformed("PT_INTERP is outside mapped segments"),
                    )?,
                );
            }
            ElfProgramType::GNU_EH_FRAME => {
                self.eh_frame_hdr = Some(
                    NonNull::new(self.segments.get_mut_ptr(phdr.p_vaddr())).ok_or(
                        ParsePhdrError::malformed("PT_GNU_EH_FRAME is outside mapped segments"),
                    )?,
                );
            }
            ElfProgramType::TLS => {
                let tls_image = self
                    .segments
                    .get_slice::<u8>(phdr.p_vaddr(), phdr.p_filesz());
                self.tls_info = Some(TlsInfo::new(phdr, tls_image));
            }
            _ => {}
        };
        Ok(())
    }

    /// Parse all program headers and collect the builder state they describe.
    pub(crate) fn parse_phdrs(&mut self, phdrs: &[ElfPhdr<L>]) -> Result<()> {
        for phdr in phdrs {
            self.parse_phdr(phdr)?;
        }
        Ok(())
    }

    /// Create program headers from the parsed data
    ///
    /// This method creates the appropriate program header representation
    /// based on whether they are mapped in memory or need to be stored
    /// in a vector.
    ///
    /// # Arguments
    /// * `phdrs` - Slice of program headers
    ///
    /// # Returns
    /// An ElfPhdrs enum containing either mapped or vector-based headers
    pub(crate) fn create_phdrs(&self, phdrs: &[ElfPhdr<L>]) -> ElfPhdrs<L> {
        let (phdr_start, phdr_end) = self.ehdr.phdr_range();
        let phdr_size = phdr_end - phdr_start;

        // Get mapped program headers or create them from loaded segments
        self.phdr_mmap
            .or_else(|| {
                phdrs
                    .iter()
                    .filter(|phdr| phdr.program_type() == ElfProgramType::LOAD)
                    .find_map(|phdr| {
                        let seg_start = phdr.p_offset();
                        let seg_end = seg_start + phdr.p_filesz();
                        if seg_start <= phdr_start && phdr_end <= seg_end {
                            return Some(self.segments.get_slice::<ElfPhdr<L>>(
                                phdr.p_vaddr() + (phdr_start - seg_start),
                                phdr_size,
                            ));
                        }
                        None
                    })
            })
            .map(ElfPhdrs::Mmap)
            .unwrap_or_else(|| ElfPhdrs::Vec(Vec::from(phdrs)))
    }
}

impl<H, D, Arch> LoaderInner<H, D, Arch>
where
    H: LoadHook<Arch::Layout>,
    D: 'static,
    Arch: crate::relocation::RelocationArch,
{
    pub(crate) fn lifecycle_handlers(&self) -> (DynLifecycleHandler, DynLifecycleHandler) {
        (self.init_fn.clone(), self.fini_fn.clone())
    }

    #[inline]
    pub(crate) fn force_static_tls(&self) -> bool {
        self.force_static_tls
    }

    #[inline]
    pub(crate) fn page_size<M: Mmap>(&self) -> Result<PageSize> {
        let required = M::page_size();
        let page_size = self.page_size.unwrap_or(required);
        if page_size.bytes() < required.bytes()
            || !page_size.bytes().is_multiple_of(required.bytes())
        {
            return Err(MmapError::InvalidPageSize {
                configured: page_size.bytes(),
                required: required.bytes(),
            }
            .into());
        }

        Ok(page_size)
    }

    #[inline]
    pub(crate) fn initialize_dynamic(&mut self, dynamic: &mut RawDynamic<D, Arch>) -> Result<()> {
        (self.dynamic_initializer)(dynamic)
    }
}

impl<H, D, Arch> LoaderInner<H, D, Arch>
where
    H: LoadHook<Arch::Layout>,
    D: 'static,
    Arch: crate::relocation::RelocationArch,
{
    pub(crate) fn create_builder<M, Tls>(
        &mut self,
        ehdr: ElfHeader<Arch::Layout>,
        phdrs: &[ElfPhdr<Arch::Layout>],
        mut object: impl ElfReader,
        user_data: D,
    ) -> Result<ImageBuilder<'_, H, M, Tls, D, Arch::Layout>>
    where
        M: Mmap,
        Tls: TlsResolver,
    {
        let path = PathBuf::from(object.path());
        let (init_fn, fini_fn) = self.lifecycle_handlers();
        let page_size = self.page_size::<M>()?.bytes();
        let mut phdr_segments =
            ProgramSegments::new(phdrs, ehdr.is_dylib(), object.as_fd().is_some(), page_size);
        let segments = phdr_segments.load_segments::<M>(&mut object)?;
        phdr_segments.mprotect::<M>()?;

        Ok(ImageBuilder::new(
            &mut self.hook,
            segments,
            path,
            ehdr,
            init_fn,
            fini_fn,
            self.force_static_tls,
            page_size,
            user_data,
        ))
    }

    pub(crate) fn create_scan_builder(
        &self,
        ehdr: ElfHeader<Arch::Layout>,
        phdrs: &[ElfPhdr<Arch::Layout>],
        object: impl ElfReader + 'static,
    ) -> ScanBuilder<Arch::Layout> {
        let path = PathBuf::from(object.path());

        ScanBuilder::new(path, ehdr, phdrs.into(), Box::new(object))
    }
}