rld/src/formats/elf/object.rs

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

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

pub struct ElfObject {
    object_index: usize,
    origin: PathBuf,
    data: Mmap,
    ehdr: Ehdr64,
    shdrs: Vec<Shdr64>,
    symbols: Vec<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_name(&self, name_idx: usize, sh_index: usize) -> Result<&str, 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 0..count {
                    let start = offset + s_i * ent_size;
                    let s64 = Symbol64::deserialize(&self.data, start)
                        .map_err(|err| Error::ParseError(err))?;
                    self.symbols.push(s64);
                }
            }
            // self.symbols.push();
            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
        }
    }

    fn make_symbol(&self, s64: &Symbol64) -> Symbol {
        let binding = match s64.get_bind() {
            Bind::Global => Binding::Global,
            Bind::Local => Binding::Local,
            Bind::Weak => Binding::Weak,
            _ => panic!("Unexpected binding type encountered on symbol"), // this is screened!
        };

        Symbol {
            object_index: self.object_index,
            binding,
            address: s64.st_value,
            size: s64.st_size,
        }
    }
}

impl Relocatable for ElfObject {
    fn new(origin: PathBuf, object_index: usize) -> Result<Self, Error> {
        ElfObject::new(origin, 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 symbols(&self) -> SymbolIterBox {
        Box::new(self.symbols.iter().map(move |s| {
            // TODO: change SymbolIterBox to allow errors!
            let name = self
                .resolve_name(s.st_name as usize, s.st_other as usize)
                .expect("Symbol name resolution");
            (name, self.make_symbol(s))
        }))
    }
}

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)
}