//! A tiny ELF64 loader: just enough to load a statically linked user program
//! (execve and spawn use it).
//!
//! 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 crate::mm::{Kind, Mm, Vma, PROT_EXEC, PROT_READ, PROT_WRITE};
use alloc::vec::Vec;
use x86_64::structures::paging::{Page, PageTableFlags, 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))
}

/// What `load` tells execve about the program (for the auxiliary vector).
pub struct LoadInfo {
    pub entry: u64,
    /// End of the highest segment: the heap (brk) starts on the next page.
    pub end: u64,
    /// Where the program headers are in the loaded image, if a segment
    /// includes them (AT_PHDR). Otherwise execve copies them to the stack.
    pub phdr: Option<u64>,
    /// File offset, entry size and count of the program headers.
    pub phoff: u64,
    pub phent: u64,
    pub phnum: u64,
}

/// Parse and check everything (headers, segment bounds, permissions) without
/// mapping anything: execve calls this *before* it destroys the old program,
/// so a bad file can still be reported as an error.
fn plan(data: &[u8]) -> Result<(LoadInfo, Vec<Segment>), ElfError> {
    let (entry, segments) = parse(data)?;
    if !(USER_START..USER_END).contains(&entry) {
        return Err(ElfError::OutsideUserSpace);
    }
    let segments = segments.collect::<Result<Vec<Segment>, ElfError>>()?;
    let (phoff, phent, phnum) = (read_u64(data, 32)?, read_u16(data, 54)? as u64, read_u16(data, 56)? as u64);
    let mut info = LoadInfo { entry, end: USER_START, phdr: None, phoff, phent, phnum };
    // The previous segment's last page and its flags.
    let mut previous: Option<(u64, u32)> = None;
    for seg in &segments {
        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);
        }
        // Segments must stay below the user stack and its guard zone.
        if seg.vaddr < USER_START || mem_end > crate::mm::USER_STACK_ZONE_START {
            return Err(ElfError::OutsideUserSpace);
        }
        if seg.memsz == 0 {
            continue;
        }
        // Segments come sorted by address. If this one starts on the page
        // where the previous one ended, both would get one page with one
        // set of permissions: e.g. .bss on a read-only .rodata page, and the
        // program dies on its first write. Refuse such files up front.
        let perms = seg.flags & (PF_W | PF_X);
        if previous.is_some_and(|(last_page, f)| last_page >= seg.vaddr & !0xfff && f != perms) {
            return Err(ElfError::SharedPage);
        }
        previous = Some(((mem_end - 1) & !0xfff, perms));
        info.end = info.end.max(mem_end);
        // Are the program headers inside this segment's file bytes?
        let ph_end = phoff + phent * phnum;
        if seg.offset <= phoff && ph_end <= file_end {
            info.phdr = Some(seg.vaddr + (phoff - seg.offset));
        }
    }
    Ok((info, segments))
}

/// Load the ELF image `data` into the (empty) address space `mm`: map and
/// fill every segment (eagerly -- programs are small), and record each as
/// a region.
pub fn load(data: &[u8], mm: &mut Mm) -> Result<LoadInfo, ElfError> {
    let (info, segments) = plan(data)?;
    let nx_enabled = memory::nx_enabled();
    for seg in segments.iter().filter(|s| s.memsz > 0) {
        let mem_end = seg.vaddr + seg.memsz;
        let mut flags = PageTableFlags::empty();
        let mut prot = PROT_READ;
        if seg.flags & PF_W != 0 {
            flags |= PageTableFlags::WRITABLE;
            prot |= PROT_WRITE;
        }
        if seg.flags & PF_X == 0 && nx_enabled {
            flags |= PageTableFlags::NO_EXECUTE; // data can't be run as code
        }
        if seg.flags & PF_X != 0 {
            prot |= PROT_EXEC;
        }

        // 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 may not be the active one).
        let first = Page::<Size4KiB>::containing_address(VirtAddr::new(seg.vaddr));
        let last = Page::<Size4KiB>::containing_address(VirtAddr::new(mem_end - 1));
        for page in Page::range_inclusive(first, last) {
            let frame = memory::map_user_page(mm.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.
        }
        // The region (a page shared with the previous segment -- same
        // permissions, checked above -- already belongs to that one).
        let mut start = first.start_address().as_u64();
        let end = last.start_address().as_u64() + 4096;
        if let Some(prev) = mm.vmas().last() {
            start = start.max(prev.end);
        }
        if start < end {
            mm.insert(Vma::new(start, end, prot, Kind::Program));
        }
    }
    Ok(info)
}

/// Just validate the file (execve before the point of no return, tests).
/// Returns the entry point.
pub fn check(data: &[u8]) -> Result<u64, ElfError> {
    plan(data).map(|(info, _)| info.entry)
}