onelf-rt 0.2.5

//! Environment variable setup for the running package.
//!
//! Sets `ONELF_*` variables and computes a `lib_path` string for the
//! dynamic linker's `--library-path` flag. Also auto-detects and configures
//! paths for graphics drivers (OpenGL/EGL/Vulkan/VA-API).
//!
//! `LD_LIBRARY_PATH` is intentionally NOT set for ELF entrypoints. It
//! would be inherited by every child process the packed app spawns,
//! including host binaries (`/bin/sh`, `ssh`, etc.), which corrupts them
//! by mixing the bundled libc/libcrypto with the host loader. Instead,
//! the lib path is passed via `--library-path` on a single linker
//! invocation.

use std::env;
use std::path::Path;

/// Set up environment variables and return a colon-joined lib path
/// string for use with the dynamic linker's `--library-path` flag.
///
/// For shebang scripts (non-ELF target), returns an empty string: the
/// kernel hands off to a host interpreter linked against the host glibc,
/// and pointing it at our bundled libs would mix two glibcs in one
/// process. Scripts that need bundled libs must export `LD_LIBRARY_PATH`
/// themselves before execing bundled binaries.
pub fn setup_env(
    onelf_dir: &str,
    argv0: &str,
    exec_path: &str,
    entrypoint_name: &str,
    mode: &str,
    lib_subpath: &str,
    target_path: &str,
) -> String {
    let launch_dir = env::current_dir()
        .ok()
        .and_then(|p| p.to_str().map(String::from))
        .unwrap_or_default();

    // SAFETY: the runtime is single-threaded at this point (before exec)
    unsafe {
        env::set_var("ONELF_DIR", onelf_dir);
        env::set_var("ONELF_ARGV0", argv0);
        env::set_var("ONELF_EXEC", exec_path);
        env::set_var("ONELF_ENTRYPOINT", entrypoint_name);
        env::set_var("ONELF_LAUNCH_DIR", &launch_dir);
        env::set_var("ONELF_MODE", mode);
    }

    if onelf_dir.is_empty() {
        return String::new();
    }

    let pkg = Path::new(onelf_dir);

    let target_is_elf = is_elf_file(target_path);

    let mut lib_path = String::new();

    // Build the library search path for ELF entrypoints. Order:
    //   <bundled lib dirs> : <existing LD_LIBRARY_PATH> : <host driver/system dirs>
    // Bundled libs win, but GPU / libGL / libcuda / libvulkan and other
    // host-provided userspace drivers are still discoverable. Our bundled
    // ld.so has its baked-in paths scrubbed, so drivers that normally
    // live in /usr/lib (or /run/opengl-driver/lib on NixOS) have to be
    // added here explicitly or Cycles/OptiX and similar features won't
    // find their driver libraries.
    if target_is_elf && !lib_subpath.is_empty() {
        let lib_paths: Vec<String> = lib_subpath
            .split(':')
            .map(|p| pkg.join(p).to_string_lossy().to_string())
            .collect();
        let lib_str = lib_paths.join(":");
        if !lib_str.is_empty() {
            let mut parts: Vec<String> = Vec::new();
            parts.push(lib_str);
            // Preserve the user's pre-existing LD_LIBRARY_PATH as a middle
            // layer, but don't propagate it to the child env.
            let existing = env::var("LD_LIBRARY_PATH").unwrap_or_default();
            if !existing.is_empty() {
                parts.push(existing);
            }
            let host_paths = host_driver_paths();
            if !host_paths.is_empty() {
                parts.push(host_paths.join(":"));
            }
            lib_path = parts.join(":");

            // Set LD_LIBRARY_PATH as a fallback. Ideally we'd only pass
            // bundled paths via --library-path on the linker invocation
            // (which avoids polluting child processes), but onelf's
            // AT_EXECFN bootstrap doesn't drive a linker invocation and
            // doesn't pass --library-path itself, so binaries with
            // bootstrap injected rely on LD_LIBRARY_PATH. The trade-off:
            // it gets inherited by every child process the app spawns.
            unsafe {
                env::set_var("LD_LIBRARY_PATH", &lib_path);
            }

            // Auto-set LIBGL_DRIVERS_PATH and LIBVA_DRIVERS_PATH if any lib dir
            // contains a dri/ subdirectory (both use the same paths)
            let dri_paths: Vec<String> = lib_paths
                .iter()
                .map(|p| Path::new(p).join("dri").to_string_lossy().to_string())
                .filter(|p| Path::new(p).is_dir())
                .collect();
            if !dri_paths.is_empty() {
                let joined = dri_paths.join(":");
                if env::var("LIBGL_DRIVERS_PATH").is_err() {
                    unsafe {
                        env::set_var("LIBGL_DRIVERS_PATH", &joined);
                    }
                }
                if env::var("LIBVA_DRIVERS_PATH").is_err() {
                    unsafe {
                        env::set_var("LIBVA_DRIVERS_PATH", &joined);
                    }
                }
            }

            // Auto-set GBM_BACKENDS_PATH if any lib dir contains a gbm/ subdirectory
            if env::var("GBM_BACKENDS_PATH").is_err() {
                let gbm_paths: Vec<String> = lib_paths
                    .iter()
                    .map(|p| Path::new(p).join("gbm").to_string_lossy().to_string())
                    .filter(|p| Path::new(p).is_dir())
                    .collect();
                if !gbm_paths.is_empty() {
                    unsafe {
                        env::set_var("GBM_BACKENDS_PATH", gbm_paths.join(":"));
                    }
                }
            }
        }
    }

    // Prepend package's share/ to XDG_DATA_DIRS so bundled GSettings schemas,
    // icons, mime types, etc. are discoverable by GLib/GTK. Host dirs are kept
    // so system themes, schemas, and desktop integrations still work.
    setup_xdg_data_dirs(pkg);

    // EGL vendor discovery: merge bundled + host dirs so both Mesa
    // and proprietary drivers (NVIDIA, AMD) are visible to libglvnd.
    if env::var("__EGL_VENDOR_LIBRARY_DIRS").is_err() {
        let mut egl_dirs: Vec<String> = Vec::new();
        let egl_dir = pkg.join("share/glvnd/egl_vendor.d");
        if egl_dir.is_dir() {
            egl_dirs.push(egl_dir.to_string_lossy().into_owned());
        }
        for d in &[
            "/run/opengl-driver/share/glvnd/egl_vendor.d",
            "/etc/glvnd/egl_vendor.d",
            "/usr/share/glvnd/egl_vendor.d",
        ] {
            if Path::new(d).is_dir() {
                egl_dirs.push((*d).to_string());
            }
        }
        if !egl_dirs.is_empty() {
            unsafe {
                env::set_var("__EGL_VENDOR_LIBRARY_DIRS", egl_dirs.join(":"));
            }
        }
    }

    // Auto-set DRIRC_CONFIGDIR if package has DRI config files
    if env::var("DRIRC_CONFIGDIR").is_err() {
        let drirc_dir = pkg.join("share/drirc.d");
        if drirc_dir.is_dir() {
            unsafe {
                env::set_var("DRIRC_CONFIGDIR", drirc_dir.to_string_lossy().as_ref());
            }
        }
    }

    // Auto-set LIBDRM_IDS_PATH if package has libdrm data
    if env::var("LIBDRM_IDS_PATH").is_err() {
        let libdrm_dir = pkg.join("share/libdrm");
        if libdrm_dir.is_dir() {
            unsafe {
                env::set_var("LIBDRM_IDS_PATH", libdrm_dir.to_string_lossy().as_ref());
            }
        }
    }

    // Auto-set XKB_CONFIG_ROOT if package has xkb data
    if env::var("XKB_CONFIG_ROOT").is_err() {
        let xkb_dir = pkg.join("share/X11/xkb");
        if xkb_dir.is_dir() {
            unsafe {
                env::set_var("XKB_CONFIG_ROOT", xkb_dir.to_string_lossy().as_ref());
            }
        }
    }

    // Auto-set LIBDECOR_PLUGIN_DIR if package has libdecor plugins
    if env::var("LIBDECOR_PLUGIN_DIR").is_err() {
        let libdecor_dir = pkg.join("share/libdecor/plugins-1");
        if libdecor_dir.is_dir() {
            unsafe {
                env::set_var(
                    "LIBDECOR_PLUGIN_DIR",
                    libdecor_dir.to_string_lossy().as_ref(),
                );
            }
        }
    }

    // Vulkan ICD discovery: use VK_ADD_DRIVER_FILES to *append* our
    // bundled ICD configs to the loader's default search. Setting
    // VK_DRIVER_FILES would *replace* the search and cut off host GPU
    // drivers (e.g. NVIDIA's ICD on NixOS at /run/opengl-driver/...).
    // Also include well-known host ICD paths that the bundled loader
    // can't find on its own (its compiled-in /etc and /usr paths are
    // scrubbed).
    if env::var("VK_DRIVER_FILES").is_err() && env::var("VK_ADD_DRIVER_FILES").is_err() {
        let mut icd_dirs: Vec<String> = Vec::new();

        let vk_dir = pkg.join("share/vulkan/icd.d");
        if vk_dir.is_dir() {
            icd_dirs.push(vk_dir.to_string_lossy().into_owned());
        }
        // Host ICD locations that the scrubbed loader can't reach.
        for d in &[
            "/run/opengl-driver/share/vulkan/icd.d",
            "/etc/vulkan/icd.d",
            "/usr/share/vulkan/icd.d",
        ] {
            if Path::new(d).is_dir() {
                icd_dirs.push((*d).to_string());
            }
        }

        if !icd_dirs.is_empty() {
            let mut all_files: Vec<String> = Vec::new();
            for dir in &icd_dirs {
                if let Ok(entries) = std::fs::read_dir(dir) {
                    for e in entries.filter_map(|e| e.ok()) {
                        if e.path().extension().map_or(false, |ext| ext == "json") {
                            all_files.push(e.path().to_string_lossy().into_owned());
                        }
                    }
                }
            }
            if !all_files.is_empty() {
                unsafe {
                    env::set_var("VK_DRIVER_FILES", all_files.join(":"));
                }
            }
        }
    }

    lib_path
}

/// Check whether `path` is an ELF file (first four bytes `\x7fELF`).
/// Scripts (shebang `#!`) return false; missing files also return false.
fn is_elf_file(path: &str) -> bool {
    use std::io::Read;
    let Ok(mut f) = std::fs::File::open(path) else {
        return false;
    };
    let mut buf = [0u8; 4];
    matches!(f.read(&mut buf), Ok(4)) && buf == *b"\x7fELF"
}

/// Return the host's driver / system library directories that exist on
/// this system, in descending priority order. These are appended to the
/// package's `LD_LIBRARY_PATH` so the bundled loader can locate host-
/// provided GPU userspace drivers (libcuda, libvulkan, libGL, libva)
/// without picking up host libraries that would clash with the bundle's
/// own copies — those come first in the search path.
fn host_driver_paths() -> Vec<String> {
    // NixOS exposes all GPU userspace drivers under /run/opengl-driver,
    // populated from the active nixpkgs `hardware.graphics` closure.
    // On other distros, drivers live in the standard multiarch or lib64
    // paths alongside the rest of the system's runtime libraries.
    const CANDIDATES: &[&str] = &[
        "/run/opengl-driver/lib",
        "/run/opengl-driver-32/lib",
        "/usr/lib/x86_64-linux-gnu",
        "/usr/lib64",
        "/usr/lib",
        "/lib/x86_64-linux-gnu",
        "/lib64",
    ];
    CANDIDATES
        .iter()
        .filter(|p| Path::new(p).is_dir())
        .map(|s| (*s).to_string())
        .collect()
}

/// Prepend the package's `share/` to `XDG_DATA_DIRS` so GLib/GTK can find
/// bundled GSettings schemas, icons, MIME types, etc. Host dirs are preserved
/// so system themes and desktop integrations still work.
fn setup_xdg_data_dirs(pkg: &Path) {
    let share = pkg.join("share");
    if !share.is_dir() {
        return;
    }

    let pkg_share = share.to_string_lossy();
    let existing = env::var("XDG_DATA_DIRS").unwrap_or_default();

    let new_val = if existing.is_empty() {
        // XDG spec default when unset is /usr/local/share:/usr/share
        format!("{pkg_share}:/usr/local/share:/usr/share")
    } else {
        format!("{pkg_share}:{existing}")
    };

    unsafe {
        env::set_var("XDG_DATA_DIRS", new_val);
    }
}

/// Apply custom environment variables from `.onelf/env` data.
/// Each line is `KEY=VALUE`. Values containing `${ONELF_DIR}` are
/// expanded to the package root at runtime.
pub fn apply_custom_env(env_data: &[u8], onelf_dir: &str) {
    let Ok(text) = std::str::from_utf8(env_data) else {
        return;
    };
    for line in text.lines() {
        let line = line.trim();
        if line.is_empty() || line.starts_with('#') {
            continue;
        }
        if let Some((key, val)) = line.split_once('=') {
            let expanded = val.replace("${ONELF_DIR}", onelf_dir);
            unsafe {
                env::set_var(key.trim(), expanded.trim());
            }
        }
    }
}