//! A tiny ELF64 loader: just enough to load a statically linked user program.
//!
//! An ELF file starts with a header saying what it is (64-bit, little endian,
//! x86_64, executable) and where the *program headers* are. Each program
//! header of type PT_LOAD says: "map `memsz` bytes at virtual address
//! `vaddr`, fill the first `filesz` of them from file offset `offset`, zero
//! the rest (that's .bss), with permissions R/W/X". That's all a loader for a
//! static executable has to do; then jump to `entry`.
//!
//! Every number we read comes from an untrusted file, so all offsets and sizes
//! are bounds-checked. Reference: `man 5 elf`, https://wiki.osdev.org/ELF

use crate::memory::{self, USER_END, USER_START};
use x86_64::structures::paging::{Page, PageTableFlags, PhysFrame, Size4KiB};
use x86_64::VirtAddr;

#[derive(Debug, PartialEq, Eq)]
pub enum ElfError {
    TooShort,
    NotElf,
    Not64BitLittleEndian,
    NotX86_64,
    NotExecutable,
    BadProgramHeader,
    OutsideUserSpace,
    /// Two segments with different permissions share a page (we can only
    /// give a page one set of permissions). A linker script problem.
    SharedPage,
    OutOfMemory,
}

const PT_LOAD: u32 = 1;
const PF_X: u32 = 1;
const PF_W: u32 = 2;

fn read_u16(data: &[u8], at: usize) -> Result<u16, ElfError> {
    let bytes = data.get(at..at + 2).ok_or(ElfError::TooShort)?;
    Ok(u16::from_le_bytes(bytes.try_into().unwrap()))
}
fn read_u32(data: &[u8], at: usize) -> Result<u32, ElfError> {
    let bytes = data.get(at..at + 4).ok_or(ElfError::TooShort)?;
    Ok(u32::from_le_bytes(bytes.try_into().unwrap()))
}
fn read_u64(data: &[u8], at: usize) -> Result<u64, ElfError> {
    let bytes = data.get(at..at + 8).ok_or(ElfError::TooShort)?;
    Ok(u64::from_le_bytes(bytes.try_into().unwrap()))
}

/// One PT_LOAD program header.
struct Segment {
    flags: u32,
    offset: u64,
    vaddr: u64,
    filesz: u64,
    memsz: u64,
}

/// Check the ELF header; return (entry point, PT_LOAD segments).
fn parse(data: &[u8]) -> Result<(u64, impl Iterator<Item = Result<Segment, ElfError>> + '_), ElfError> {
    if data.len() < 64 {
        return Err(ElfError::TooShort);
    }
    if &data[0..4] != b"\x7fELF" {
        return Err(ElfError::NotElf);
    }
    if data[4] != 2 || data[5] != 1 {
        return Err(ElfError::Not64BitLittleEndian); // EI_CLASS=ELFCLASS64, EI_DATA=LSB
    }
    if read_u16(data, 18)? != 0x3e {
        return Err(ElfError::NotX86_64); // e_machine = EM_X86_64
    }
    if read_u16(data, 16)? != 2 {
        return Err(ElfError::NotExecutable); // e_type = ET_EXEC (static, non-PIE)
    }
    let entry = read_u64(data, 24)?;
    let phoff = read_u64(data, 32)? as usize;
    let phentsize = read_u16(data, 54)? as usize;
    let phnum = read_u16(data, 56)? as usize;
    if phentsize < 56 {
        return Err(ElfError::BadProgramHeader);
    }

    let segments = (0..phnum).filter_map(move |i| {
        let ph = phoff.checked_add(i.checked_mul(phentsize)?)?;
        let segment = (|| {
            Ok(Segment {
                flags: read_u32(data, ph + 4)?,
                offset: read_u64(data, ph + 8)?,
                vaddr: read_u64(data, ph + 16)?,
                filesz: read_u64(data, ph + 32)?,
                memsz: read_u64(data, ph + 40)?,
            })
        })();
        match read_u32(data, ph) {
            Ok(PT_LOAD) => Some(segment),
            Ok(_) => None, // other kinds (GNU_STACK, ...) we ignore
            Err(e) => Some(Err(e)),
        }
    });
    Ok((entry, segments))
}

/// Load the ELF image `data` into address space `l4`. Returns the entry point.
pub fn load(data: &[u8], l4: PhysFrame) -> Result<u64, ElfError> {
    let (entry, segments) = parse(data)?;
    if !(USER_START..USER_END).contains(&entry) {
        return Err(ElfError::OutsideUserSpace);
    }
    let nx_enabled = memory::nx_enabled();
    // The previous segment's last page and its flags.
    let mut previous: Option<(Page<Size4KiB>, PageTableFlags)> = None;

    for segment in segments {
        let seg = segment?;
        let file_end = seg.offset.checked_add(seg.filesz).ok_or(ElfError::BadProgramHeader)?;
        let mem_end = seg.vaddr.checked_add(seg.memsz).ok_or(ElfError::BadProgramHeader)?;
        if seg.filesz > seg.memsz || file_end > data.len() as u64 {
            return Err(ElfError::BadProgramHeader);
        }
        if seg.vaddr < USER_START || mem_end > USER_END {
            return Err(ElfError::OutsideUserSpace);
        }
        if seg.memsz == 0 {
            continue;
        }

        let mut flags = PageTableFlags::empty();
        if seg.flags & PF_W != 0 {
            flags |= PageTableFlags::WRITABLE;
        }
        if seg.flags & PF_X == 0 && nx_enabled {
            flags |= PageTableFlags::NO_EXECUTE; // data can't be run as code
        }

        // Map every page the segment touches, then copy its file bytes in.
        // We write through the kernel's physical-memory window (the new
        // address space isn't active, so its user addresses mean nothing here).
        let first = Page::<Size4KiB>::containing_address(VirtAddr::new(seg.vaddr));
        let last = Page::<Size4KiB>::containing_address(VirtAddr::new(mem_end - 1));
        // Segments come sorted by address. If this one starts on the page
        // where the previous one ended, map_user_page would keep the old
        // page's permissions: e.g. .bss on a read-only .rodata page, and the
        // program dies on its first write. Refuse such files up front.
        if previous.is_some_and(|(page, f)| page >= first && f != flags) {
            return Err(ElfError::SharedPage);
        }
        previous = Some((last, flags));
        for page in Page::range_inclusive(first, last) {
            let frame = memory::map_user_page(l4, page, flags).ok_or(ElfError::OutOfMemory)?;

            // Which part of this page is covered by file bytes?
            let page_start = page.start_address().as_u64();
            let copy_start = page_start.max(seg.vaddr);
            let copy_end = (page_start + 4096).min(seg.vaddr + seg.filesz);
            if copy_start < copy_end {
                let src = &data[(seg.offset + (copy_start - seg.vaddr)) as usize..]
                    [..(copy_end - copy_start) as usize];
                let dst = memory::phys_to_virt(frame.start_address()) + (copy_start - page_start);
                unsafe {
                    core::ptr::copy_nonoverlapping(src.as_ptr(), dst.as_mut_ptr(), src.len());
                }
            }
            // Bytes past filesz (.bss) stay zero: map_user_page zeroes new pages.
        }
    }
    Ok(entry)
}

/// Just validate the headers (used by tests).
pub fn check(data: &[u8]) -> Result<u64, ElfError> {
    let (entry, segments) = parse(data)?;
    for s in segments {
        s?;
    }
    Ok(entry)
}