rld/src/formats/elf/output.rs

use std::{
    convert::TryInto,
    fs::File,
    io::BufWriter,
    io::{Seek, SeekFrom, Write},
    mem::size_of,
    path::{Path, PathBuf},
};

use elf_utilities::{
    header::Ehdr64,
    section::{build_string_table, Shdr64, Type as ShType},
    segment::{Phdr64, Type as SeType, PF_R, PF_W, PF_X},
};

use crate::{
    common::{Loadable, Output, SegmentType},
    error::Error,
};

use super::ElfObject;

pub struct ElfOutput {
    destination: PathBuf,
    writer: BufWriter<File>,
}

impl ElfOutput {
    pub fn new(destination: PathBuf) -> Result<Self, Error> {
        let str_path = expand_path(&destination)?;
        let writer = make_file_writer(str_path, 0o755)?;

        let result = Self {
            destination,
            writer,
        };

        Ok(result)
    }
}

impl Output<ElfObject> for ElfOutput {
    fn finalize(mut self, objects: &[ElfObject], loadable: &Loadable) -> Result<PathBuf, Error> {
        const EHS: u64 = size_of::<Ehdr64>() as u64;
        const PHS: u16 = size_of::<Phdr64>() as u16;
        const SHS: u16 = size_of::<Shdr64>() as u16;

        let strtab = make_strtab();
        let mut ehdr = make_elf_header();
        ehdr.e_shnum = Loadable::section_count() as u16 + 1; // +1 for .shstrtab
        ehdr.e_phnum = Loadable::segment_count() as u16;
        ehdr.e_shstrndx = ehdr.e_shnum - 1; // .shstrab is always last
        ehdr.e_entry = get_start_offset(loadable)?;
        ehdr.e_phentsize = PHS;
        ehdr.e_shentsize = SHS;
        ehdr.e_phoff = EHS;
        ehdr.e_shoff =
            ehdr.e_phoff + (usize::from(ehdr.e_phentsize) * usize::from(ehdr.e_phnum)) as u64;

        let mut offset = 0;

        // write ELF header
        offset += self.writer.write(&ehdr.to_le_bytes())?;

        // write program header table
        // contains .text + .rodata as one segment header
        let ph = make_program_header(offset, loadable.program_size(), SegmentType::Text);
        offset += self.writer.write(&ph.to_le_bytes())?;
        // contains .data as one segment header
        let ph = make_program_header(offset, loadable.data_size(), SegmentType::Data);
        offset += self.writer.write(&ph.to_le_bytes())?;
        // contains .bss as one segment header
        let ph = make_program_header(offset, loadable.bss_size(), SegmentType::Bss);
        offset += self.writer.write(&ph.to_le_bytes())?;

        eprintln!("SH start: {}", offset);
        let mut name_idx = 0;
        // write section header table (text + rodata + data + bss)
        for sections in loadable.segment_sections() {
            let sh = make_section_header(
                strtab.1[name_idx],
                ShType::ProgBits,
                0,
                sections.data_size(),
            );
            name_idx += 1;
            offset += self.writer.write(&sh.to_le_bytes())?;
        }

        // .shstrtab as last section header written (+ the name data including)
        {
            // write .shstrtab header
            let data_offset = (offset + usize::from(SHS)) as u64;
            let strtab_header = make_section_header(
                strtab.1[name_idx],
                ShType::StrTab,
                data_offset,
                strtab.0.len() as u64,
            );
            offset += self.writer.write(&strtab_header.to_le_bytes())?;
            // write .shstrtab data
            offset += self.writer.write(&strtab.0)?;
        }

        // program data (segments)
        offset += pad_to_next_page(&mut self.writer, offset)?;
        eprintln!("Prog start: {}", offset);
        // write section/segment data
        for bytes in loadable.program_bytes(objects) {
            offset += self.writer.write(bytes?)?;
        }
        offset += pad_to_next_page(&mut self.writer, offset)?;

        eprintln!("Data start: {}", offset);
        for bytes in loadable.data_bytes(objects) {
            offset += self.writer.write(bytes?)?;
        }

        self.writer.flush()?;

        Ok(self.destination)
    }
}

fn make_file_writer(output_filename: &str, permission: u32) -> Result<BufWriter<File>, Error> {
    use std::os::unix::fs::OpenOptionsExt;

    let file = std::fs::OpenOptions::new()
        .create(true)
        .truncate(true)
        .write(true)
        .mode(permission)
        .open(output_filename)?;

    Ok(BufWriter::new(file))
}

fn make_elf_header() -> Ehdr64 {
    use elf_utilities::header::{Class, Data, Machine, Type, Version, OSABI};

    let mut ehdr = Ehdr64::default();
    ehdr.set_elf_type(Type::Exec);
    ehdr.set_class(Class::Bit64);
    ehdr.set_machine(Machine::X8664);
    ehdr.set_object_version(Version::Current);
    ehdr.set_file_version(Version::Current);
    ehdr.set_osabi(OSABI::Linux);
    ehdr.set_data(Data::LSB2);
    ehdr.e_ehsize = size_of::<Ehdr64>() as u16;

    ehdr
}

fn make_section_header(name_idx: usize, stype: ShType, offset: u64, size: u64) -> Shdr64 {
    let mut h = Shdr64::default();
    h.set_type(stype);

    // section index used for name and section indexing (currently same)
    let idx = 1 + (name_idx as u32); // add 0 at start
                                     // name index
    h.sh_name = idx;
    // link index
    h.sh_link = 0;
    h.sh_size = size;
    // filled at finalize()
    h.sh_offset = offset;

    h
}

fn make_program_header(offset: usize, size: u64, which: SegmentType) -> Phdr64 {
    let mut header = Phdr64::default();

    header.set_type(SeType::Load);
    header.p_filesz = size;
    header.p_memsz = size;
    header.p_offset = offset as u64;

    match which {
        SegmentType::Text => header.p_flags = PF_R | PF_X,
        SegmentType::Data => header.p_flags = PF_R | PF_W,
        SegmentType::Bss => {
            header.p_flags = PF_R;
            header.p_filesz = 0;
        }
    }

    header
}

// strtab as bytes + indexes to individual strings
fn make_strtab() -> (Vec<u8>, Vec<usize>) {
    let mut section_names = Vec::from(Loadable::section_names());
    section_names.push(".shstrtab");

    let strtab_bytes = build_string_table(section_names, true);

    let mut indexes = Vec::new();
    let mut on_string = false;
    for (i, byte) in strtab_bytes.iter().enumerate() {
        if *byte == 0 {
            on_string = false;
            continue;
        }

        if on_string {
            continue;
        }

        assert!(i > 0, "First byte of .shstrab not 0");
        indexes.push(i - 1);
        on_string = true;
    }

    (strtab_bytes, indexes)
}

fn get_start_offset(loadable: &Loadable) -> Result<u64, Error> {
    use crate::error::Trace;

    loadable.start_offset.ok_or_else(|| {
        let trace = Trace {
            message: "Program entrypoint not found",
            ..Default::default()
        };

        Error::LinkingError(trace)
    })
}

fn expand_path(path: &Path) -> Result<&str, Error> {
    use std::io::Error as IOError;
    use std::io::ErrorKind;

    path.to_str().ok_or_else(|| {
        let ioe = IOError::new(ErrorKind::Other, "Path expansion fail");
        let boxed = Box::new(ioe);
        Error::IOError(boxed)
    })
}

fn pad_to_next_page(writer: &mut BufWriter<File>, offset: usize) -> Result<usize, Error> {
    let page_size = page_size::get();

    let padding = page_size - (offset % page_size);

    eprintln!("Padding from: {} with: {}", offset, padding);
    writer.seek(SeekFrom::Current(padding.try_into()?))?;

    Ok(padding)
}