yggdrasil/userspace/dyn-loader/src/main.rs

#![feature(yggdrasil_os)]
use std::{
    collections::HashMap,
    env,
    path::{Path, PathBuf},
    process::ExitCode,
};

use error::Error;
use state::State;
use yggdrasil_rt::process::ProgramArgumentInner;

use crate::object::Object;

mod builtins;
mod error;
mod object;
mod relocation;
mod state;

struct Config {
    library_path: Vec<PathBuf>,
    trace_libraries: bool,
}

impl Config {
    pub fn from_env() -> Result<Self, Error> {
        let library_path = if let Ok(paths) = env::var("LD_LIBRARY_PATH") {
            paths.split(':').map(PathBuf::from).collect()
        } else {
            vec!["/lib".into()]
        };
        let trace_libraries = env::var("LD_TRACE_LOADED_OBJECTS")
            .map(|v| v == "1")
            .unwrap_or(false);

        Ok(Self {
            library_path,
            trace_libraries,
        })
    }
}

static mut STATE: Option<State> = None;

fn resolve_library(config: &Config, name: &str) -> Result<PathBuf, Error> {
    // TODO temporary hack for hash-tagged libstd
    if name.starts_with("libstd-") && name.ends_with(".so") {
        return resolve_library(config, "libstd.so");
    }

    for lib_dir in config.library_path.iter() {
        let path = lib_dir.join(name);

        if path.exists() {
            return Ok(path);
        }
    }

    Err(Error::LibraryNotFound(name.into()))
}

fn run<P: AsRef<Path>>(path: P, args: &[String]) -> Result<(), Error> {
    let mut state = State::default();
    let mut libs = HashMap::new();

    let config = Config::from_env()?;

    state.insert_linker_builtins();

    // Open and load the main object
    let mut main_object = Object::open(path)?;
    main_object.load(&mut state)?;

    main_object.collect_dependencies(&mut libs, &|name| resolve_library(&config, name))?;

    if config.trace_libraries {
        println!("Main object: {}", main_object.path.display());
        for (item, _) in libs.iter() {
            println!("* {}", item.display());
        }
        return Ok(());
    }

    // Load the libraries first
    for (_, lib) in libs.iter_mut() {
        lib.load(&mut state)?;
    }

    // Relocate the libraries
    for (_, lib) in libs.iter_mut() {
        lib.relocate(&mut state)?;
    }
    // Then relocate the main object
    main_object.relocate(&mut state)?;

    debug_trace!("Load finished");

    if !state.undefined_references.is_empty() {
        for item in state.undefined_references.iter() {
            eprintln!("Undefined reference to {:?}", item);
        }
        return Err(Error::UndefinedReference);
    }

    let args = args.iter().map(|s| s.as_str()).collect::<Vec<_>>();
    // TODO pass environment to the program
    let envs = vec![];

    let arg = Box::new(ProgramArgumentInner {
        args: &args,
        env: &envs,
    });

    let entry = main_object.entry()?.ok_or(Error::NoEntryPoint)?;

    // Store object as linker's global state and enter the program
    unsafe {
        STATE = Some(state);

        // FIXME TLS relocation and allocation is really broken for aarch64, so this hack ensures
        //       code doesn't try to add the TLS provided by the kernel to this linker to whatever
        //       is returned from __dl_tlsdesc_static.
        //
        //       This breaks TLS for the loader itself, but at least allows proper relocations
        //       against local TLS data to be performed directly without having to go and collect
        //       all TLS locals into one single blob.
        #[cfg(target_arch = "aarch64")]
        std::arch::asm!("msr tpidr_el0, xzr");

        entry(Box::leak(arg) as *mut _ as usize);
    }
}

fn main() -> ExitCode {
    let args: Vec<_> = env::args().collect();
    assert_ne!(args.len(), 0);

    if args.len() == 1 {
        eprintln!(
            "{}\n\n{} PROGRAM [ARGS...]",
            r#"
This program is the Yggdrasil OS dynamic executable linker/loader.
It is not meant to be called directly without arguments, but to be
used as an interpreter for dynamically-linked programs instead.

If needed, the program can still be invoked like follows:
"#
            .trim(),
            args[0]
        );
        return ExitCode::SUCCESS;
    };

    let args = &args[1..];
    let program = PathBuf::from(&args[0]);

    match run(program, args) {
        // Normal execution doesn't return here, but if LD_TRACE_LOADED_OBJECTS is set,
        // the loader will exit after printing everything
        Ok(()) => ExitCode::SUCCESS,
        Err(error) => {
            eprintln!("Error: {}", error);
            ExitCode::FAILURE
        }
    }
}