yggdrasil/kernel/libk/src/task/binary/mod.rs

use alloc::{borrow::ToOwned, string::String, sync::Arc, vec, vec::Vec};
use bytemuck::Pod;
use kernel_arch::task::{TaskContext, UserContextInfo};
use libk_mm::{
    pointer::PhysicalRefMut,
    process::{ProcessAddressSpace, VirtualRangeBacking},
    table::MapAttributes,
};
use libk_util::io::{Read, Seek};
use yggdrasil_abi::{
    error::Error,
    io::SeekFrom,
    path::Path,
    process::{auxv, AuxValue, ProcessGroupId},
};

use crate::{
    task::{
        mem::ForeignPointer,
        process::{Process, ProcessCreateInfo, ProcessImage},
        thread::Thread,
        TaskContextImpl,
    },
    vfs::IoContext,
};

use self::elf::ElfKind;

pub mod elf;

pub type LoadedProcess = (Arc<Process>, Arc<Thread>);

#[derive(Debug)]
pub struct LoadOptions<'e, P: AsRef<Path>> {
    pub parent: Option<Arc<Process>>,
    pub group_id: ProcessGroupId,
    pub path: P,
    pub args: &'e [&'e str],
    pub envs: &'e [&'e str],
    pub single_step: bool,
    pub disable_aslr: bool,
}

struct ArgPlacer<'a> {
    buffer: &'a mut [u8],
    position: usize,
}

impl<'a> ArgPlacer<'a> {
    fn new(buffer: &'a mut [u8]) -> Self {
        Self {
            buffer,
            position: 0,
        }
    }

    fn align(&mut self, align: usize) -> Result<(), Error> {
        debug_assert!(align.is_power_of_two());
        let aligned = (self.position + align - 1) & !(align - 1);
        if aligned > self.buffer.len() {
            return Err(Error::InvalidArgument);
        }
        self.position = aligned;
        Ok(())
    }

    fn put_str(&mut self, s: &str) -> Result<usize, Error> {
        if self.position + s.len() >= self.buffer.len() {
            return Err(Error::InvalidArgument);
        }
        let off = self.position;
        self.buffer[off..off + s.len()].copy_from_slice(s.as_bytes());
        self.buffer[off + s.len()] = 0;
        self.position += s.len() + 1;
        Ok(off)
    }

    fn put<T: Pod>(&mut self, v: &T) -> Result<usize, Error> {
        if self.position + size_of::<T>() > self.buffer.len() {
            return Err(Error::InvalidArgument);
        }
        let off = self.position;
        self.buffer[off..off + size_of::<T>()].copy_from_slice(bytemuck::bytes_of(v));
        self.position += size_of::<T>();
        Ok(off)
    }

    fn put_aux_array(&mut self, s: &[AuxValue]) -> Result<usize, Error> {
        self.align(size_of::<u64>())?;
        let off = self.position;
        for item in s {
            self.put(&item.tag)?;
            self.put(&item.val)?;
        }
        self.put(&auxv::NULL)?;
        self.put(&0u64)?;
        Ok(off)
    }

    fn put_ptr_array(&mut self, s: &[usize]) -> Result<usize, Error> {
        self.align(size_of::<usize>())?;
        let off = self.position;
        for item in s {
            self.put(item)?;
        }
        self.put(&0usize)?;
        Ok(off)
    }
}

// args, envs are passed as Vec to ensure kernel ownership
#[allow(clippy::ptr_arg)]
fn setup_program_env(
    space: &ProcessAddressSpace,
    virt: usize,
    real_program: Option<&str>,
    args: &Vec<String>,
    envs: &Vec<String>,
    aux: &[AuxValue],
) -> Result<usize, Error> {
    // TODO growing buffer
    let phys_page = space.map_single_anon(
        virt,
        MapAttributes::USER_READ | MapAttributes::USER_WRITE | MapAttributes::NON_GLOBAL,
    )?;
    let mut buffer = unsafe { PhysicalRefMut::map_slice(phys_page, 4096) };
    let mut placer = ArgPlacer::new(&mut buffer);
    let mut arg_ptrs = vec![];
    let mut env_ptrs = vec![];

    for arg in args.iter() {
        let ptr = placer.put_str(arg)? + virt;
        arg_ptrs.push(ptr);
    }
    for env in envs.iter() {
        let ptr = placer.put_str(env)? + virt;
        env_ptrs.push(ptr);
    }

    let argv = placer.put_ptr_array(&arg_ptrs)? + virt;
    let envp = placer.put_ptr_array(&env_ptrs)? + virt;
    let auxv = placer.put_aux_array(aux)? + virt;
    let real_program = match real_program {
        Some(path) => placer.put_str(path)? + virt,
        _ => 0,
    };

    placer.align(size_of::<usize>())?;

    // Put ProgramArgumentInner struct
    let arg_address = placer.position + virt;

    placer.put(&argv)?;
    placer.put(&envp)?;
    placer.put(&auxv)?;
    placer.put(&real_program)?;

    Ok(arg_address)
}

fn setup_context<P>(
    options: &LoadOptions<P>,
    space: &ProcessAddressSpace,
    image: &ProcessImage,
    image_path: &str,
    args: &Vec<String>,
    envs: &Vec<String>,
) -> Result<TaskContextImpl, Error>
where
    P: AsRef<Path>,
{
    const USER_STACK_PAGES: usize = 32;

    // TODO dynamically map the stack
    let virt_stack_base = 0x3000000;
    // 0x1000 of guard page
    let virt_args_base = virt_stack_base + (USER_STACK_PAGES + 1) * 0x1000;

    space.map(
        virt_stack_base,
        USER_STACK_PAGES * 0x1000,
        VirtualRangeBacking::anonymous(),
        MapAttributes::USER_WRITE | MapAttributes::USER_READ | MapAttributes::NON_GLOBAL,
    )?;

    log::debug!(
        "stack: {:#x}..{:#x}",
        virt_stack_base,
        virt_stack_base + USER_STACK_PAGES * 0x1000
    );

    let mut auxv = vec![];
    if let Some(tls_image) = image.tls_image.as_ref() {
        if tls_image.master_copy_base != 0 {
            auxv.push(AuxValue {
                tag: auxv::TLS_MASTER_COPY,
                val: tls_image.master_copy_base as _,
            });
            auxv.push(AuxValue {
                tag: auxv::TLS_DATA_SIZE,
                val: tls_image.data_size as _,
            });
        } else {
            assert_eq!(tls_image.data_size, 0);
        }
        auxv.push(AuxValue {
            tag: auxv::TLS_ALIGN,
            val: tls_image.align as _,
        });
        auxv.push(AuxValue {
            tag: auxv::TLS_FULL_SIZE,
            val: tls_image.full_size as _,
        });
    }

    let argument = setup_program_env(space, virt_args_base, Some(image_path), args, envs, &auxv)?;

    let user_sp =
        virt_stack_base + USER_STACK_PAGES * 0x1000 - TaskContextImpl::USER_STACK_EXTRA_ALIGN;

    // Fill with some sentinel value to detect stack underflows
    let ptr = user_sp as *mut usize;

    #[allow(clippy::reversed_empty_ranges)]
    for i in 0..TaskContextImpl::USER_STACK_EXTRA_ALIGN / 8 {
        unsafe {
            ptr.add(i).write_foreign_volatile(space, 0xDEADC0DE);
        }
    }

    let ptr = unsafe { Thread::setup_stack_header(space, ptr, argument)? };

    // let tls_address = elf::clone_tls(space, image)?;

    log::debug!(
        "argument = {:#x}, user_sp = {:#x}, stack: {:#x}..{:#x}",
        argument,
        user_sp,
        virt_stack_base,
        virt_stack_base + USER_STACK_PAGES * 0x1000 - TaskContextImpl::USER_STACK_EXTRA_ALIGN
    );
    let (address_space, asid) = space.as_address_with_asid();
    TaskContext::user(UserContextInfo {
        entry: image.entry,
        argument,
        stack_pointer: ptr.addr(),
        thread_pointer: 0,
        address_space,
        asid,
        single_step: options.single_step,
    })
}

fn setup_binary<S, P>(
    options: &LoadOptions<P>,
    name: S,
    space: ProcessAddressSpace,
    image: ProcessImage,
    image_path: &str,
    args: &Vec<String>,
    envs: &Vec<String>,
) -> Result<LoadedProcess, Error>
where
    S: Into<String>,
    P: AsRef<Path>,
{
    let context = setup_context(options, &space, &image, image_path, args, envs)?;
    let info = ProcessCreateInfo {
        name,
        context,
        group_id: options.group_id,
        space: Arc::new(space),
        image: Some(image),
    };
    let (process, main) = Process::spawn(options.parent.as_ref(), info);
    Ok((process, main))
}

fn xxx_load_program<P: AsRef<Path>>(
    space: &ProcessAddressSpace,
    ioctx: &mut IoContext,
    path: P,
    args: Vec<String>,
    envs: Vec<String>,
    aslr: bool,
) -> Result<(ProcessImage, Vec<String>, Vec<String>), Error> {
    let mut head = [0; 256];
    let path = path.as_ref();
    let file = ioctx.open_executable(path)?;

    file.seek(SeekFrom::Start(0))?;
    let count = file.read(&mut head)?;
    let head = &head[..count];

    // Recognize #!
    if let Some(shebang) = head.strip_prefix(b"#!")
        && let Some((shebang, _)) = shebang.split_once(|&ch| ch == b'\n')
    {
        let shebang = core::str::from_utf8(shebang).map_err(|_| Error::InvalidFile)?;
        let mut shebang_args = shebang.split(' ').map(|s| s.to_owned()).collect::<Vec<_>>();
        if shebang_args.is_empty() || shebang_args.len() >= 8 {
            return Err(Error::UnrecognizedExecutable);
        }
        shebang_args.extend_from_slice(&args);

        return xxx_load_program(
            space,
            ioctx,
            shebang_args[0].clone(),
            shebang_args,
            envs,
            aslr,
        );
    }

    file.seek(SeekFrom::Start(0))?;

    log::debug!("open_elf_file({:?})", path);
    if head.starts_with(b"\x7FELF") {
        match elf::open_elf_file(file)? {
            ElfKind::Static(elf, file) => {
                let image = elf::load_elf_from_file(space, elf, file, aslr)?;

                Ok((image, args, envs))
            }
            ElfKind::Dynamic(interp) => {
                let interp = interp.unwrap_or_else(|| "/libexec/dyn-loader".into());
                let mut interp_args = Vec::new();
                interp_args.push(interp.as_str().into());
                interp_args.extend_from_slice(&args);
                xxx_load_program(space, ioctx, interp, interp_args, envs, aslr)
            }
        }
    } else {
        Err(Error::UnrecognizedExecutable)
    }
}

/// Loads a program from given `path`
pub fn load<P: AsRef<Path>>(
    ioctx: &mut IoContext,
    options: &LoadOptions<P>,
) -> Result<LoadedProcess, Error> {
    let path = options.path.as_ref();
    let args = options.args.iter().map(|&s| s.to_owned()).collect();
    let envs = options.envs.iter().map(|&s| s.to_owned()).collect();

    let space = ProcessAddressSpace::new()?;
    let (image, args, envs) =
        xxx_load_program(&space, ioctx, path, args, envs, !options.disable_aslr)?;
    let real_path = path.display();
    let name = match real_path.rsplit_once('/') {
        Some((_, name)) => name,
        None => real_path,
    };
    setup_binary(options, name, space, image, real_path, &args, &envs)
}

pub fn load_into<P: AsRef<Path>>(
    options: &LoadOptions<P>,
    args: Vec<String>,
    envs: Vec<String>,
) -> Result<(TaskContextImpl, ProcessImage), Error> {
    let process = options.parent.as_ref().ok_or(Error::InvalidOperation)?;
    let mut io = process.io.lock();
    // Have to make the Path owned, going to drop the address space from which it came
    let path = options.path.as_ref().to_owned();
    let space = process.space();
    space.clear()?;
    let (image, args, envs) = xxx_load_program(
        &space,
        io.ioctx_mut(),
        &path,
        args,
        envs,
        !options.disable_aslr,
    )?;
    let context = setup_context(options, &space, &image, path.display(), &args, &envs)?;

    Ok((context, image))
}