curiOS/kernel/src/x86_64/backtrace.rs

use core::ffi::{CStr, c_void};

use crate::{
    limine::LimineFile,
    memory::{VirtAddr, VirtAddrTranslationExt},
    serial_println,
    sync::LazyLock,
    x86_64::{Context, VirtAddrExt},
};

unsafe extern "C" {
    static __kernel_start: c_void;
    static __kernel_end: c_void;
}

pub struct Frame {
    pub rip: u64,
    pub rbp: u64,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Elf64SymBind {
    Local = 0,
    Global = 1,
    Weak = 2,
    Loos = 10,
    Hios = 12,
    Loproc = 13,
    Hiproc = 15,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Elf64SymType {
    Notype = 0,
    Object = 1,
    Func = 2,
    Section = 3,
    File = 4,
    Common = 5,
    Tls = 6,
    Loos = 10,
    Hios = 12,
    Loproc = 13,
    Hiproc = 15,
}

#[repr(C)]
#[derive(Debug)]
struct Elf64Sym {
    name: u32,
    info: u8,
    other: u8,
    shndx: u16,
    value: u64,
    size: u64,
}

unsafe impl plain::Plain for Elf64Sym {}

#[repr(u32)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum SectionHeaderType {
    Null = 0,
    Progbits = 1,
    Symtab = 2,
    Strtab = 3,
    Rela = 4,
    Hash = 5,
    Dynamic = 6,
    Note = 7,
    Nobits = 8,
    Rel = 9,
    Shlib = 10,
    Dynsym = 11,
    InitArray = 14,
    FiniArray = 15,
    PreinitArray = 16,
    Group = 17,
    SymtabShndx = 18,
    Num = 19,
    GnuAttributes = 0x6ffffff5,
    GnuHash = 0x6ffffff6,
    GnuLiblist = 0x6ffffff7,
    Checksum = 0x6ffffff8,
    Losunwind = 0x6ffffffa,
    Amd64Unwind = 0x7000001,
    OsSpecific(u32),
    ProcSpecific(u32),
    UserSpecific(u32),
    Unknown = 0xffffffff,
}

impl SectionHeaderType {
    fn from_u32(value: u32) -> Self {
        match value {
            0 => Self::Null,
            1 => Self::Progbits,
            2 => Self::Symtab,
            3 => Self::Strtab,
            4 => Self::Rela,
            5 => Self::Hash,
            6 => Self::Dynamic,
            7 => Self::Note,
            8 => Self::Nobits,
            9 => Self::Rel,
            10 => Self::Shlib,
            11 => Self::Dynsym,
            14 => Self::InitArray,
            15 => Self::FiniArray,
            16 => Self::PreinitArray,
            17 => Self::Group,
            18 => Self::SymtabShndx,
            19 => Self::Num,
            0x6ffffff5 => Self::GnuAttributes,
            0x6ffffff6 => Self::GnuHash,
            0x6ffffff7 => Self::GnuLiblist,
            0x6ffffff8 => Self::Checksum,
            0x6ffffffa => Self::Losunwind,
            0x7000001 => Self::Amd64Unwind,
            0x60000000..=0x6fffffff => Self::OsSpecific(value),
            0x70000000..=0x7fffffff => Self::ProcSpecific(value),
            0x80000000..=0x8fffffff => Self::UserSpecific(value),
            _ => Self::Unknown,
        }
    }

    fn into_u32(self) -> u32 {
        match self {
            Self::Null => 0,
            Self::Progbits => 1,
            Self::Symtab => 2,
            Self::Strtab => 3,
            Self::Rela => 4,
            Self::Hash => 5,
            Self::Dynamic => 6,
            Self::Note => 7,
            Self::Nobits => 8,
            Self::Rel => 9,
            Self::Shlib => 10,
            Self::Dynsym => 11,
            Self::InitArray => 14,
            Self::FiniArray => 15,
            Self::PreinitArray => 16,
            Self::Group => 17,
            Self::SymtabShndx => 18,
            Self::Num => 19,
            Self::GnuAttributes => 0x6ffffff5,
            Self::GnuHash => 0x6ffffff6,
            Self::GnuLiblist => 0x6ffffff7,
            Self::Checksum => 0x6ffffff8,
            Self::Losunwind => 0x6ffffffa,
            Self::Amd64Unwind => 0x7000001,
            Self::OsSpecific(value) | Self::ProcSpecific(value) | Self::UserSpecific(value) => {
                value
            }
            Self::Unknown => 0xffffffff,
        }
    }
}

#[repr(C)]
#[derive(Debug)]
struct SectionHeader {
    name: u32,
    sh_type: u32,
    flags: u64,
    addr: u64,
    offset: u64,
    size: u64,
    link: u32,
    info: u32,
    addralign: u64,
    entsize: u64,
}

unsafe impl plain::Plain for SectionHeader {}

impl SectionHeader {
    fn bytes<'a>(&self, bytes: &'a [u8]) -> Option<&'a [u8]> {
        let offset = self.offset as usize;
        let size = self.size as usize;

        if offset + size > bytes.len() {
            return None;
        }

        Some(&bytes[offset..offset + size])
    }
}

#[repr(C)]
#[derive(Debug)]
struct Elf64Ehdr {
    ident: [u8; 16],
    kind: u16,
    machine: u16,
    version: u32,
    entry: u64,
    phoff: u64,
    shoff: u64,
    flags: u32,
    ehsize: u16,
    phentsize: u16,
    phnum: u16,
    shentsize: u16,
    shnum: u16,
    shstrndx: u16,
}

unsafe impl plain::Plain for Elf64Ehdr {}

struct SectionHeaders<'a> {
    headers: &'a [SectionHeader],
}

impl<'a> SectionHeaders<'a> {
    fn symtab_hdr(&self) -> Option<&'a SectionHeader> {
        self.headers
            .iter()
            .find(|header| SectionHeaderType::from_u32(header.sh_type) == SectionHeaderType::Symtab)
    }
}

#[derive(Default)]
struct Strtab<'a> {
    bytes: &'a [u8],
}

impl<'a> Strtab<'a> {
    fn len(&self) -> usize {
        self.bytes.len()
    }
    fn get_str(&self, offset: u32) -> Option<&'a str> {
        let offset = offset as usize;

        if offset >= self.bytes.len() {
            return None;
        }

        let bytes = &self.bytes[offset..];

        let nul_pos = bytes.iter().position(|&b| b == 0)?;

        core::str::from_utf8(&bytes[..nul_pos]).ok()
    }
}

impl Elf64Ehdr {
    fn section_headers<'a>(&self, bytes: &'a [u8]) -> SectionHeaders<'a> {
        let shoff = self.shoff as usize;
        // let shentsize = self.shentsize as usize;
        let shnum = self.shnum as usize;

        unsafe {
            SectionHeaders {
                headers: core::slice::from_raw_parts(
                    bytes.as_ptr().byte_add(shoff).cast::<SectionHeader>(),
                    shnum,
                ),
            }
        }
    }

    fn shstrtab<'a>(&self, bytes: &'a [u8]) -> Option<Strtab<'a>> {
        let shstrndx = self.shstrndx as usize;
        let section_headers = self.section_headers(bytes);

        if shstrndx >= section_headers.headers.len() {
            return None;
        }

        let shstrtab_header = &section_headers.headers[shstrndx];
        serial_println!("shstrtab_header: {:#?}", shstrtab_header);
        let offset = shstrtab_header.offset as usize;
        let size = shstrtab_header.size as usize;

        if offset + size > bytes.len() {
            return None;
        }

        Some(Strtab {
            bytes: &bytes[offset..offset + size],
        })
    }

    fn find_shdr_by_name<'a>(&self, bytes: &'a [u8], name: &str) -> Option<&'a SectionHeader> {
        let shstrtab = self.shstrtab(bytes)?;
        let section_headers = self.section_headers(bytes);

        section_headers
            .headers
            .iter()
            .find(|header| shstrtab.get_str(header.name) == Some(name))
    }

    fn strtab<'a>(&self, bytes: &'a [u8]) -> Option<Strtab<'a>> {
        let strtab = self.find_shdr_by_name(bytes, ".strtab")?;

        let offset = strtab.offset as usize;
        let size = strtab.size as usize;

        if offset + size > bytes.len() {
            return None;
        }

        Some(Strtab {
            bytes: &bytes[offset..offset + size],
        })
    }

    fn symtab<'a>(&self, bytes: &'a [u8]) -> Option<&'a [Elf64Sym]> {
        let section_headers = self.section_headers(bytes);
        let symtab_header = section_headers.symtab_hdr()?;

        let offset = symtab_header.offset as usize;
        let size = symtab_header.size as usize;

        if offset + size > bytes.len() {
            return None;
        }

        let num_symbols = size / core::mem::size_of::<Elf64Sym>();

        unsafe {
            Some(core::slice::from_raw_parts(
                bytes.as_ptr().byte_add(offset).cast::<Elf64Sym>(),
                num_symbols,
            ))
        }
    }
}

impl Elf64Sym {
    fn symbols() -> &'static [Elf64Sym] {
        static SYMBOLS: LazyLock<&'static [Elf64Sym]> = LazyLock::new(|| {
            if let Some(bytes) = crate::limine::EXECUTABLE_FILE_REQUEST
                .file()
                .map(LimineFile::bytes)
            {
                if let Ok(ehdr) = plain::from_bytes::<Elf64Ehdr>(bytes) {
                    if let Some(symtab) = ehdr.symtab(bytes) {
                        return symtab;
                    }
                }
            }
            &[]
        });

        SYMBOLS.as_ref()
    }

    fn find(addr: u64) -> Option<&'static Elf64Sym> {
        Self::symbols()
            .iter()
            .find(|&sym| sym.range().contains(&addr))
            .map(|v| v as _)
    }

    fn find_type(addr: u64, kind: Elf64SymType) -> Option<&'static Elf64Sym> {
        Self::symbols()
            .iter()
            .find(|&sym| sym.kind() == kind && sym.range().contains(&addr))
            .map(|v| v as _)
    }

    fn bind(&self) -> Elf64SymBind {
        match self.info >> 4 {
            0 => Elf64SymBind::Local,
            1 => Elf64SymBind::Global,
            2 => Elf64SymBind::Weak,
            10 => Elf64SymBind::Loos,
            12 => Elf64SymBind::Hios,
            13 => Elf64SymBind::Loproc,
            15 => Elf64SymBind::Hiproc,
            _ => panic!("Invalid bind value"),
        }
    }
    fn kind(&self) -> Elf64SymType {
        match self.info & 0xf {
            0 => Elf64SymType::Notype,
            1 => Elf64SymType::Object,
            2 => Elf64SymType::Func,
            3 => Elf64SymType::Section,
            4 => Elf64SymType::File,
            5 => Elf64SymType::Common,
            6 => Elf64SymType::Tls,
            10 => Elf64SymType::Loos,
            12 => Elf64SymType::Hios,
            13 => Elf64SymType::Loproc,
            15 => Elf64SymType::Hiproc,
            _ => panic!("Invalid type value"),
        }
    }

    fn range(&self) -> core::ops::Range<u64> {
        self.value..(self.value.saturating_add(self.size))
    }

    fn name(&self) -> Option<&str> {
        static STRTAB: LazyLock<Strtab<'static>> = LazyLock::new(|| {
            if let Some(bytes) = crate::limine::EXECUTABLE_FILE_REQUEST
                .file()
                .map(LimineFile::bytes)
                && let Ok(ehdr) = plain::from_bytes::<Elf64Ehdr>(bytes)
            {
                serial_println!("ELF64 Header: {:#?}", ehdr);
                return ehdr.strtab(bytes).unwrap_or_default();
            }

            Strtab::default()
        });

        let len = STRTAB.len();

        if self.name == 0 || self.name as usize >= len {
            return None;
        }

        STRTAB.get_str(self.name)
    }
}

pub fn backtrace(ctx: &Context) -> impl Iterator<Item = Frame> {
    Backtrace {
        frame: Some(Frame {
            rip: ctx.rip,
            rbp: ctx.registers.rbp,
        }),
    }
}

pub fn debug_backtrace(ctx: &Context) {
    serial_println!("{ctx}");
    for frame in backtrace(ctx) {
        if let Some(sym) = Elf64Sym::find(frame.rip) {
            if sym.kind() == Elf64SymType::Func
                && sym.shndx != 0
                && let Some(name) = sym.name()
            {
                serial_println!(
                    "0x{:016x} - {} + 0x{:x}",
                    frame.rip,
                    rustc_demangle::demangle(name),
                    frame.rip - sym.value
                );
            } else {
                serial_println!("{sym:#?}");
                serial_println!(
                    "0x{:016x} - <unknown> + 0x{:x}",
                    frame.rip,
                    frame.rip - sym.value
                );
            }
        } else {
            serial_println!("0x{:016x} - <unknown>", frame.rip);
        }
    }
}

struct Backtrace {
    frame: Option<Frame>,
}

impl Iterator for Backtrace {
    type Item = Frame;

    fn next(&mut self) -> Option<Self::Item> {
        let frame = self.frame.take()?;

        let next_frame = {
            let rsp = VirtAddr::from(frame.rbp);
            if rsp.is_canonical() && rsp.is_mapped() {
                // The stack frame layout is as follows:
                // | Our Function    |
                // +-----------------+
                // | Previous RBP    | <- RBP points here
                // +-----------------+
                // | Return Address  | <- RBP + 8 points here
                // +-----------------+
                // | Parent Function |
                //
                let rip = unsafe { *(rsp.as_ptr::<u64>().add(1)) };
                let rbp = unsafe { *(rsp.as_ptr::<u64>()) };

                Some(Frame { rip, rbp })
            } else {
                None
            }
        };

        self.frame = next_frame;

        Some(frame)
    }
}