rld/src/formats/elf/output.rs

use std::{fs::File, io::BufWriter, io::Write, mem::size_of, path::PathBuf};

use elf_utilities::{
    file::ELF64,
    header::Ehdr64,
    section::{build_string_table, Contents64, Section64, Shdr64},
    segment::{Phdr64, PF_R, PF_W, PF_X},
};

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

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

pub struct ElfOutput<'data> {
    destination: PathBuf,
    file: ELF64,
    segment_data: SegmentData<'data>,
}

impl<'data> ElfOutput<'data> {
    pub fn new(destination: PathBuf) -> Self {
        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::Intel386);
        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);

        Self {
            destination,
            file: elf,
            segment_data: SegmentData::new(),
        }
    }

    pub fn from_object(object: &ElfObject, destination: PathBuf) -> Self {
        let other = object.elf();

        Self {
            destination,
            file: elf_bin_from_object(other),
            segment_data: SegmentData::new(),
        }
    }

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

        for (name, sections) in self.segment_data.iter_mut() {
            let mut data = Vec::new();

            for t in sections.iter_mut() {
                if let Some(iter) = &mut t.data {
                    data.extend(iter.as_mut());
                }
            }

            if data.is_empty() {
                continue;
            }

            result += data.len() as u64;
            names_len += name.len();
            names.push(name);
            let section = Section64 {
                name: name.into(),
                header: make_section_header(&name, names_len, data.len()),
                contents: Contents64::Raw(data),
            };
            self.file.add_section(section);
        }

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

        Ok(result)
    }

    fn populate_segment(&mut self, offset: &mut u64, size: u64, which: SegmentType) -> u64 {
        use elf_utilities::segment::{Segment64, Type};

        if size == 0 {
            return 0;
        }

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

        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_memsz = 0;
            }
        }

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

        *offset += size;

        size
    }

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

        Ok(result)
    }
}

impl<'data> Output<'data> for ElfOutput<'data> {
    fn allocate(&mut self, size: u64) -> Result<u64, Error> {
        Ok(size)
    }

    fn append_section(&mut self, section: Section<'data>) -> Result<(), Error> {
        self.segment_data.append_section(section)
    }

    fn finalize(mut self) -> Result<PathBuf, Error> {
        use std::io::Error as IOError;
        use std::io::ErrorKind;

        const EHS: u64 = size_of::<Ehdr64>() as u64;
        const PHS: u64 = size_of::<Phdr64>() as u64;
        const SHS: u16 = size_of::<Shdr64>() as u16;

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

        self.file.ehdr.e_phentsize = PHS as u16;
        self.file.ehdr.e_shentsize = SHS;
        self.file.ehdr.e_phoff = EHS; // TODO: align
        self.file.ehdr.e_shoff = self.file.ehdr.e_phoff
            + u64::from(self.file.ehdr.e_phnum * self.file.ehdr.e_phentsize)
            + segment_data_size;
        self.file.ehdr.e_shnum = self.file.sections.len() as u16;
        self.file.ehdr.e_entry = 0x100; // TODO

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

        let mut writer = file_writer(str_path, 0o755)?;
        let mut offset = 0;

        offset += 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 += writer.write(&seg.header.to_le_bytes())?;
        }

        // write section/segment data
        for sec in self.file.sections.iter_mut() {
            sec.header.sh_offset = offset as u64;
            offset += writer.write(&sec.to_le_bytes())?;
        }

        // write section header table
        for sec in self.file.sections.iter() {
            eprintln!("NAme index: {}", sec.header.sh_name);
            offset += writer.write(&sec.header.to_le_bytes())?;
        }
        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)
        .read(true)
        .write(true)
        .mode(permission)
        .open(output_filename)?;

    Ok(BufWriter::new(file))
}

// init new ELF64 from an object file meant as executable output
fn elf_bin_from_object(other: &ELF64) -> ELF64 {
    use elf_utilities::header::Type;

    let mut elf = ELF64::default();
    elf.ehdr.set_elf_type(Type::Exec);
    elf.ehdr.e_ehsize = other.ehdr.e_ehsize;
    elf.ehdr.e_version = other.ehdr.e_version;
    elf.ehdr.e_ident = other.ehdr.e_ident;
    elf.ehdr.e_machine = other.ehdr.e_machine;

    elf
}

fn make_section_header(name: &str, names_len: usize, size: usize) -> Shdr64 {
    use elf_utilities::section::Type;

    let mut h = Shdr64::default();
    match name {
        ".shstrtab" => h.set_type(Type::StrTab),
        ".bss" => h.set_type(Type::NoBits),
        _ => h.set_type(Type::ProgBits),
    }

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

    h
}