#![no_std]

pub const BOOT_INFO_MAGIC: [u8; 8] = *b"TURKOSBI";
pub const BOOT_INFO_VERSION: u32 = 1;
pub const BOOT_INFO_PAGE_BYTES: u64 = 4096;
pub const BOOT_INFO_MAP_CAPACITY: u64 = 64 * 1024;
pub const MAX_EXIT_BOOT_SERVICES_ATTEMPTS: u8 = 3;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum MapCopyError {
    Length,
    Stride,
    Capacity,
}

/// Copy into caller-owned backing; do not retain a firmware buffer reference.
/// Validate all lengths before touching the target and clear its unused tail.
pub fn copy_owned_memory_map(
    source: &[u8],
    target: &mut [u8],
    stride: usize,
) -> Result<usize, MapCopyError> {
    if stride < 40 || stride % 8 != 0 {
        return Err(MapCopyError::Stride);
    }
    if source.is_empty() || source.len() % stride != 0 {
        return Err(MapCopyError::Length);
    }
    if source.len() > BOOT_INFO_MAP_CAPACITY as usize || source.len() > target.len() {
        return Err(MapCopyError::Capacity);
    }
    target[..source.len()].copy_from_slice(source);
    target[source.len()..].fill(0);
    Ok(source.len())
}

#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct ExitBootServicesRetry {
    used: u8,
}

impl ExitBootServicesRetry {
    pub const fn new() -> Self {
        Self { used: 0 }
    }

    pub fn next_attempt(&mut self) -> Option<u8> {
        if self.used == MAX_EXIT_BOOT_SERVICES_ATTEMPTS {
            return None;
        }
        self.used += 1;
        Some(self.used)
    }
}

/// Wire layout shared by the UEFI application and the future kernel handoff.
/// Unpopulated payload fields remain zero until their separate acceptance gates.
#[repr(C)]
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct BootInfoV1 {
    pub magic: [u8; 8],
    pub version: u32,
    pub struct_bytes: u32,
    pub memory_map_ptr: u64,
    pub memory_map_bytes: u64,
    pub descriptor_bytes: u32,
    pub descriptor_version: u32,
    pub kernel_phys_start: u64,
    pub kernel_file_bytes: u64,
    pub dtb_phys_addr: u64,
    pub flags: u64,
}

const _: () = assert!(core::mem::size_of::<BootInfoV1>() == 72);
const _: () = assert!(core::mem::align_of::<BootInfoV1>() == 8);
// K501: freeze every V1 field offset, not only the aggregate size. A future
// field reorder with the same 72-byte size must not silently change the wire ABI.
const _: () = assert!(core::mem::offset_of!(BootInfoV1, magic) == 0);
const _: () = assert!(core::mem::offset_of!(BootInfoV1, version) == 8);
const _: () = assert!(core::mem::offset_of!(BootInfoV1, struct_bytes) == 12);
const _: () = assert!(core::mem::offset_of!(BootInfoV1, memory_map_ptr) == 16);
const _: () = assert!(core::mem::offset_of!(BootInfoV1, memory_map_bytes) == 24);
const _: () = assert!(core::mem::offset_of!(BootInfoV1, descriptor_bytes) == 32);
const _: () = assert!(core::mem::offset_of!(BootInfoV1, descriptor_version) == 36);
const _: () = assert!(core::mem::offset_of!(BootInfoV1, kernel_phys_start) == 40);
const _: () = assert!(core::mem::offset_of!(BootInfoV1, kernel_file_bytes) == 48);
const _: () = assert!(core::mem::offset_of!(BootInfoV1, dtb_phys_addr) == 56);
const _: () = assert!(core::mem::offset_of!(BootInfoV1, flags) == 64);

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum BootInfoPointerError {
    Missing,
    Unaligned,
    AddressOverflow,
}

/// K503: validate only the numeric address before any unsafe dereference.
/// This does not prove that the address is mapped, owned, or kernel-readable.
pub fn validate_boot_info_pointer(address: u64) -> Result<(), BootInfoPointerError> {
    if address == 0 {
        return Err(BootInfoPointerError::Missing);
    }
    if address % core::mem::align_of::<BootInfoV1>() as u64 != 0 {
        return Err(BootInfoPointerError::Unaligned);
    }
    if address
        .checked_add(core::mem::size_of::<BootInfoV1>() as u64)
        .is_none()
    {
        return Err(BootInfoPointerError::AddressOverflow);
    }
    Ok(())
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum BootInfoHeaderError {
    Magic,
    Version,
    Size,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum BootInfoMemoryMapError {
    Missing,
    DescriptorSize,
    DescriptorVersion,
    Length,
    AddressOverflow,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum BootInfoHandoffError {
    Header,
    Pointer,
    Map,
    Range,
    Overlap,
    Fields,
}

fn checked_range(start: u64, bytes: u64, outer_start: u64, outer_end: u64) -> Option<u64> {
    let end = start.checked_add(bytes)?;
    (bytes != 0 && start >= outer_start && end <= outer_end).then_some(end)
}

fn disjoint(a_start: u64, a_end: u64, b_start: u64, b_end: u64) -> bool {
    a_end <= b_start || b_end <= a_start
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum KernelFileLengthError {
    Empty,
    TooLarge,
    DeclaredMismatch,
    ShortRead,
    OverRead,
}

pub fn validate_kernel_file_length(
    expected: u64,
    declared: u64,
    actual_read: Option<u64>,
    capacity: u64,
) -> Result<(), KernelFileLengthError> {
    if expected == 0 || declared == 0 {
        return Err(KernelFileLengthError::Empty);
    }
    if expected > capacity || declared > capacity {
        return Err(KernelFileLengthError::TooLarge);
    }
    if declared != expected {
        return Err(KernelFileLengthError::DeclaredMismatch);
    }
    if let Some(actual) = actual_read {
        if actual < expected {
            return Err(KernelFileLengthError::ShortRead);
        }
        if actual > expected {
            return Err(KernelFileLengthError::OverRead);
        }
    }
    Ok(())
}

impl BootInfoV1 {
    pub const fn empty() -> Self {
        Self {
            magic: BOOT_INFO_MAGIC,
            version: BOOT_INFO_VERSION,
            struct_bytes: core::mem::size_of::<Self>() as u32,
            memory_map_ptr: 0,
            memory_map_bytes: 0,
            descriptor_bytes: 0,
            descriptor_version: 0,
            kernel_phys_start: 0,
            kernel_file_bytes: 0,
            dtb_phys_addr: 0,
            flags: 0,
        }
    }

    pub fn validate_header(&self) -> Result<(), BootInfoHeaderError> {
        if self.magic != BOOT_INFO_MAGIC {
            return Err(BootInfoHeaderError::Magic);
        }
        if self.version != BOOT_INFO_VERSION {
            return Err(BootInfoHeaderError::Version);
        }
        if self.struct_bytes != core::mem::size_of::<Self>() as u32 {
            return Err(BootInfoHeaderError::Size);
        }
        Ok(())
    }

    /// Checks metadata without dereferencing an address supplied by firmware.
    /// Ownership and descriptor contents are separate gates.
    pub fn validate_memory_map_metadata(&self) -> Result<(), BootInfoMemoryMapError> {
        if self.memory_map_ptr == 0 || self.memory_map_bytes == 0 {
            return Err(BootInfoMemoryMapError::Missing);
        }
        if self.descriptor_bytes < 40 || self.descriptor_bytes % 8 != 0 {
            return Err(BootInfoMemoryMapError::DescriptorSize);
        }
        if self.descriptor_version != 1 {
            return Err(BootInfoMemoryMapError::DescriptorVersion);
        }
        if self.memory_map_ptr % 8 != 0
            || self.memory_map_bytes % u64::from(self.descriptor_bytes) != 0
        {
            return Err(BootInfoMemoryMapError::Length);
        }
        if self
            .memory_map_ptr
            .checked_add(self.memory_map_bytes)
            .is_none()
        {
            return Err(BootInfoMemoryMapError::AddressOverflow);
        }
        Ok(())
    }

    /// Validate all numeric bounds before the kernel reads the firmware map.
    /// The caller must first prove that the 72-byte header itself is readable.
    pub fn validate_handoff(
        &self,
        pointer: u64,
        ram_start: u64,
        ram_bytes: u64,
        kernel_start: u64,
        kernel_bytes: u64,
        dtb_start: u64,
        dtb_bytes: u64,
    ) -> Result<(), BootInfoHandoffError> {
        self.validate_header()
            .map_err(|_| BootInfoHandoffError::Header)?;
        validate_boot_info_pointer(pointer).map_err(|_| BootInfoHandoffError::Pointer)?;
        self.validate_memory_map_metadata()
            .map_err(|_| BootInfoHandoffError::Map)?;
        if pointer % BOOT_INFO_PAGE_BYTES != 0
            || self.memory_map_ptr % BOOT_INFO_PAGE_BYTES != 0
            || self.memory_map_bytes > BOOT_INFO_MAP_CAPACITY
            || self.kernel_phys_start != kernel_start
            || self.dtb_phys_addr != dtb_start
            || self.kernel_file_bytes == 0
            || self.kernel_file_bytes > kernel_bytes
            || self.flags != 0
        {
            return Err(BootInfoHandoffError::Fields);
        }
        let ram_end = ram_start
            .checked_add(ram_bytes)
            .ok_or(BootInfoHandoffError::Range)?;
        let info_end = checked_range(pointer, BOOT_INFO_PAGE_BYTES, ram_start, ram_end)
            .ok_or(BootInfoHandoffError::Range)?;
        let map_end = checked_range(
            self.memory_map_ptr,
            BOOT_INFO_MAP_CAPACITY,
            ram_start,
            ram_end,
        )
        .ok_or(BootInfoHandoffError::Range)?;
        let kernel_end = checked_range(kernel_start, kernel_bytes, ram_start, ram_end)
            .ok_or(BootInfoHandoffError::Range)?;
        let dtb_end = checked_range(dtb_start, dtb_bytes, ram_start, ram_end)
            .ok_or(BootInfoHandoffError::Range)?;
        for (a_start, a_end, b_start, b_end) in [
            (pointer, info_end, self.memory_map_ptr, map_end),
            (pointer, info_end, kernel_start, kernel_end),
            (pointer, info_end, dtb_start, dtb_end),
            (self.memory_map_ptr, map_end, kernel_start, kernel_end),
            (self.memory_map_ptr, map_end, dtb_start, dtb_end),
            // K586 fuzz: K502 also requires the loaded Image and the DTB to be disjoint.
            (kernel_start, kernel_end, dtb_start, dtb_end),
        ] {
            if !disjoint(a_start, a_end, b_start, b_end) {
                return Err(BootInfoHandoffError::Overlap);
            }
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn k504_k515_handoff_bounds_and_ownership_are_fail_closed() {
        let mut info = BootInfoV1::empty();
        info.memory_map_ptr = 0x4f00_0000;
        info.memory_map_bytes = 48;
        info.descriptor_bytes = 48;
        info.descriptor_version = 1;
        info.kernel_phys_start = 0x4008_0000;
        info.kernel_file_bytes = 0x20_000;
        info.dtb_phys_addr = 0x4e00_0000;
        let check = |value: &BootInfoV1, pointer| {
            value.validate_handoff(
                pointer,
                0x4000_0000,
                0x1000_0000,
                0x4008_0000,
                0x30_000,
                0x4e00_0000,
                0x2000,
            )
        };
        assert_eq!(check(&info, 0x4f10_0000), Ok(()));
        assert_eq!(check(&info, 0), Err(BootInfoHandoffError::Pointer));
        assert_eq!(check(&info, 0x4f10_0008), Err(BootInfoHandoffError::Fields));
        let mut bad = info;
        bad.version = 2;
        assert_eq!(check(&bad, 0x4f10_0000), Err(BootInfoHandoffError::Header));
        bad = info;
        bad.memory_map_bytes = BOOT_INFO_MAP_CAPACITY + 32;
        assert_eq!(check(&bad, 0x4f10_0000), Err(BootInfoHandoffError::Fields));
        bad = info;
        bad.memory_map_ptr = 0x4f10_0000;
        assert_eq!(check(&bad, 0x4f10_0000), Err(BootInfoHandoffError::Overlap));
        bad = info;
        bad.memory_map_ptr = 0x4008_0000;
        assert_eq!(check(&bad, 0x4f10_0000), Err(BootInfoHandoffError::Overlap));
        bad = info;
        bad.memory_map_ptr = 0x4fff_0000;
        assert_eq!(check(&bad, 0x4f10_0000), Ok(()));
        bad.memory_map_ptr = 0x4fff_1000;
        assert_eq!(check(&bad, 0x4f10_0000), Err(BootInfoHandoffError::Range));
    }

    #[test]
    fn k501_v1_layout_keeps_field_offsets_and_empty_payload() {
        assert_eq!(core::mem::size_of::<BootInfoV1>(), 72);
        assert_eq!(core::mem::align_of::<BootInfoV1>(), 8);
        assert_eq!(core::mem::offset_of!(BootInfoV1, memory_map_ptr), 16);
        assert_eq!(core::mem::offset_of!(BootInfoV1, descriptor_bytes), 32);
        assert_eq!(core::mem::offset_of!(BootInfoV1, kernel_phys_start), 40);
        assert_eq!(core::mem::offset_of!(BootInfoV1, dtb_phys_addr), 56);
        assert_eq!(core::mem::offset_of!(BootInfoV1, flags), 64);
        let empty = BootInfoV1::empty();
        assert_eq!(empty.validate_header(), Ok(()));
        assert_eq!(empty.memory_map_ptr, 0);
        assert_eq!(empty.kernel_phys_start, 0);
        assert_eq!(empty.dtb_phys_addr, 0);
    }

    #[test]
    fn k503_pointer_numeric_gate_rejects_zero_unaligned_and_overflow() {
        assert_eq!(validate_boot_info_pointer(0x4000_1000), Ok(()));
        assert_eq!(
            validate_boot_info_pointer(0),
            Err(BootInfoPointerError::Missing)
        );
        assert_eq!(
            validate_boot_info_pointer(0x4000_1001),
            Err(BootInfoPointerError::Unaligned)
        );
        assert_eq!(
            validate_boot_info_pointer(u64::MAX - 7),
            Err(BootInfoPointerError::AddressOverflow)
        );
        assert_eq!(validate_boot_info_pointer(0x4000_1000), Ok(()));
    }

    #[test]
    fn k484_versioned_header_accepts_v1_and_rejects_incompatible_inputs() {
        let original = BootInfoV1::empty();
        assert_eq!(original.validate_header(), Ok(()));
        assert_eq!(core::mem::size_of::<BootInfoV1>(), 72);
        let mut wrong = original;
        wrong.magic[0] ^= 1;
        assert_eq!(wrong.validate_header(), Err(BootInfoHeaderError::Magic));
        wrong = original;
        wrong.version = 0;
        assert_eq!(wrong.validate_header(), Err(BootInfoHeaderError::Version));
        wrong.version = 2;
        assert_eq!(wrong.validate_header(), Err(BootInfoHeaderError::Version));
        wrong = original;
        wrong.struct_bytes = 16;
        assert_eq!(wrong.validate_header(), Err(BootInfoHeaderError::Size));
        wrong.struct_bytes = 80;
        assert_eq!(wrong.validate_header(), Err(BootInfoHeaderError::Size));
        assert_eq!(
            original.validate_header(),
            Ok(()),
            "rejections do not mutate source"
        );
    }

    #[test]
    fn k485_memory_map_metadata_rejects_malformed_spans_without_dereferencing_them() {
        let mut info = BootInfoV1::empty();
        assert_eq!(
            info.validate_memory_map_metadata(),
            Err(BootInfoMemoryMapError::Missing)
        );
        info.memory_map_ptr = 0x1000;
        info.memory_map_bytes = 80;
        info.descriptor_bytes = 40;
        info.descriptor_version = 1;
        assert_eq!(info.validate_memory_map_metadata(), Ok(()));
        let valid = info;
        info.memory_map_ptr = 0;
        assert_eq!(
            info.validate_memory_map_metadata(),
            Err(BootInfoMemoryMapError::Missing)
        );
        info = valid;
        info.memory_map_ptr = 0x1001;
        assert_eq!(
            info.validate_memory_map_metadata(),
            Err(BootInfoMemoryMapError::Length)
        );
        info = valid;
        info.memory_map_ptr = u64::MAX - 7;
        assert_eq!(
            info.validate_memory_map_metadata(),
            Err(BootInfoMemoryMapError::AddressOverflow)
        );
        info = valid;
        info.memory_map_bytes = 0;
        assert_eq!(
            info.validate_memory_map_metadata(),
            Err(BootInfoMemoryMapError::Missing)
        );
        info = valid;
        info.memory_map_bytes = 79;
        assert_eq!(
            info.validate_memory_map_metadata(),
            Err(BootInfoMemoryMapError::Length)
        );
        info = valid;
        info.descriptor_bytes = 32;
        assert_eq!(
            info.validate_memory_map_metadata(),
            Err(BootInfoMemoryMapError::DescriptorSize)
        );
        info.descriptor_bytes = 41;
        assert_eq!(
            info.validate_memory_map_metadata(),
            Err(BootInfoMemoryMapError::DescriptorSize)
        );
        info = valid;
        info.descriptor_version = 0;
        assert_eq!(
            info.validate_memory_map_metadata(),
            Err(BootInfoMemoryMapError::DescriptorVersion)
        );
        info.descriptor_version = 2;
        assert_eq!(
            info.validate_memory_map_metadata(),
            Err(BootInfoMemoryMapError::DescriptorVersion)
        );
    }

    #[test]
    fn k486_exit_retry_budget_stops_after_three_invalid_keys() {
        let mut budget = ExitBootServicesRetry::new();
        assert_eq!(budget.next_attempt(), Some(1));
        assert_eq!(budget.next_attempt(), Some(2));
        assert_eq!(budget.next_attempt(), Some(3));
        for _ in 0..1000 {
            assert_eq!(budget.next_attempt(), None);
        }
        let mut second = ExitBootServicesRetry::new();
        assert_eq!(second.next_attempt(), Some(1));
    }

    #[test]
    fn k487_kernel_file_rejects_empty_oversize_declared_and_actual_short_reads() {
        const CAP: u64 = 1_048_576;
        assert_eq!(
            validate_kernel_file_length(44_000, 44_000, None, CAP),
            Ok(())
        );
        assert_eq!(
            validate_kernel_file_length(44_000, 44_000, Some(44_000), CAP),
            Ok(())
        );
        assert_eq!(
            validate_kernel_file_length(0, 0, None, CAP),
            Err(KernelFileLengthError::Empty)
        );
        assert_eq!(
            validate_kernel_file_length(44_000, 0, None, CAP),
            Err(KernelFileLengthError::Empty)
        );
        assert_eq!(
            validate_kernel_file_length(CAP + 1, CAP + 1, None, CAP),
            Err(KernelFileLengthError::TooLarge)
        );
        assert_eq!(
            validate_kernel_file_length(44_000, CAP + 1, None, CAP),
            Err(KernelFileLengthError::TooLarge)
        );
        assert_eq!(
            validate_kernel_file_length(44_000, 43_999, None, CAP),
            Err(KernelFileLengthError::DeclaredMismatch)
        );
        assert_eq!(
            validate_kernel_file_length(44_000, 44_001, None, CAP),
            Err(KernelFileLengthError::DeclaredMismatch)
        );
        assert_eq!(
            validate_kernel_file_length(44_000, 44_000, Some(43_999), CAP),
            Err(KernelFileLengthError::ShortRead)
        );
        assert_eq!(
            validate_kernel_file_length(44_000, 44_000, Some(44_001), CAP),
            Err(KernelFileLengthError::OverRead)
        );
    }
}