//! User and group id mapping for the sandbox's user namespaces.
//!
//! Two modes:
//! - **single**: the runner's uid and gid map to container 0 and nothing else. Jobs can't
//!   `chown` to other ids.
//! - **subuid**: container 0 maps to the runner's uid, and container `1..=n` map to the range
//!   granted in `/etc/subuid` and `/etc/subgid`. Writing a map with more than one line needs the
//!   setuid helpers `newuidmap` and `newgidmap` from shadow-utils, so that is what the helper
//!   runs.
//!
//! The namespace entry functions ([`enter_mapped_userns`]) must run in a single-threaded process.
//! The sandbox helpers are single-threaded, and [`enter_mapped_userns`] forks in subuid mode.

use std::ffi::OsString;
use std::fs::{self, File};
use std::io::{self, Read, Write};
use std::os::unix::fs::PermissionsExt;
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};

use nix::errno::Errno;
use nix::fcntl::OFlag;
use nix::mount::{MsFlags, mount};
use nix::sched::{CloneFlags, unshare};
use nix::sys::wait::{WaitStatus, waitpid};
use nix::unistd::{ForkResult, Pid, Uid, User, fork, getgid, getpid, getuid, pipe2};
use serde::{Deserialize, Serialize};

/// The most ids a subordinate range may contribute. This is the usual size of a range.
const MAX_SUBID_COUNT: u32 = 65536;
/// Subordinate ranges smaller than this are not used at all.
const MIN_SUBID_COUNT: u32 = 1000;
/// Where the setuid mapping tools are looked for, in order.
const TOOL_DIRS: [&str; 4] = ["/usr/bin", "/bin", "/usr/sbin", "/sbin"];

/// How the sandbox maps ids. Set by `[sandbox] userns` in the runner config.
#[derive(Clone, Copy, Debug, Default, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum UsernsMode {
    /// Subuid mode when the host supports it, otherwise single.
    #[default]
    Auto,
    /// Only the runner's own uid and gid.
    Single,
    /// Subordinate ids. A host that can't provide them is an error, not a fallback.
    Subuid,
}

/// The id mapping a sandbox uses. The helper receives it as JSON.
///
/// `sub_uid` and `sub_gid` are `(start, count)` ranges on the host. Both are set in subuid mode.
/// Both must have the same count, so container ids `1..=count` are valid for uids and gids.
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct IdMap {
    pub uid: u32,
    pub gid: u32,
    pub sub_uid: Option<(u32, u32)>,
    pub sub_gid: Option<(u32, u32)>,
}

/// Errors from resolving or entering an id mapping.
#[derive(Debug, thiserror::Error)]
pub enum IdMapError {
    /// Subuid mode was asked for, but the host can't provide it.
    #[error("subuid mode unavailable: {0}")]
    Unavailable(String),
    /// `unshare(CLONE_NEWUSER)` failed.
    #[error("user namespaces unavailable: {0}")]
    Namespace(String),
    /// Writing or mapping ids failed, or making the mounts private.
    #[error("{0}")]
    Map(String),
}

impl IdMap {
    /// The runner's own uid and gid, with no subordinate ranges.
    pub fn single() -> Self {
        Self {
            uid: getuid().as_raw(),
            gid: getgid().as_raw(),
            sub_uid: None,
            sub_gid: None,
        }
    }

    /// Resolves `mode` to a concrete map. See the module docs for the rules.
    pub fn resolve(mode: UsernsMode) -> Result<Self, IdMapError> {
        match mode {
            UsernsMode::Single => Ok(Self::single()),
            UsernsMode::Subuid => subuid_map().map_err(IdMapError::Unavailable),
            UsernsMode::Auto => Ok(subuid_map().unwrap_or_else(|reason| {
                tracing::info!("user namespace ids: using the runner's own id only, {reason}");
                Self::single()
            })),
        }
    }

    /// Whether the map has subordinate ranges.
    pub fn is_subuid(&self) -> bool {
        self.sub_uid.is_some() && self.sub_gid.is_some()
    }

    /// How many container ids are mapped, including container 0. `1 + n` in subuid mode.
    pub fn size(&self) -> u32 {
        match (self.sub_uid, self.sub_gid) {
            (Some((_, n_uid)), Some((_, n_gid))) => 1 + n_uid.min(n_gid),
            _ => 1,
        }
    }

    /// A short description for logs, for example `subuid 524288+65536`.
    pub fn describe(&self) -> String {
        match (self.sub_uid, self.sub_gid) {
            (Some((start, count)), Some(_)) => format!("subuid {start}+{count}"),
            _ => format!("single uid {}", self.uid),
        }
    }
}

impl Default for IdMap {
    fn default() -> Self {
        Self::single()
    }
}

/// Parses `/etc/subuid` or `/etc/subgid` text. It returns the `(start, count)` ranges whose owner
/// is `user` (by name) or `uid` (by number), in file order. Comments, blank lines and malformed
/// lines are skipped.
pub fn parse_subid(content: &str, user: &str, uid: u32) -> Vec<(u32, u32)> {
    content
        .lines()
        .filter_map(|line| {
            let line = line.trim();
            if line.is_empty() || line.starts_with('#') {
                return None;
            }
            let mut fields = line.split(':');
            let (Some(owner), Some(start), Some(count), None) =
                (fields.next(), fields.next(), fields.next(), fields.next())
            else {
                return None;
            };
            let matches =
                (!user.is_empty() && owner == user) || owner.parse::<u32>().ok() == Some(uid);
            if !matches {
                return None;
            }
            let start: u32 = start.parse().ok()?;
            let count: u32 = count.parse().ok()?;
            if count == 0 || start.checked_add(count).is_none() {
                return None;
            }
            Some((start, count))
        })
        .collect()
}

/// The range to use. It is the first one with at least 65536 ids. Failing that, it is the first
/// one with at least 1000, which the caller reports as a warning. Smaller ranges are not used.
fn pick_range(ranges: &[(u32, u32)]) -> Option<(u32, u32)> {
    ranges
        .iter()
        .copied()
        .find(|&(_, count)| count >= MAX_SUBID_COUNT)
        .or_else(|| {
            ranges
                .iter()
                .copied()
                .find(|&(_, count)| count >= MIN_SUBID_COUNT)
        })
}

/// Builds the subuid map for the current user, or says why it can't.
fn subuid_map() -> Result<IdMap, String> {
    for tool in ["newuidmap", "newgidmap"] {
        if find_tool(tool).is_none() {
            return Err(format!("{tool} is not installed (shadow-utils)"));
        }
    }
    let uid = getuid().as_raw();
    let gid = getgid().as_raw();
    let name = User::from_uid(Uid::from_raw(uid))
        .ok()
        .flatten()
        .map(|user| user.name)
        .unwrap_or_default();
    let who = if name.is_empty() {
        format!("uid {uid}")
    } else {
        name.clone()
    };

    // /etc/subgid entries are also keyed by the user, so the uid is the key there too.
    let uid_ranges = find_ranges("/etc/subuid", &name, uid)?;
    let gid_ranges = find_ranges("/etc/subgid", &name, uid)?;
    let (uid_start, uid_count) = pick_range(&uid_ranges)
        .ok_or_else(|| format!("no usable range in /etc/subuid for {who}"))?;
    let (gid_start, gid_count) = pick_range(&gid_ranges)
        .ok_or_else(|| format!("no usable range in /etc/subgid for {who}"))?;
    let count = uid_count.min(gid_count).min(MAX_SUBID_COUNT);
    if count < MAX_SUBID_COUNT {
        tracing::warn!(
            count,
            "the subordinate id range for {who} is smaller than {MAX_SUBID_COUNT}, \
             so jobs get only {count} ids"
        );
    }
    Ok(IdMap {
        uid,
        gid,
        sub_uid: Some((uid_start, count)),
        sub_gid: Some((gid_start, count)),
    })
}

fn find_ranges(path: &str, user: &str, uid: u32) -> Result<Vec<(u32, u32)>, String> {
    match fs::read_to_string(path) {
        Ok(text) => Ok(parse_subid(&text, user, uid)),
        Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(Vec::new()),
        Err(e) => Err(format!("cannot read {path}: {e}")),
    }
}

fn find_tool(name: &str) -> Option<PathBuf> {
    TOOL_DIRS
        .iter()
        .map(|dir| Path::new(dir).join(name))
        .find(|path| path.is_file())
}

/// Enters a new user and mount namespace whose ids are mapped as `map` says. The process is root
/// inside it, and its mounts are private, so they don't propagate to the host.
///
/// Must be called from a single-threaded process, because subuid mode forks. The caller is the
/// sandbox helper, which has no other threads.
pub fn enter_mapped_userns(map: &IdMap) -> Result<(), IdMapError> {
    match (map.sub_uid, map.sub_gid) {
        (Some((uid_start, count_uid)), Some((gid_start, count_gid))) => {
            let count = count_uid.min(count_gid).min(MAX_SUBID_COUNT);
            enter_subuid(map, uid_start, gid_start, count)?;
        }
        _ => enter_single(map)?,
    }
    mount(
        None::<&str>,
        "/",
        None::<&str>,
        MsFlags::MS_REC | MsFlags::MS_PRIVATE,
        None::<&str>,
    )
    .map_err(|e| IdMapError::Map(format!("cannot make mounts private: {e}")))
}

/// Single mode: unshare, deny setgroups, and write the one-line maps directly.
fn enter_single(map: &IdMap) -> Result<(), IdMapError> {
    unshare(CloneFlags::CLONE_NEWUSER | CloneFlags::CLONE_NEWNS)
        .map_err(|e| IdMapError::Namespace(e.to_string()))?;
    fs::write("/proc/self/setgroups", "deny")
        .map_err(|e| IdMapError::Map(format!("cannot write setgroups: {e}")))?;
    fs::write("/proc/self/gid_map", format!("0 {} 1\n", map.gid))
        .map_err(|e| IdMapError::Map(format!("cannot write gid_map: {e}")))?;
    fs::write("/proc/self/uid_map", format!("0 {} 1\n", map.uid))
        .map_err(|e| IdMapError::Map(format!("cannot write uid_map: {e}")))
}

/// What the forked child needs to write the maps.
struct MapJob {
    parent: i32,
    uid: u32,
    gid: u32,
    uid_start: u32,
    gid_start: u32,
    count: u32,
}

/// Subuid mode. The parent unshares, and a forked child writes the maps with `newuidmap` and
/// `newgidmap`. The child can't do this itself, because the maps belong to the parent's new
/// namespace.
///
/// The child waits on a pipe until the parent has unshared. The parent then writes one byte, and
/// the child runs the two tools. The child reports its failure through a second pipe, which the
/// parent reads to EOF before it reaps the child.
fn enter_subuid(map: &IdMap, uid_start: u32, gid_start: u32, count: u32) -> Result<(), IdMapError> {
    let pipe_err = |e: Errno| IdMapError::Map(format!("cannot create pipe: {e}"));
    let (go_r, go_w) = pipe2(OFlag::O_CLOEXEC).map_err(pipe_err)?;
    let (err_r, err_w) = pipe2(OFlag::O_CLOEXEC).map_err(pipe_err)?;
    let job = MapJob {
        parent: getpid().as_raw(),
        uid: map.uid,
        gid: map.gid,
        uid_start,
        gid_start,
        count,
    };

    // SAFETY: the helper is single-threaded (see the module docs), so the child is a copy of a
    // process with one thread, and no lock held by another thread can block it. The child only
    // reads a pipe, runs std::process::Command, and exits.
    match unsafe { fork() }.map_err(|e| IdMapError::Map(format!("cannot fork: {e}")))? {
        ForkResult::Child => {
            drop(go_w);
            drop(err_r);
            let mut go = File::from(go_r);
            let mut err = File::from(err_w);
            let code = match write_maps_in_child(&mut go, &job) {
                Ok(()) => 0,
                Err(message) => {
                    let _ = err.write_all(message.as_bytes());
                    1
                }
            };
            std::process::exit(code)
        }
        ForkResult::Parent { child } => {
            drop(go_r);
            drop(err_w);
            let mut go_w = File::from(go_w);
            let mut err_r = File::from(err_r);
            let signalled = unshare(CloneFlags::CLONE_NEWUSER | CloneFlags::CLONE_NEWNS)
                .map_err(|e| IdMapError::Namespace(e.to_string()))
                .and_then(|()| {
                    go_w.write_all(b"1").map_err(|e| {
                        IdMapError::Map(format!("cannot signal the mapping helper: {e}"))
                    })
                });
            // Closing the pipe tells the child to exit if it is still waiting.
            drop(go_w);
            let mut message = Vec::new();
            let _ = err_r.read_to_end(&mut message);
            let status = wait_for_child(child);
            signalled?;
            match status {
                Ok(WaitStatus::Exited(_, 0)) => {}
                Ok(_) => {
                    return Err(IdMapError::Map(format!(
                        "cannot map user ids: {}",
                        String::from_utf8_lossy(&message).trim()
                    )));
                }
                Err(e) => {
                    return Err(IdMapError::Map(format!(
                        "waiting for the mapping helper: {e}"
                    )));
                }
            }
            // newgidmap leaves setgroups allowed. Drop the supplementary groups the process had
            // before the namespace existed.
            nix::unistd::setgroups(&[])
                .map_err(|e| IdMapError::Map(format!("cannot clear groups: {e}")))
        }
    }
}

/// Runs in the forked child. It waits for the parent's byte, then maps the uids and gids.
fn write_maps_in_child(go: &mut File, job: &MapJob) -> Result<(), String> {
    let mut byte = [0u8; 1];
    loop {
        match go.read(&mut byte) {
            Ok(1) => break,
            Ok(_) => return Err("user namespace setup was cancelled".into()),
            Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
            Err(e) => return Err(format!("cannot read from the parent: {e}")),
        }
    }
    let parent = job.parent.to_string();
    let count = job.count.to_string();
    let uid_start = job.uid_start.to_string();
    let gid_start = job.gid_start.to_string();
    let uid = job.uid.to_string();
    let gid = job.gid.to_string();
    run_map_tool(
        "newuidmap",
        &[&parent, "0", &uid, "1", "1", &uid_start, &count],
    )?;
    run_map_tool(
        "newgidmap",
        &[&parent, "0", &gid, "1", "1", &gid_start, &count],
    )
}

fn run_map_tool(tool: &str, args: &[&str]) -> Result<(), String> {
    let path = find_tool(tool).ok_or_else(|| format!("{tool} is not installed"))?;
    let output = Command::new(&path)
        .args(args)
        .stdin(Stdio::null())
        .output()
        .map_err(|e| format!("cannot run {tool}: {e}"))?;
    if output.status.success() {
        Ok(())
    } else {
        Err(format!(
            "{tool} failed ({}): {}",
            output.status,
            String::from_utf8_lossy(&output.stderr).trim()
        ))
    }
}

fn wait_for_child(child: Pid) -> Result<WaitStatus, Errno> {
    loop {
        match waitpid(child, None) {
            Err(Errno::EINTR) => continue,
            other => return other,
        }
    }
}

/// `<exe> __sandbox-rm <path>`, with the [`IdMap`] JSON on stdin. Enters the namespace, makes the
/// directories under `path` writable, and removes it. Exit 0 on success, or 125 with a
/// `forge-sandbox:` line on stderr.
pub fn rm_helper_main(args: &[OsString]) -> i32 {
    match rm_helper(args) {
        Ok(()) => 0,
        Err(message) => {
            eprintln!("forge-sandbox: {message}");
            125
        }
    }
}

fn rm_helper(args: &[OsString]) -> Result<(), String> {
    let [path] = args else {
        return Err("usage: __sandbox-rm <path>".into());
    };
    let path = PathBuf::from(path);
    if !path.is_absolute() {
        return Err("path must be absolute".into());
    }
    let map: IdMap = serde_json::from_reader(io::stdin().lock())
        .map_err(|e| format!("cannot read the id map: {e}"))?;
    enter_mapped_userns(&map).map_err(|e| e.to_string())?;
    remove_local(&path).map_err(|e| format!("cannot remove {}: {e}", path.display()))
}

/// Removes `path` (a file, symlink or directory tree) as the current user. Missing is not an
/// error. If the first attempt fails, the directories under `path` get `u+rwx` and it is retried.
fn remove_local(path: &Path) -> io::Result<()> {
    match remove_once(path) {
        Ok(()) => Ok(()),
        Err(first) => {
            let is_dir = fs::symlink_metadata(path).is_ok_and(|m| m.is_dir());
            if !is_dir {
                return Err(first);
            }
            make_dirs_writable(path)?;
            remove_once(path)
        }
    }
}

fn remove_once(path: &Path) -> io::Result<()> {
    match fs::symlink_metadata(path) {
        Ok(meta) if meta.is_dir() => fs::remove_dir_all(path),
        Ok(_) => fs::remove_file(path),
        Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(()),
        Err(e) => Err(e),
    }
}

/// Adds `u+rwx` to every directory under `root`, without following symlinks.
fn make_dirs_writable(root: &Path) -> io::Result<()> {
    let mut stack = vec![root.to_owned()];
    while let Some(dir) = stack.pop() {
        let mode = fs::symlink_metadata(&dir)?.permissions().mode() & 0o7777;
        if mode & 0o700 != 0o700 {
            fs::set_permissions(&dir, fs::Permissions::from_mode(mode | 0o700))?;
        }
        for entry in fs::read_dir(&dir)? {
            let entry = entry?;
            if entry.file_type()?.is_dir() {
                stack.push(entry.path());
            }
        }
    }
    Ok(())
}

/// Removes `path` when its files may be owned by subordinate ids. Single mode removes it
/// directly. Subuid mode runs `<exe> __sandbox-rm <path>` with the map on stdin, because the
/// runner can't delete files owned by ids it doesn't have. A missing path is `Ok`.
pub fn remove_tree(map: &IdMap, path: &Path, exe: &Path) -> io::Result<()> {
    if !map.is_subuid() {
        return remove_local(path);
    }
    match fs::symlink_metadata(path) {
        Ok(_) => {}
        Err(e) if e.kind() == io::ErrorKind::NotFound => return Ok(()),
        Err(e) => return Err(e),
    }
    let json = serde_json::to_vec(map).map_err(io::Error::other)?;
    let mut child = Command::new(exe)
        .arg("__sandbox-rm")
        .arg(path)
        .stdin(Stdio::piped())
        .stdout(Stdio::null())
        .stderr(Stdio::piped())
        .spawn()?;
    {
        // The helper reads stdin to EOF, so the pipe must be closed before waiting.
        let mut stdin = child
            .stdin
            .take()
            .ok_or_else(|| io::Error::other("no stdin for the removal helper"))?;
        stdin.write_all(&json)?;
    }
    let output = child.wait_with_output()?;
    if output.status.success() {
        Ok(())
    } else {
        Err(io::Error::other(format!(
            "cannot remove {}: {}",
            path.display(),
            String::from_utf8_lossy(&output.stderr).trim()
        )))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// Whether `unshare --map-auto` works here. The tests that need subordinate ids skip when it
    /// doesn't.
    fn map_auto_works() -> bool {
        Command::new("unshare")
            .args(["--map-auto", "--map-root-user", "true"])
            .stdout(Stdio::null())
            .stderr(Stdio::null())
            .status()
            .is_ok_and(|s| s.success())
    }

    #[test]
    fn parse_subid_matches_name_and_uid() {
        let text = "# comment\n\nalice:100000:65536\ntahinli:524288:65536\n\
                    1000:200000:1000\nbob:x:5\nbad line\ncarol:1:2:3\nzero:9:0\n";
        assert_eq!(
            parse_subid(text, "tahinli", 1000),
            vec![(524288, 65536), (200000, 1000)]
        );
        // With no name, only the numeric owner matches.
        assert_eq!(parse_subid(text, "", 1000), vec![(200000, 1000)]);
        assert_eq!(parse_subid(text, "carol", 1), Vec::new());
        assert_eq!(parse_subid(text, "nobody", 4242), Vec::new());
    }

    #[test]
    fn parse_subid_skips_ranges_that_overflow() {
        let text = "tahinli:4294967295:2\ntahinli:10:5\n";
        assert_eq!(parse_subid(text, "tahinli", 1000), vec![(10, 5)]);
    }

    #[test]
    fn pick_range_prefers_full_ranges() {
        assert_eq!(pick_range(&[(1, 500), (2, 1000)]), Some((2, 1000)));
        assert_eq!(
            pick_range(&[(1, 500), (9, 65536), (3, 70000)]),
            Some((9, 65536))
        );
        assert_eq!(pick_range(&[(1, 500)]), None);
        assert_eq!(pick_range(&[]), None);
    }

    #[test]
    fn size_and_describe() {
        let single = IdMap {
            uid: 1000,
            gid: 1000,
            sub_uid: None,
            sub_gid: None,
        };
        assert!(!single.is_subuid());
        assert_eq!(single.size(), 1);
        assert_eq!(single.describe(), "single uid 1000");

        let sub = IdMap {
            uid: 1000,
            gid: 1000,
            sub_uid: Some((524288, 65536)),
            sub_gid: Some((524288, 65536)),
        };
        assert!(sub.is_subuid());
        assert_eq!(sub.size(), 65537);
        assert_eq!(sub.describe(), "subuid 524288+65536");

        let uneven = IdMap {
            sub_gid: Some((600000, 10)),
            ..sub
        };
        assert_eq!(uneven.size(), 11);
    }

    #[test]
    fn default_is_single() {
        let map = IdMap::default();
        assert_eq!(map, IdMap::single());
        assert!(!map.is_subuid());
    }

    #[test]
    fn auto_resolves_to_subuid_on_this_host() {
        if !map_auto_works() {
            eprintln!("skipping: unshare --map-auto is not available");
            return;
        }
        let map = IdMap::resolve(UsernsMode::Auto).expect("auto never fails");
        assert!(map.is_subuid(), "expected subuid, got {}", map.describe());
    }
}