//! A tiny ELF64 loader: just enough to load a statically linked user program
//! (execve and spawn use it) -- classic or position independent.
//!
//! 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`.
//!
//! Stage 19: **static PIE** ("position-independent executable"). Modern
//! toolchains make PIE files by default (e_type ET_DYN): their addresses
//! start at 0 and the program can run anywhere, because it fixes up its own
//! pointers at start-up (musl's `rcrt1.o` does that). We pick a *load
//! bias* -- a random page-aligned base, so every run lands somewhere else
//! (address space layout randomization) -- add it to every address in the
//! file, and tell the program in its auxiliary vector (AT_PHDR, AT_ENTRY).
//! A PIE file that names a *dynamic loader* (PT_INTERP, e.g.
//! /lib/ld-musl-x86_64.so.1) needs shared libraries we don't have: ENOEXEC.
//! Other program headers we look at: PT_GNU_STACK (may the stack be
//! executable?), PT_TLS and PT_GNU_RELRO (just accepted: the C library
//! sets up thread-local storage and makes RELRO read-only by itself).
//!
//! 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,
    /// A dynamically linked program (PT_INTERP): needs a dynamic loader.
    NeedsInterpreter,
}

const PT_LOAD: u32 = 1;
const PT_INTERP: u32 = 3;
const PT_TLS: u32 = 7;
const PT_GNU_STACK: u32 = 0x6474_e551;
const PT_GNU_RELRO: u32 = 0x6474_e552;
const ET_EXEC: u16 = 2;
const ET_DYN: u16 = 3;
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,
}

/// What the ELF header and the program headers say.
struct Parsed {
    /// ET_DYN (PIE): may be loaded anywhere.
    relocatable: bool,
    entry: u64,
    /// The PT_LOAD segments, sorted by address.
    segments: Vec<Segment>,
    /// PT_GNU_STACK asks for an executable stack.
    exec_stack: bool,
}

/// Check the ELF header and read the program headers.
fn parse(data: &[u8]) -> Result<Parsed, 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
    }
    let relocatable = match read_u16(data, 16)? {
        ET_EXEC => false, // a classic static executable: fixed addresses
        ET_DYN => true,   // PIE (stage 19)
        _ => return Err(ElfError::NotExecutable),
    };
    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 mut parsed = Parsed { relocatable, entry, segments: Vec::new(), exec_stack: false };
    for i in 0..phnum {
        let ph = i.checked_mul(phentsize).and_then(|o| o.checked_add(phoff)).ok_or(ElfError::BadProgramHeader)?;
        let flags = read_u32(data, ph + 4)?;
        match read_u32(data, ph)? {
            PT_LOAD => parsed.segments.push(Segment {
                flags,
                offset: read_u64(data, ph + 8)?,
                vaddr: read_u64(data, ph + 16)?,
                filesz: read_u64(data, ph + 32)?,
                memsz: read_u64(data, ph + 40)?,
            }),
            PT_INTERP => return Err(ElfError::NeedsInterpreter),
            PT_GNU_STACK => parsed.exec_stack = flags & PF_X != 0,
            PT_TLS | PT_GNU_RELRO => {} // the C library handles these
            _ => {}                     // others (NOTE, PHDR, ...) we ignore
        }
    }
    Ok(parsed)
}

/// 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,
    /// Load bias: what was added to every address in the file (0 for a
    /// classic ET_EXEC, the random base for a PIE).
    pub bias: u64,
    /// PT_GNU_STACK said the stack may hold code (old programs need it).
    pub exec_stack: bool,
}

/// Where a PIE goes: 4 MiB above the start of user space plus a random
/// number of pages (up to 1 GiB), so the heap after it and mmap from the
/// top still have hundreds of GiB between them.
const PIE_BASE: u64 = USER_START + 0x40_0000;
const PIE_RANDOM_PAGES: u64 = 0x4_0000;

/// A random load bias for a PIE whose lowest segment starts at `low`.
fn pie_bias(low: u64) -> u64 {
    let mut random = [0u8; 8];
    crate::process::random_bytes(&mut random);
    PIE_BASE + (u64::from_le_bytes(random) % PIE_RANDOM_PAGES) * 4096 - (low & !0xfff)
}

/// 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.
///
/// `random`: pick a random bias for a PIE (else the lowest one, enough to
/// check the file).
fn plan(data: &[u8], random: bool) -> Result<(LoadInfo, Vec<Segment>), ElfError> {
    let Parsed { relocatable, entry, mut segments, exec_stack } = parse(data)?;
    let mut bias = 0;
    if relocatable {
        let low = segments.iter().map(|s| s.vaddr).min().unwrap_or(0);
        if low >= USER_START {
            return Err(ElfError::OutsideUserSpace); // not really position independent
        }
        bias = if random { pie_bias(low) } else { PIE_BASE - (low & !0xfff) };
        for seg in &mut segments {
            seg.vaddr = seg.vaddr.wrapping_add(bias);
        }
    }
    let entry = entry.wrapping_add(bias);
    if !(USER_START..USER_END).contains(&entry) {
        return Err(ElfError::OutsideUserSpace);
    }
    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, bias, exec_stack };
    // 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::MMAP_TOP {
            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, true)?;
    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, false).map(|(info, _)| info.entry)
}