// Copyright 2023-Present Jonas Kruckenberg
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

use core::alloc::Layout;
use core::range::Range;
use core::slice;

use arrayvec::ArrayVec;
use loader_api::TlsTemplate;
use mem_core::{FrameAllocator, PhysicalAddress, VirtualAddress};
use object::LittleEndian;
use object::elf::{
    DT_REL, DT_RELA, DT_RELAENT, DT_RELASZ, DT_RELENT, DT_RELR, DT_RELRENT, DT_RELRSZ, DT_RELSZ,
    Dyn64, FileHeader64, PF_R, PF_W, PF_X, PT_DYNAMIC, PT_GNU_RELRO, PT_LOAD, PT_TLS,
    ProgramHeader64, R_RISCV_RELATIVE, Rela64,
};
use object::read::elf::{Dyn, FileHeader, ProgramHeader, Rela};

use crate::{Error, ensure};

// 64 segments ought to be enough for anyone!
const MAX_LOAD_SEGMENTS: usize = 64;

pub trait ImageSource {
    /// Total length of the backing image, in bytes.
    fn len(&self) -> u64;
    /// Fill `dst` from the image starting at `offset`.
    ///
    /// # Errors
    ///
    /// Should return `Err` if the number of bytes read DID NOT equal
    /// `dst.len()`.
    fn read_at(&mut self, offset: u64, dst: &mut [u8]) -> crate::Result<()>;
}

pub struct KernelImage {
    /// The size of the kernel in memory
    size: usize,
    /// The entrypoint as reported by the ELF, offset relative to the in-memory
    /// image
    pub entry: usize,
    pub load: ArrayVec<LoadSegment, MAX_LOAD_SEGMENTS>,
    tls: TlsTemplate,
    dynamic: ProgramHeader64<LittleEndian>,
    /// The range (offsets into in-memory image) that must be marked as
    /// read-only after applying relocations.
    relro: Range<usize>,
}

pub struct Kernel<S> {
    kernel: S,
    debug_info: Option<S>,
    image: KernelImage,
}

impl<S: ImageSource> Kernel<S> {
    /// Parse the located kernel payload.
    /// When this completes successfully we have validated the file structure
    /// and parsed out relevant information.
    pub fn from_parts(mut kernel: S, debug_info: Option<S>, granule: usize) -> crate::Result<Self> {
        let mut buf = [0; 4096];
        kernel.read_at(0, &mut buf)?;

        let header = FileHeader64::<LittleEndian>::parse(buf.as_slice())?;
        log::debug!("kernel header {header:?}");

        ensure!(header.is_supported(), Error::MalformedImage);
        ensure!(header.is_class_64(), Error::MalformedImage);
        ensure!(header.is_little_endian(), Error::MalformedImage);

        #[cfg(debug_assertions)]
        {
            let phtable_end = header.e_phoff(LittleEndian)
                + u64::from(header.e_phnum(LittleEndian))
                    * u64::from(header.e_phentsize(LittleEndian));
            assert!(phtable_end <= buf.len() as u64);
        }

        let mut load = ArrayVec::new();
        let mut dynamic = None;
        let mut tls = None;
        let mut relro = None;

        for ph in header.program_headers(LittleEndian, buf.as_slice())? {
            match ph.p_type(LittleEndian) {
                PT_LOAD => {
                    load.try_push(parse_load_segment(ph, granule)?)
                        .map_err(|_| Error::MalformedImage)?;
                }
                PT_TLS => {
                    ensure!(
                        tls.replace(parse_tls_segment(ph)?).is_none(),
                        Error::MalformedImage
                    );
                }
                PT_DYNAMIC => ensure!(dynamic.replace(*ph).is_none(), Error::MalformedImage),
                PT_GNU_RELRO => {
                    let start = usize::try_from(ph.p_vaddr(LittleEndian))?;
                    let len = usize::try_from(ph.p_memsz(LittleEndian))?;
                    let range = Range::from(start..start + len);

                    ensure!(relro.replace(range).is_none(), Error::MalformedImage);
                }
                _ => continue,
            }
        }

        ensure!(
            load.is_sorted_by_key(|seg| seg.image_offset),
            Error::MalformedImage
        );
        ensure!(
            load.first().ok_or(Error::MalformedImage)?.image_offset == 0,
            Error::MalformedImage
        );

        let size = {
            let last_seg = load.last().ok_or(Error::MalformedImage)?;
            let image_end = last_seg
                .image_offset
                .checked_add(last_seg.mem_size)
                .ok_or(Error::FieldOutOfRange)?;

            image_end
                .checked_next_multiple_of(granule)
                .ok_or(Error::FieldOutOfRange)?
        };

        log::trace!("kernel image size: {size} bytes");

        Ok(Kernel {
            kernel,
            debug_info,
            image: KernelImage {
                size,
                entry: usize::try_from(header.e_entry(LittleEndian))?,
                load,
                tls: tls.ok_or(Error::MissingSegment)?,
                dynamic: dynamic.ok_or(Error::MissingSegment)?,
                relro: relro.ok_or(Error::MissingSegment)?,
            },
        })
    }

    /// Stage the kernel into physical memory.
    /// When this completes successfully we have allocated physical memory and
    /// initialized it (from disk or zeroed as appropriate).
    pub fn stage(
        mut self,
        frame_alloc: &impl FrameAllocator,
        granule: usize,
    ) -> crate::Result<StagedKernel> {
        let alloc = |size: usize| -> crate::Result<&mut [u8]> {
            log::trace!("allocating {size} physmem bytes");

            let base =
                frame_alloc.allocate_contiguous(Layout::from_size_align(size, granule).unwrap())?;

            assert!(!base.is_null());

            // Safety: we just allocated the memory, `allocate_contiguous` ensures the
            // memory range is valid and initialized
            Ok(unsafe { slice::from_raw_parts_mut(base.as_mut_ptr(), size) })
        };

        let staging = alloc(self.image.size)?;

        for seg in &self.image.load {
            // NB: explain the filesz<->memsz difference, BSS and how if the BSS ends before
            // the page boundary we potentially end up with uninitialized memory being
            // mapped which is of course not a great idea from a security POV.
            // We therefore just zero everything up to the next page boundary.
            let span = seg.unaligned_mem_range();
            let seg_end = span
                .end
                .checked_next_multiple_of(granule)
                .ok_or(Error::FieldOutOfRange)?;

            log::trace!("attempting to read {:?}", span.start..seg_end);
            let dst = staging
                .get_mut(span.start..seg_end)
                .ok_or(Error::FieldOutOfRange)?;

            if seg.file_size > 0 {
                self.kernel
                    .read_at(seg.file_offset, &mut dst[..seg.file_size])?;
            }

            dst[seg.file_size..].fill(0); // BSS + over-zero to page boundary
        }

        log::trace!("staging kernel debug info...");
        let debug_info = if let Some(mut debug_info) = self.debug_info {
            let size = usize::try_from(debug_info.len())?;
            let buf = alloc(size)?;
            debug_info.read_at(0, buf)?;

            Some(buf)
        } else {
            None
        };

        Ok(StagedKernel {
            kernel: staging,
            debug_info,
            image: self.image,
        })
    }
}

pub struct StagedKernel {
    kernel: &'static mut [u8],
    debug_info: Option<&'static mut [u8]>,
    image: KernelImage,
}

struct RelaInfo {
    offset: usize,
    count: usize,
}

impl StagedKernel {
    pub fn size(&self) -> usize {
        self.image.size
    }

    pub fn tls_template(&self) -> &TlsTemplate {
        &self.image.tls
    }

    /// Apply relocations to the staged kernel.
    /// When this completes successfully the kernel images is ready for
    /// execution _within_ its address space.
    pub fn relocate(self, kernel_virt: Range<VirtualAddress>) -> crate::Result<RelocatedKernel> {
        let rela_info = self.parse_rela()?.ok_or(Error::MalformedImage)?;

        for i in 0..rela_info.count {
            let off = rela_info.offset + i * size_of::<Rela64<LittleEndian>>();
            let raw = self.kernel.get(off..).ok_or(Error::FieldOutOfRange)?;
            let (rela, _) = object::pod::from_bytes::<Rela64<LittleEndian>>(raw)
                .map_err(|_| Error::MalformedImage)?;

            match rela.r_type(LittleEndian, false) {
                R_RISCV_RELATIVE => {
                    // dynamic relocation offsets are relative to the virtual
                    // layout of the elf, not the physical file
                    let target = usize::try_from(rela.r_offset(LittleEndian))?;
                    let value = kernel_virt
                        .start
                        .offset(isize::try_from(rela.r_addend(LittleEndian))?);
                    ensure!(kernel_virt.contains(&value), Error::MalformedImage);

                    self.kernel
                        .get_mut(target..target + 8)
                        .ok_or(Error::FieldOutOfRange)?
                        .copy_from_slice(&value.get().to_le_bytes());
                }
                ty => {
                    log::error!("unsupported relocation type {ty}");
                    return Err(Error::MalformedImage);
                }
            }
        }

        Ok(RelocatedKernel {
            kernel: self.kernel,
            debug_info: self.debug_info,
            image: self.image,
            virt: kernel_virt,
        })
    }

    fn parse_rela(&self) -> crate::Result<Option<RelaInfo>> {
        // NB: not using `object`s `.dynamic` method here because we access the segment
        // through the staging buffer which requires memory-offsets, not
        // file-offsets like `.dynamic` uses.
        let fields = {
            let vaddr = usize::try_from(self.image.dynamic.p_vaddr(LittleEndian))?;
            let filesz = usize::try_from(self.image.dynamic.p_filesz(LittleEndian))?;

            object::pod::slice_from_all_bytes::<Dyn64<LittleEndian>>(
                self.kernel
                    .get(vaddr..vaddr + filesz)
                    .ok_or(Error::FieldOutOfRange)?,
            )
            .map_err(|_| Error::MalformedImage)?
        };

        let mut rela = None; // Address of Rela relocs
        let mut rela_size = None; // Total size of Rela relocs
        let mut rela_ent = None; // Size of one Rela reloc

        for field in fields {
            let property = match field.tag(LittleEndian) {
                DT_RELA => &mut rela,
                DT_RELASZ => &mut rela_size,
                DT_RELAENT => &mut rela_ent,
                DT_REL | DT_RELSZ | DT_RELENT => {
                    log::error!("REL relocations are not supported");
                    return Err(Error::MalformedImage);
                }
                DT_RELR | DT_RELRSZ | DT_RELRENT => {
                    log::error!("RELR relocations are not supported");
                    return Err(Error::MalformedImage);
                }
                _ => continue,
            };

            let val = usize::try_from(field.val(LittleEndian))?;
            let prev = property.replace(val);
            ensure!(prev.is_none(), Error::MalformedImage);
        }

        if rela.is_none() && (rela_size.is_some() || rela_ent.is_some()) {
            log::error!("Rela entry is missing but RelaSize or RelaEnt have been provided");
            return Err(Error::MalformedImage);
        }

        // No RELA field means no relocations at all
        let Some(offset) = rela else {
            return Ok(None);
        };

        let total_size = rela_size.ok_or(Error::MalformedImage)?;
        let entry_size = rela_ent.ok_or(Error::MalformedImage)?;

        ensure!(
            entry_size == size_of::<Rela64<LittleEndian>>(),
            Error::MalformedImage
        );

        Ok(Some(RelaInfo {
            offset,
            count: total_size / entry_size,
        }))
    }
}

pub struct RelocatedKernel {
    kernel: &'static mut [u8],
    debug_info: Option<&'static mut [u8]>,
    image: KernelImage,
    virt: Range<VirtualAddress>,
}

impl RelocatedKernel {
    /// Physical base address of the staged kernel image.
    ///
    /// Valid only while UEFI boot services are active: before
    /// `ExitBootServices` the staging buffer is identity-mapped, so the
    /// pointer's address *is* its physical address.
    pub fn phys_base(&self) -> PhysicalAddress {
        PhysicalAddress::new(self.kernel.as_ptr().addr())
    }

    pub fn entry(&self) -> VirtualAddress {
        let entry = self.virt.start.add(self.image.entry);
        debug_assert!(self.virt.contains(&entry));
        entry
    }

    pub fn debug_info_phys(&self) -> Option<Range<PhysicalAddress>> {
        let range = self.debug_info.as_ref()?.as_ptr_range();
        Some(Range::from(
            PhysicalAddress::from_ptr(range.start)..PhysicalAddress::from_ptr(range.end),
        ))
    }

    pub fn load_segments(&self) -> impl ExactSizeIterator<Item = &'_ LoadSegment> {
        self.image.load.iter()
    }

    pub fn tls_template(&self) -> &TlsTemplate {
        &self.image.tls
    }

    pub fn relro_range(&self) -> Range<usize> {
        self.image.relro
    }

    /// Allocate the boot hart's TLS block and copy the `.tdata` template into
    /// it.
    ///
    /// This copies out of the staging buffer rather than the file, and must run
    /// *after* relocation: `.tdata` can itself carry `R_RISCV_RELATIVE`
    /// relocations, and those are only applied to the staging buffer. Copying
    /// from disk would yield an unrelocated — and therefore corrupt — template.
    pub fn instantiate_tls_block(
        &self,
        frame_alloc: &impl FrameAllocator,
        page_size: usize,
    ) -> crate::Result<&'static mut [u8]> {
        let seg = &self.image.tls;

        let tls_size = seg.mem_size.next_multiple_of(page_size);

        log::trace!("allocating {tls_size} bytes for boot hart TLS block");
        let block = frame_alloc
            .allocate_contiguous(Layout::from_size_align(tls_size, page_size).unwrap())?;

        // Safety: we just allocated with `allocate_contiguous`.
        let block = unsafe { slice::from_raw_parts_mut(block.as_mut_ptr(), tls_size) };

        // `.tdata` lives at `image_offset` in the staged (and now relocated) image.
        let tdata = self
            .kernel
            .get(seg.image_offset..seg.image_offset + seg.file_size)
            .ok_or(Error::FieldOutOfRange)?;
        block[..seg.file_size].copy_from_slice(tdata);
        block[seg.file_size..].fill(0); // .tbss + over-zero to page boundary

        Ok(block)
    }
}

/// Access permissions of a `PT_LOAD` segment, decoded from the ELF `p_flags`.
#[derive(Debug, Clone, Copy)]
pub enum Permissions {
    ReadOnly,
    ReadWrite,
    ReadExecute,
}

pub struct LoadSegment {
    /// Offset into the in-memory image where this segment starts
    pub image_offset: usize,
    /// Size of the segment in memory
    pub mem_size: usize,
    /// Offset into the on-disk file where this segment starts
    pub file_offset: u64,
    /// Size of the segment on disk
    pub file_size: usize,
    pub perms: Permissions,
}

impl LoadSegment {
    /// Byte range this segment occupies in the in-memory image — the
    /// file-backed bytes followed by the zero-initialized BSS tail.
    pub fn unaligned_mem_range(&self) -> Range<usize> {
        Range::from(self.image_offset..self.image_offset + self.mem_size)
    }
}

fn parse_load_segment(
    ph: &ProgramHeader64<LittleEndian>,
    granule: usize,
) -> crate::Result<LoadSegment> {
    ensure!(
        usize::try_from(ph.p_align(LittleEndian))? == granule,
        Error::MalformedImage
    );

    let image_offset = usize::try_from(ph.p_vaddr(LittleEndian))?;
    let file_offset = ph.p_offset(LittleEndian);

    let file_size = usize::try_from(ph.p_filesz(LittleEndian))?;
    let mem_size = usize::try_from(ph.p_memsz(LittleEndian))?;

    const RW: u32 = PF_R | PF_W;
    const RX: u32 = PF_R | PF_X;
    let flags = ph.p_flags(LittleEndian);
    let perms = match flags {
        PF_R => Permissions::ReadOnly,
        RW => Permissions::ReadWrite,
        RX => Permissions::ReadExecute,
        _ => return Err(Error::InvalidSegmentFlags(flags)),
    };

    Ok(LoadSegment {
        image_offset,
        mem_size,
        file_offset,
        file_size,
        perms,
    })
}

fn parse_tls_segment(ph: &ProgramHeader64<LittleEndian>) -> crate::Result<TlsTemplate> {
    let image_offset = usize::try_from(ph.p_vaddr(LittleEndian))?;

    let file_size = usize::try_from(ph.p_filesz(LittleEndian))?;
    let mem_size = usize::try_from(ph.p_memsz(LittleEndian))?;

    let align = usize::try_from(ph.p_align(LittleEndian))?;

    Ok(TlsTemplate {
        image_offset,
        mem_size,
        file_size,
        align,
    })
}
