yggdrasil/tools/gentables/src/x86_64.rs

use core::fmt;
use std::{
    io::{Read, Seek},
    mem::offset_of,
};

use bitflags::bitflags;
use bytemuck::Zeroable;
use elf::{
    abi::{PF_W, PF_X, PT_LOAD},
    endian::AnyEndian,
    ElfStream,
};
use memtables::{FixedTables, KERNEL_L3_COUNT};

use crate::{GenData, GenError};

bitflags! {
    #[derive(Clone, Copy)]
    struct PageFlags: u64 {
        const PRESENT = 1 << 0;
        const WRITABLE = 1 << 1;
        const NX = 1 << 63;
    }
}

pub struct X8664Builder<F: Seek + Read> {
    elf: ElfStream<AnyEndian, F>,
    data: GenData,
    tables: FixedTables,

    l0i: usize,
    l1i: usize,
    start_l2i: usize,
    end_l2i: usize,
}

impl PageFlags {
    fn from_elf(flags: u32) -> Self {
        let mut out = Self::empty();
        if flags & PF_W != 0 {
            out |= Self::WRITABLE;
        }
        if flags & PF_X == 0 {
            out |= Self::NX;
        }
        out
    }
}

impl fmt::Display for PageFlags {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "r{}{}",
            if self.contains(Self::WRITABLE) {
                'w'
            } else {
                '-'
            },
            if self.contains(Self::NX) { '-' } else { 'x' }
        )
    }
}

impl<F: Seek + Read> X8664Builder<F> {
    pub fn new(elf: ElfStream<AnyEndian, F>, data: GenData) -> Result<Self, GenError> {
        let l2_aligned_start = data.kernel_start & !0x1FFFFF;
        let l2_aligned_end = (data.kernel_end + 0x1FFFFF) & !0x1FFFFF;

        if l2_aligned_end <= l2_aligned_start {
            todo!();
        }

        if (l2_aligned_end - l2_aligned_start) as usize >= KERNEL_L3_COUNT * 0x200000 {
            return Err(GenError::KernelTooLarge(
                l2_aligned_start..l2_aligned_end,
                l2_aligned_end - l2_aligned_start,
                (KERNEL_L3_COUNT * 0x20000) as u64,
            ));
        }

        let l0i = (data.kernel_start >> 39) as usize & 0x1FF;
        let l1i = (data.kernel_start >> 30) as usize & 0x1FF;
        let start_l2i = (l2_aligned_start >> 21) as usize & 0x1FF;
        let end_l2i = (l2_aligned_end >> 21) as usize & 0x1FF;

        Ok(Self {
            elf,
            data,
            tables: FixedTables::zeroed(),

            l0i,
            l1i,
            start_l2i,
            end_l2i,
        })
    }

    pub fn build(mut self) -> Result<(FixedTables, u64), GenError> {
        // L0 -> L1
        let l1_physical_address =
            self.data.table_physical_address + offset_of!(FixedTables, kernel_l1) as u64;
        self.tables.l0.data[self.l0i] =
            l1_physical_address | (PageFlags::PRESENT | PageFlags::WRITABLE).bits();

        // L1 -> L2
        let l2_physical_address =
            self.data.table_physical_address + offset_of!(FixedTables, kernel_l2) as u64;
        self.tables.kernel_l1.data[self.l1i] =
            l2_physical_address | (PageFlags::PRESENT | PageFlags::WRITABLE).bits();

        // L2 -> L3s
        for i in 0..KERNEL_L3_COUNT {
            let l3_physical_address = self.data.table_physical_address
                + (offset_of!(FixedTables, kernel_l3s) + 0x1000 * i) as u64;

            self.tables.kernel_l2.data[i + self.start_l2i] =
                l3_physical_address | (PageFlags::PRESENT | PageFlags::WRITABLE).bits();
        }

        for (i, segment) in self.elf.segments().into_iter().enumerate() {
            if segment.p_type != PT_LOAD
                || segment.p_vaddr != segment.p_paddr + self.data.kernel_virt_offset
            {
                continue;
            }

            let aligned_virt_start = segment.p_vaddr & !0xFFF;
            let aligned_virt_end = (segment.p_vaddr + segment.p_memsz + 0xFFF) & !0xFFF;
            let aligned_phys_start = segment.p_paddr & !0xFFF;
            let count = (aligned_virt_end - aligned_virt_start) / 0x1000;

            let flags = PageFlags::from_elf(segment.p_flags);

            println!(
                "{}: {:#x?} -> {:#x} {}",
                i,
                aligned_virt_start..aligned_virt_end,
                aligned_phys_start,
                flags
            );

            Self::map_segment(
                self.start_l2i,
                &mut self.tables,
                aligned_virt_start,
                aligned_phys_start,
                count as usize,
                flags,
            )?;
        }

        Ok((self.tables, self.data.table_offset))
    }

    fn map_segment(
        l2i_offset: usize,
        tables: &mut FixedTables,
        vaddr: u64,
        paddr: u64,
        count: usize,
        flags: PageFlags,
    ) -> Result<(), GenError> {
        for index in 0..count {
            let address = vaddr + index as u64 * 0x1000;
            let page = paddr + index as u64 * 0x1000;

            let entry = page | (PageFlags::PRESENT | flags).bits();

            let l2i = (address >> 21) as usize & 0x1FF - l2i_offset;
            let l3i = (address >> 12) as usize & 0x1FF;

            let l3 = &mut tables.kernel_l3s[l2i];

            if l3.data[l3i] != 0 {
                if l3.data[l3i] != entry {
                    todo!();
                } else {
                    continue;
                }
            }

            l3.data[l3i] = entry;
        }

        Ok(())
    }
}