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::Shdr64,
    symbol::{Bind, Symbol64},
};
use memmap::Mmap;

pub struct ElfObject {
    object_index: usize,
    origin: PathBuf,
    data: Mmap,
    ehdr: Ehdr64,
    sh_name_offset: usize,
}

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 shstrtab = parse_shstrtab(&data, &ehdr)?;

        let result = ElfObject {
            object_index,
            origin,
            data,
            ehdr,
            sh_name_offset: shstrtab.sh_offset as usize,
        };

        Ok(result)
    }

    fn section_name(&self, shdr: &Shdr64) -> Result<&str, Error> {
        let idx: usize = shdr.sh_name as usize;
        let start = self.sh_name_offset + 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()))
    }

    fn make_section(&self, offset: usize, sh: &Shdr64) -> Option<Result<Section, Error>> {
        let name = match self.section_name(sh) {
            Ok(n) => n,
            Err(err) => return Some(Err(err)),
        };

        let mut si = SectionInfo {
            object_index: self.object_index,
            file_size: sh.sh_size,
            data_size: sh.sh_size,
            offset: offset as u64,
        };

        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, strtab: &[u8]) -> 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,
            name: parse_strtab_name(&strtab, s64.st_name),
            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 = (0..self.ehdr.e_shnum).into_iter().filter_map(move |i| {
            let i_usize = usize::from(i);
            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 shr: Result<Shdr64, Error> =
                bincode::deserialize(&self.data[offset..offset + sh_size])
                    .map_err(|err| Error::ParseError(err));

            match shr {
                Err(err) => Some(Err(err)),
                Ok(sh) => self.make_section(offset, &sh),
            }
        });

        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 {
        // if let Some(strtab_section) = self
        //     .elf
        //     .first_section_by(|s64| s64.header.get_type() == Type::StrTab && s64.name == ".strtab")
        // {
        //     let strtab = match &strtab_section.contents {
        //         Contents64::Raw(bytes) => bytes,
        //         _ => panic!("Unexpected strtab content type"),
        //     };

        //     let iter = self
        //         .elf
        //         .sections
        //         .iter()
        //         .filter_map(move |s| match &s.contents {
        //             Contents64::Symbols(symbols) => {
        //                 Some(symbols.iter().filter_map(move |sym| match sym.get_bind() {
        //                     Bind::Global | Bind::Local | Bind::Weak => {
        //                         Some(self.make_symbol(sym, strtab))
        //                     }
        //                     _ => None,
        //                 }))
        //             }
        //             _ => None,
        //         })
        //         .flatten();

        //     Box::new(iter)
        // } else {
        Box::new(std::iter::empty())
        //}
    }
}

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

fn parse_shstrtab(data: &Mmap, ehdr: &Ehdr64) -> Result<Shdr64, Error> {
    let idx: usize = ehdr.e_shstrndx.into();
    let sh_start: usize = ehdr.e_shoff as usize;
    let sh_size: usize = ehdr.e_shentsize.into();
    let offset: usize = sh_start + sh_size * idx;

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