yggdrasil/userspace/lib/ygglibc/src/env.rs

#![allow(non_upper_case_globals, static_mut_refs)]

use core::{
    ffi::{c_char, CStr},
    ptr::{null_mut, NonNull},
    slice::memchr,
};

use alloc::{borrow::ToOwned, ffi::CString, vec::Vec};
use yggdrasil_rt::process::{thread_local, ProgramArgumentInner};
// use yggdrasil_rt::process::ProgramArgumentInner;

use crate::{
    allocator::{c_alloc, c_free}, error::EResult, headers::{
        errno,
        string::{mem::memcpy, str::strlen},
    }, thread::{self, tls}, util::PointerExt
};

#[no_mangle]
pub static mut environ: *mut *mut c_char = null_mut();
static mut shadow: Vec<*mut c_char> = Vec::new();

unsafe fn get_key(var: NonNull<c_char>) -> Option<&'static [u8]> {
    let cstr = CStr::from_ptr(var.as_ptr()).to_bytes();
    let eq_pos = memchr::memchr(b'=', cstr)?;
    Some(&cstr[..eq_pos])
}

unsafe fn entry_from_key_value(key: &[u8], value: &[u8]) -> EResult<NonNull<c_char>> {
    let entry = c_alloc(key.len() + value.len() + 2, 1, false)?.cast::<c_char>();
    memcpy(entry.as_ptr().cast(), key.as_ptr().cast(), key.len());
    memcpy(
        entry.add(key.len() + 1).as_ptr().cast(),
        value.as_ptr().cast(),
        value.len(),
    );
    entry.add(key.len()).write(b'=' as _);
    entry.add(key.len() + value.len() + 1).write(0);
    EResult::Ok(entry)
}

unsafe fn entry_from_raw(raw: &[u8]) -> EResult<NonNull<c_char>> {
    let entry = c_alloc(raw.len() + 1, 1, false)?.cast::<c_char>();
    memcpy(entry.as_ptr().cast(), raw.as_ptr().cast(), raw.len());
    entry.add(raw.len()).write(0);
    EResult::Ok(entry)
}

unsafe fn reclaim_env() -> EResult<()> {
    if environ == shadow.as_mut_ptr() {
        // environ already points to ygglibc-owned list
        return EResult::Ok(());
    }

    todo!()
}

unsafe fn push_env(str: NonNull<c_char>) -> EResult<()> {
    reclaim_env()?;

    if shadow.try_reserve(1).is_err() {
        return EResult::Err(errno::ENOMEM);
    }

    let entry = shadow.last_mut().unwrap();
    assert!(entry.is_null());
    *entry = str.as_ptr();

    // Won't fail
    shadow.push(null_mut());

    // Fixup the pointer, shadow might've got reallocated
    environ = shadow.as_mut_ptr();

    EResult::Ok(())
}

pub unsafe fn get_env(name: &[u8]) -> EResult<Option<NonNull<c_char>>> {
    if name.is_empty() || name.contains(&b'=') {
        return EResult::Err(errno::EINVAL);
    }

    let Some(mut iter) = NonNull::new(environ) else {
        // User set environ to NULL
        return EResult::Ok(None);
    };

    while let Some(entry) = NonNull::new(iter.read()) {
        let Some(key) = get_key(entry) else { continue };

        if key == name {
            // Safe: key != None means there is at least a '=' after whatever "key" points at
            let value = entry.add(key.len() + 1);
            return EResult::Ok(Some(value));
        }

        iter = iter.add(1);
    }

    EResult::Ok(None)
}

pub unsafe fn set_env(name: &[u8], value: NonNull<c_char>, overwrite: bool) -> EResult<()> {
    if name.is_empty() || name.contains(&b'=') {
        return EResult::Err(errno::EINVAL);
    }

    let value = CStr::from_ptr(value.as_ptr()).to_bytes();

    if let Some(mut iter) = NonNull::new(environ) {
        // Try to locate the key=value pair in the existing env
        while let Some(entry) = NonNull::new(iter.read()) {
            let Some(key) = get_key(entry) else { continue };

            if key == name {
                if !overwrite {
                    return EResult::Ok(());
                }

                // Safe: key != None means there is at least a '=' after whatever "key" points at
                let old_value = entry.add(key.len() + 1);

                if strlen(old_value.as_ptr()) >= value.len() {
                    // Old allocation can contain `value`
                    memcpy(
                        old_value.as_ptr().cast(),
                        value.as_ptr().cast(),
                        value.len(),
                    );
                    old_value.add(value.len()).write(0);
                } else {
                    // Need to allocate a new entry
                    let new_entry = entry_from_key_value(name, value)?;
                    c_free(entry.cast());
                    iter.write(new_entry.as_ptr());
                }

                return EResult::Ok(());
            }

            iter = iter.add(1);
        }
    }

    // Key not found or environ == NULL, need to push a new entry
    let entry = entry_from_key_value(name, value)?;
    push_env(entry)
}

pub unsafe fn remove_env(name: &[u8]) -> EResult<bool> {
    if name.is_empty() || name.contains(&b'=') {
        return EResult::Err(errno::EINVAL);
    }

    reclaim_env()?;

    for i in 0..shadow.len() {
        let Some(entry) = NonNull::new(shadow[i]) else {
            continue;
        };
        let Some(key) = get_key(entry) else { continue };

        if key == name {
            // Pop entry at [i]
            shadow.remove(i);
            environ = shadow.as_mut_ptr();
            c_free(entry.cast());
            return EResult::Ok(true);
        }
    }

    EResult::Ok(false)
}

pub unsafe fn put_env(str: NonNull<c_char>) -> EResult<()> {
    let bytes = CStr::from_ptr(str.as_ptr()).to_bytes();
    if bytes.is_empty() || !bytes.contains(&b'=') {
        return EResult::Err(errno::EINVAL);
    }
    let name = get_key(str).unwrap();

    if let Some(mut iter) = NonNull::new(environ) {
        while let Some(entry) = NonNull::new(iter.read()) {
            let Some(key) = get_key(entry) else { continue };

            if key == name {
                // Replace the entry
                iter.write(str.as_ptr());
                c_free(entry.cast());
                return EResult::Ok(());
            }

            iter = iter.add(1);
        }
    }

    // Push a new entry
    push_env(str)
}

unsafe fn setup_env<'a>(envs: impl Iterator<Item = &'a CStr>) {
    for env in envs {
        // Make a malloc()ed, NULL-terminated string
        let entry = entry_from_raw(env.to_bytes()).expect("Couldn't allocate env variable");
        shadow.push(entry.as_ptr());
    }
    shadow.push(null_mut());
    environ = shadow.as_mut_ptr();
}

pub fn handle_kernel_argument(arg: usize) -> Vec<CString> {
    let arg_ptr: *const ProgramArgumentInner = core::ptr::with_exposed_provenance(arg);
    let arg = unsafe { arg_ptr.ensure() };

    let mut args = Vec::new();

    for arg in arg.args() {
        let c_arg = arg.to_owned();
        args.push(c_arg);
    }

    unsafe { setup_env(arg.envs()) };

    // Setup TLS from argument
    thread::init_main_thread(arg);

    args
}