rld/src/formats/elf/object.rs

use std::{
    fmt::Display,
    fs::File,
    path::{Path, PathBuf},
};

use crate::common::{Relocatable, Section, SectionInfo};
use crate::{common::SectionIterBox, error::Error};
use elf_utilities::{
    header::Ehdr64,
    section::{Shdr64, Type},
    symbol::{Bind, Symbol64, Type as SymType},
};
use memmap::Mmap;

pub struct ElfObject {
    object_index: usize,
    origin: PathBuf,
    data: Mmap,
    ehdr: Ehdr64,
    shdrs: Vec<Shdr64>,
    symbols: Vec<(usize, Symbol64)>,
}

impl ElfObject {
    pub fn new(origin: PathBuf, object_index: usize) -> Result<Self, Error> {
        let str_origin = origin
            .as_path()
            .to_str()
            .ok_or(Error::InvalidObjectType(404))?;

        let file = File::open(str_origin)?;
        let data = unsafe { Mmap::map(&file)? };
        let ehdr = parse_elf_header(&data)?;

        let mut result = ElfObject {
            object_index,
            origin,
            data,
            ehdr,
            shdrs: Vec::new(),
            symbols: Vec::new(),
        };

        // TODO: make paralellization from outside possible?
        result.populate()?;

        Ok(result)
    }

    fn resolve_symbol_value(&self, symbol_index: usize) -> Result<&[u8], Error> {
        if let Some((_, s64)) = self.symbols.get(symbol_index) {
            if s64.st_size == 0 {
                Ok(&[])
            } else if let Some(shdr) = self.shdrs.get(s64.st_shndx as usize) {
                let start: usize = (shdr.sh_offset + s64.st_value) as usize;
                let end: usize = start + s64.st_size as usize;

                Ok(&self.data[start..end])
            } else {
                Err(Error::InvalidSectionIndex)
            }
        } else {
            Err(Error::InvalidSymbolIndex)
        }
    }

    fn resolve_name(&self, name_idx: usize, sh_index: usize) -> Result<&str, Error> {
        if name_idx == 0 {
            let error = if sh_index == self.ehdr.e_shstrndx.into() {
                Error::InvalidSectionName
            } else {
                Error::InvalidSymbolName
            };

            return Err(error);
        }

        let strtab = self.parse_strtab(sh_index)?;
        let strtab_offset: usize = strtab.sh_offset as usize;
        let start = strtab_offset + name_idx;

        let mut i = start;
        while self.data[i] != 0 {
            i += 1;
            // sanity check
            if i - start > 1024 {
                return Err(Error::DataError(self.origin.clone()));
            }
        }

        // TODO: consider unchecked?
        std::str::from_utf8(&self.data[start..i]).map_err(|err| Error::ParseError(err.into()))
    }

    // TODO: consider caching these
    fn parse_strtab(&self, sh_index: usize) -> Result<Shdr64, Error> {
        let sh_start: usize = self.ehdr.e_shoff as usize;
        let sh_size: usize = self.ehdr.e_shentsize.into();
        let offset: usize = sh_start + sh_size * sh_index;

        bincode::deserialize(&self.data[offset..offset + sh_size])
            .map_err(|err| Error::ParseError(err))
    }

    fn populate(&mut self) -> Result<(), Error> {
        for i in 0..self.ehdr.e_shnum {
            let i_usize: usize = i.into();
            let sh_start = self.ehdr.e_shoff as usize;
            let sh_size = usize::from(self.ehdr.e_shentsize);

            let offset: usize = sh_start + i_usize * sh_size;
            let shdr: Shdr64 = bincode::deserialize(&self.data[offset..offset + sh_size])
                .map_err(|err| Error::ParseError(err))?;

            if shdr.get_type() == Type::SymTab {
                let ent_size = shdr.sh_entsize as usize;
                let offset = shdr.sh_offset as usize;
                let count = shdr.sh_size as usize / ent_size;

                for s_i in 1..count {
                    // skip null-symbol
                    let start = offset + s_i * ent_size;
                    let s64 = Symbol64::deserialize(&self.data, start)
                        .map_err(|err| Error::ParseError(err))?;

                    match s64.get_type() {
                        SymType::Object | SymType::Func | SymType::NoType => {
                            self.symbols.push((shdr.sh_link as usize, s64))
                        }
                        _ => {}
                    }
                }
            }

            self.shdrs.push(shdr);
        }

        Ok(())
    }

    fn make_section(&self, shdr: &Shdr64) -> Option<Result<Section, Error>> {
        let sh_index: usize = self.ehdr.e_shstrndx.into();
        let name = match self.resolve_name(shdr.sh_name as usize, sh_index) {
            Ok(n) => n,
            Err(err) => return Some(Err(err)),
        };

        let mut si = SectionInfo {
            object_index: self.object_index,
            file_size: shdr.sh_size,
            data_size: shdr.sh_size,
            offset: shdr.sh_offset,
        };

        if name.starts_with(".bss") {
            si.file_size = 0;
            Some(Ok(Section::Bss(si)))
        } else if name.starts_with(".text") {
            Some(Ok(Section::Text(si)))
        } else if name.starts_with(".rodata") {
            Some(Ok(Section::Data(si, true)))
        } else if name.starts_with(".data") {
            Some(Ok(Section::Data(si, false)))
        } else {
            None
        }
    }
}

impl Relocatable for ElfObject {
    fn new(origin: PathBuf, object_index: usize) -> Result<Self, Error> {
        ElfObject::new(origin, object_index)
    }

    fn object_index(&self) -> usize {
        self.object_index
    }

    fn origin(&self) -> &Path {
        &self.origin
    }

    fn sections(&self) -> SectionIterBox {
        let iter = self
            .shdrs
            .iter()
            .filter_map(move |shdr| match shdr.get_type() {
                Type::ProgBits | Type::NoBits => self.make_section(shdr),
                _ => None,
            });

        Box::new(iter)
    }

    fn bytes(&self, offset: u64, size: u64) -> Result<&[u8], Error> {
        let o = offset as usize;
        let s = size as usize;
        Ok(&self.data[o..o + s])
    }

    fn symbol_count(&self) -> usize {
        self.symbols.len()
    }

    fn symbol_name(&self, index: usize) -> Result<&str, Error> {
        if let Some(symbol) = self.symbols.get(index) {
            self.resolve_name(symbol.1.st_name as usize, symbol.0)
        } else {
            Err(Error::InvalidSymbolIndex)
        }
    }

    fn symbol_value(&self, index: usize) -> Result<&[u8], Error> {
        self.resolve_symbol_value(index)
    }

    fn symbol_file_offset(&self, index: usize) -> Result<usize, Error> {
        let (sh_index, s64) = self.symbols.get(index).ok_or(Error::InvalidSymbolIndex)?;

        let shdr = self
            .shdrs
            .get(*sh_index)
            .ok_or(Error::InvalidSectionIndex)?;
        let sh_offset: usize = shdr.sh_offset as usize;

        Ok(sh_offset)
    }

    fn symbol_needs_resolving(&self, index: usize) -> bool {
        match self.symbols[index].1.get_bind() {
            Bind::Global | Bind::Weak => true,
            _ => false,
        }
    }
}

impl Display for ElfObject {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "==={:?}===\n>Symbols:\n{}\n>Sections:\n",
            self.origin().file_name().unwrap(),
            "TODO"
        )?;

        for section in self.sections() {
            let u = section.unwrap();
            write!(f, "{}", u)?;
        }

        Ok(())
    }
}

fn parse_strtab_name(strtab: &[u8], idx: u32) -> String {
    let bytes: Vec<u8> = strtab
        .iter()
        .skip(idx as usize)
        .take_while(|byte| **byte != 0x00)
        .copied()
        .collect();

    std::str::from_utf8(&bytes)
        .expect("Symbol name parse")
        .to_string()
}

fn parse_elf_header(data: &Mmap) -> Result<Ehdr64, Error> {
    use elf_utilities::header::Type;

    let magic = &data[0..4];
    if magic != &[0x7f, 0x45, 0x4c, 0x46] {
        return Err(Error::InvalidInput);
    }

    let ehdr = Ehdr64::deserialize(data, 0).map_err(|err| Error::ParseError(err))?;
    if ehdr.get_type() != Type::Rel {
        return Err(Error::InvalidObjectType(u32::from(ehdr.e_type)));
    }

    Ok(ehdr)
}