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::{
    file::ELF64,
    header::Ehdr64,
    section::{build_string_table, Contents64, Section64, Shdr64, Type as ShType},
    segment::{Phdr64, Segment64, Type as SeType, PF_R, PF_W, PF_X},
};

use crate::{
    common::{Output, Section},
    error::Error,
};

use super::segment::*;
use super::ElfObject;

pub struct ElfOutput {
    destination: PathBuf,
    file: ELF64,
    segment_view: SegmentView,
    writer: BufWriter<File>,
}

impl ElfOutput {
    pub fn new(destination: PathBuf) -> Result<Self, Error> {
        use elf_utilities::header::{Class, Data, Machine, Type, Version, OSABI};

        let mut elf = ELF64::default();
        elf.ehdr.set_elf_type(Type::Exec);
        elf.ehdr.set_class(Class::Bit64);
        elf.ehdr.set_machine(Machine::X8664);
        elf.ehdr.set_object_version(Version::Current);
        elf.ehdr.set_file_version(Version::Current);
        elf.ehdr.set_osabi(OSABI::Linux);
        elf.ehdr.set_data(Data::LSB2);

        let str_path = expand_path(&destination)?;
        let writer = file_writer(str_path, 0o755)?;

        let result = Self {
            destination,
            file: elf,
            segment_view: SegmentView::default(),
            writer,
        };

        Ok(result)
    }

    fn populate_sections(&mut self) -> Result<(), Error> {
        let mut names = Vec::new();
        let mut name_idx = 0usize;

        for (name, sections) in self.segment_view.sections_mut() {
            let mut data_size = 0;

            for t in sections.iter() {
                data_size += t.data_size;
            }

            if data_size == 0 {
                continue;
            }

            names.push(name);
            let section = Section64 {
                name: name.into(),
                header: make_section_header(name_idx, ShType::ProgBits, data_size),
                contents: Contents64::Raw(Vec::new()), // placeholder only
            };
            name_idx += name.len() + 1;
            self.file.add_section(section);
        }

        let name = ".shstrtab";
        names.push(name);
        let string_table = build_string_table(names, true);
        let section = Section64 {
            name: name.into(),
            header: make_section_header(name_idx, ShType::StrTab, string_table.len() as u64),
            contents: Contents64::Raw(string_table),
        };
        self.file.add_section(section);

        Ok(())
    }

    fn populate_segment(&mut self, offset: &mut u64, size: u64, which: SegmentType) {
        if size == 0 {
            return;
        }

        let mut segment = Segment64 {
            header: Phdr64::default(),
        };
        segment.header.set_type(SeType::Load);
        segment.header.p_filesz = size;
        segment.header.p_memsz = size;
        segment.header.p_offset = *offset;

        let mut page_size = next_page(size);

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

        self.file.segments.push(segment);
        self.file.ehdr.e_phnum += 1;
        self.file.ehdr.e_phentsize = size_of::<Phdr64>() as u16;

        *offset += page_size;
    }

    fn populate_segments(&mut self) -> Result<u64, Error> {
        let mut offset = 0u64;
        // program header/segments
        // contains .text + .rodata as one segment
        self.populate_segment(
            &mut offset,
            self.segment_view.program_size(),
            SegmentType::Text,
        );
        // contains .data as one segment
        self.populate_segment(
            &mut offset,
            self.segment_view.data_size(),
            SegmentType::Data,
        );
        // contains .bss as one segment
        self.populate_segment(&mut offset, self.segment_view.bss_size(), SegmentType::Bss);

        Ok(offset)
    }
}

impl Output<ElfObject> for ElfOutput {
    fn process_section(&mut self, section: Section<(usize, usize)>) -> Result<(), Error> {
        self.segment_view.append_section(section)
    }

    fn finalize(mut self, objects: &[ElfObject]) -> 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;

        self.populate_sections()?;
        let page_size = self.populate_segments()?;

        self.file.ehdr.e_shnum = self.file.sections.len() as u16;
        self.file.ehdr.e_entry = 4096; // TODO
        self.file.ehdr.e_phentsize = PHS;
        self.file.ehdr.e_shentsize = SHS;
        self.file.ehdr.e_phoff = EHS;
        self.file.ehdr.e_shoff = self.file.ehdr.e_phoff
            + (usize::from(self.file.ehdr.e_phentsize) * self.file.segments.len()) as u64;

        let mut offset = 0;

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

        // write program header table
        for seg in self.file.segments.iter_mut() {
            seg.header.p_offset = offset as u64;
            offset += self.writer.write(&seg.header.to_le_bytes())?;
        }

        eprintln!("SH start: {}", offset);
        // write section header table
        for sec in self.file.sections.iter_mut() {
            if sec.header.get_type() == ShType::StrTab {
                sec.header.sh_offset = offset as u64 + u64::from(SHS);
            }
            offset += self.writer.write(&sec.header.to_le_bytes())?;
            if sec.header.get_type() == ShType::StrTab {
                offset += self.writer.write(&sec.to_le_bytes())?;
            }
        }

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

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

        self.writer.flush()?;

        Ok(self.destination)
    }
}

fn 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_section_header(name_idx: usize, stype: ShType, 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 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 next_page(size: u64) -> u64 {
    let page_size: u64 = page_size::get() as u64;

    let pages_needed = (size / page_size) + 1;

    page_size * pages_needed
}

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