yggdrasil/kernel/libk/src/module.rs

use core::mem::size_of;

use alloc::{string::String, vec, vec::Vec};
use bytemuck::{Pod, Zeroable};
use elf::{
    io_traits::{InputStream, SeekFrom},
    relocation::Rela,
    symbol::Symbol as ElfSymbol,
};
use libk_mm::PageBox;
use libk_util::{hash_table::DefaultHashTable, io::Read, sync::LockMethod, OneTimeInit};
use yggdrasil_abi::{
    error::Error,
    io::{FileMode, OpenOptions},
    path::Path,
};

use crate::{
    task::{
        binary::{
            self,
            elf::{ElfSegment, ElfSegmentType},
        },
        sync::Mutex,
    },
    vfs::{FileRef, IoContext},
};

#[derive(Clone, Copy, Debug, Zeroable, Pod)]
#[repr(C)]
pub struct Symbol {
    pub st_name: u32,
    pub st_info: u8,
    pub st_other: u8,
    pub st_shndx: u16,
    pub st_value: u64,
    pub st_size: u64,
}

#[derive(Clone, Debug)]
#[repr(C)]
pub struct ModuleInfo {
    pub name: &'static str,
    pub version: &'static str,
    pub init: fn() -> Result<(), Error>,
    pub unload: Option<fn() -> Result<(), Error>>,
}

#[derive(Debug)]
pub struct LoadedModule {
    pub image_data: PageBox<[u8]>,
    // Ref into the image_data PageBox
    pub info: &'static ModuleInfo,
}

pub fn load_kernel_symbol_table<P: AsRef<Path>>(
    ioctx: &mut IoContext,
    path: P,
) -> Result<(), Error> {
    let symbol_file = ioctx.open(None, path, OpenOptions::READ, FileMode::empty())?;
    let mut string_buffer = PageBox::new_slice(0, 4096)?;
    let mut map = DefaultHashTable::new();

    loop {
        let mut len = [0; 8];
        if symbol_file.read(&mut len)? != len.len() {
            break;
        }
        let len = usize::from_le_bytes(len);
        symbol_file.read_exact(&mut string_buffer[..len])?;
        let name = core::str::from_utf8(&string_buffer[..len]).unwrap();
        let mut value = [0; 8];
        symbol_file.read_exact(&mut value)?;
        let value = usize::from_le_bytes(value);
        map.insert(name.into(), value);
    }

    KERNEL_SYMBOL_TABLE.init(map);
    MODULE_SYMBOL_TABLE.init(Mutex::new(DefaultHashTable::new()));

    Ok(())
}

pub fn lookup_symbol(name: &str) -> Option<usize> {
    let symtab = KERNEL_SYMBOL_TABLE.get();
    if let Some(addr) = symtab.get(name).copied() {
        return Some(addr);
    }
    let symtab = MODULE_SYMBOL_TABLE.get();
    if let Some(addr) = symtab.lock().unwrap().get(name).copied() {
        return Some(addr);
    }

    None
}

pub fn load(file: FileRef) -> Result<LoadedModule, Error> {
    // TODO better dependency tracking than just throwing undefined references
    let (mut elf, mut file) = binary::elf::open_elf_direct(file)?;

    // Find out image size
    let (vaddr_min, vaddr_max) = binary::elf::elf_virtual_range(&elf);
    let mut image_data = PageBox::new_slice(0u8, vaddr_max - vaddr_min)?;

    // Load the segments
    for segment in elf
        .segments()
        .iter()
        .filter_map(ElfSegment::from_phdr)
        .filter(|seg| seg.ty == ElfSegmentType::Load)
    {
        let rel_offset = segment.vaddr as usize - vaddr_min;
        let segment_data = &mut image_data[rel_offset..rel_offset + segment.full_size as usize];

        if segment.data_size > 0 {
            file.seek(SeekFrom::Start(segment.offset)).unwrap();
            file.read_exact(&mut segment_data[..segment.data_size as usize])
                .unwrap();
        }
    }

    let mut info_struct_addr = None;

    // Extract dynamic symbols
    let (dynsym, dynstr) = elf.dynamic_symbol_table().unwrap().unwrap();
    let mut dynamic_symbols = vec![];
    for mut sym in dynsym {
        let name = dynstr.get(sym.st_name as _).unwrap();

        if sym.st_shndx == elf::abi::SHN_UNDEF && sym.st_bind() != elf::abi::STB_LOCAL {
            let Some(value) = lookup_symbol(name) else {
                log::warn!("Could not load module: undefined reference to {:?} (possibly compiled against a different libk?)", name);
                return Err(Error::InvalidArgument);
            };
            sym.st_value = value as u64;

            dynamic_symbols.push(sym);
        } else {
            // Adjust symbol value by image base
            sym.st_value += image_data.as_ptr().addr() as u64;

            if name == "MODULE" && sym.st_symtype() == elf::abi::STT_OBJECT {
                if sym.st_size != size_of::<ModuleInfo>() as u64 {
                    todo!();
                }
                info_struct_addr = Some(sym.st_value as usize);
            }

            dynamic_symbols.push(sym);
        }
    }

    // Apply relocations
    let rela_sections: Vec<_> = elf
        .section_headers()
        .iter()
        .filter(|c| c.sh_type == elf::abi::SHT_RELA)
        .copied()
        .collect();

    for rela in rela_sections {
        let rela = elf.section_data_as_relas(&rela).unwrap();
        for rela in rela {
            handle_relocation(&mut image_data, &rela, |idx| {
                dynamic_symbols
                    .get(idx as usize)
                    .ok_or(Error::InvalidArgument)
            })?;
        }
    }

    // TODO Add module symbols to export list: make sure to only export properly namespaced items
    // let module_syms = MODULE_SYMBOL_TABLE.get();
    // for (name, sym) in dynamic_symbols.iter() {
    //     if sym.st_shndx == elf::abi::SHN_UNDEF || sym.st_vis() != elf::abi::STV_DEFAULT {
    //         continue;
    //     }
    //     if !name.starts_with("_ZN") {
    //         continue;
    //     }
    //     log::info!("Export {:?} -> {:#x}", name, sym.st_value);
    //     module_syms
    //         .lock()
    //         .unwrap()
    //         .insert(name.clone(), sym.st_value as usize);
    // }

    let info = info_struct_addr
        .map(|addr| unsafe { core::mem::transmute(addr) })
        .expect("TODO return error if module is missing MODULE info struct");

    Ok(LoadedModule { image_data, info })
}

pub fn load_and_execute(file: FileRef) -> Result<(), Error> {
    let module = load(file)?;
    (module.info.init)()
}

fn handle_relocation<'a, F: FnOnce(u32) -> Result<&'a ElfSymbol, Error>>(
    image_data: &mut PageBox<[u8]>,
    rela: &Rela,
    resolve_symbol: F,
) -> Result<(), Error> {
    let value: i64 = match rela.r_type {
        // Symbol + addend
        #[cfg(target_arch = "x86_64")]
        elf::abi::R_X86_64_64 => resolve_symbol(rela.r_sym)?.st_value as i64 + rela.r_addend,
        elf::abi::R_AARCH64_ABS64 => resolve_symbol(rela.r_sym)?.st_value as i64 + rela.r_addend,

        // Image base + addend
        #[cfg(target_arch = "x86_64")]
        elf::abi::R_X86_64_RELATIVE => image_data.as_ptr().addr() as i64 + rela.r_addend,
        #[cfg(any(target_arch = "aarch64", rust_analyzer))]
        elf::abi::R_AARCH64_RELATIVE => image_data.as_ptr().addr() as i64 + rela.r_addend,

        // Symbol value
        #[cfg(target_arch = "x86_64")]
        elf::abi::R_X86_64_JUMP_SLOT | elf::abi::R_X86_64_GLOB_DAT => {
            resolve_symbol(rela.r_sym)?.st_value as i64
        }
        #[cfg(any(target_arch = "aarch64", rust_analyzer))]
        elf::abi::R_AARCH64_JUMP_SLOT | elf::abi::R_AARCH64_GLOB_DAT => {
            resolve_symbol(rela.r_sym)?.st_value as i64
        }
        _ => todo!("Unhandled relocation: {}", rela.r_type),
    };

    image_data[rela.r_offset as usize..rela.r_offset as usize + 8]
        .copy_from_slice(&value.to_ne_bytes());

    Ok(())
}

// TODO roll own hashmap
static KERNEL_SYMBOL_TABLE: OneTimeInit<DefaultHashTable<String, usize, 128>> = OneTimeInit::new();
static MODULE_SYMBOL_TABLE: OneTimeInit<Mutex<DefaultHashTable<String, usize>>> =
    OneTimeInit::new();