ecr/src/namespace.rs

use anyhow::{anyhow, Context, Result};
use nix::sched::{clone, CloneFlags};
use nix::sys::signal::Signal;
use nix::unistd::{getgid, getuid, Pid};

/// RAII wrapper that closes a raw file descriptor on drop.
/// Guarantees all pipe fds are closed on every return path, including
/// clone() and setup_user_namespace() failures.
struct AutoCloseFd(i32);

impl AutoCloseFd {
    fn raw(&self) -> i32 {
        self.0
    }
}

impl Drop for AutoCloseFd {
    fn drop(&mut self) {
        unsafe {
            libc::close(self.0);
        }
    }
}

/// Clone flags for namespace creation
const CLONE_FLAGS: CloneFlags = CloneFlags::CLONE_NEWUSER
    .union(CloneFlags::CLONE_NEWPID)
    .union(CloneFlags::CLONE_NEWNS)
    .union(CloneFlags::CLONE_NEWUTS);

/// Check if user namespaces are available
pub fn check_user_namespace() -> Result<()> {
    // Check kernel.unprivileged_userns_clone on systems that have it
    if let Ok(content) = std::fs::read_to_string("/proc/sys/kernel/unprivileged_userns_clone") {
        if content.trim() == "0" {
            return Err(anyhow!(
                "User namespaces not available\n\n\
                Enable with:\n\
                  sysctl -w kernel.unprivileged_userns_clone=1\n\n\
                Or check AppArmor profile restrictions."
            ));
        }
    }

    // Check max_user_namespaces
    if let Ok(content) = std::fs::read_to_string("/proc/sys/user/max_user_namespaces") {
        if let Ok(max) = content.trim().parse::<u32>() {
            if max == 0 {
                return Err(anyhow!(
                    "User namespaces not available\n\n\
                    Enable with:\n\
                      sysctl -w user.max_user_namespaces=10000"
                ));
            }
        }
    }

    Ok(())
}

/// Setup namespaces and run the provided function inside them.
/// Returns the child's exit code (0 on success, 128+signal when killed by a
/// signal). Setup failures are returned as Err.
pub fn setup_namespaces<F>(f: F) -> Result<i32>
where
    F: FnOnce() -> Result<()> + Send + 'static,
{
    // sync pipe: parent signals child to proceed after UID/GID mapping
    // done pipe: child signals parent it has finished (or exec'd the shell)
    // error pipe: child writes the anyhow error chain to the parent on failure.
    // The write end is O_CLOEXEC so it is automatically closed when execvp
    // succeeds — the parent then reads EOF and knows there was no error.
    //
    // All six fds are wrapped in AutoCloseFd so they are closed on every return
    // path, including clone() and setup_user_namespace() failures.
    let (parent_read, parent_write, child_read, child_write, error_read, error_write);

    unsafe {
        let mut fds: [i32; 2] = [-1, -1];

        if libc::pipe(fds.as_mut_ptr()) != 0 {
            return Err(anyhow!("Failed to create sync pipe"));
        }
        parent_read = AutoCloseFd(fds[0]);
        parent_write = AutoCloseFd(fds[1]);
        // parent_read/write auto-closed if subsequent pipes fail ↑

        if libc::pipe(fds.as_mut_ptr()) != 0 {
            return Err(anyhow!("Failed to create done pipe"));
        }
        child_read = AutoCloseFd(fds[0]);
        child_write = AutoCloseFd(fds[1]);
        // O_CLOEXEC on the write end: if execve succeeds the kernel closes cw
        // atomically, the parent's read(child_read) gets EOF immediately, and
        // waitpid becomes the real wait.  The done-pipe is then only used on
        // the error path (f() returned Err before execve was reached).
        libc::fcntl(child_write.raw(), libc::F_SETFD, libc::FD_CLOEXEC);

        if libc::pipe(fds.as_mut_ptr()) != 0 {
            return Err(anyhow!("Failed to create error pipe"));
        }
        error_read = AutoCloseFd(fds[0]);
        error_write = AutoCloseFd(fds[1]);
        // Same treatment for error_write: auto-closed on exec (no error),
        // written explicitly on the error path before the child exits.
        libc::fcntl(error_write.raw(), libc::F_SETFD, libc::FD_CLOEXEC);
    }

    // Stack for the child process
    let stack_size = crate::utils::CHILD_STACK_SIZE;
    let mut stack = vec![0u8; stack_size];

    // Wrap f in Option to allow taking it once inside the child closure
    let mut f = Some(f);

    // Extract raw fds for the child closure.  The child is a clone of the
    // parent process and gets its own copies of all open fds; the parent's
    // AutoCloseFd wrappers independently manage the parent's copies.
    let pr = parent_read.raw();
    let pw = parent_write.raw();
    let cr = child_read.raw();
    let cw = child_write.raw();
    let er = error_read.raw();
    let ew = error_write.raw();

    // Clone with new namespaces
    let pid = unsafe {
        clone(
            Box::new(move || {
                // Close unused pipe ends in the child
                libc::close(pw);
                libc::close(cr);
                libc::close(er);

                // Wait for parent to set up UID/GID mappings
                let mut buf = [0u8; 1];
                libc::read(pr, buf.as_mut_ptr() as *mut libc::c_void, 1);
                libc::close(pr);

                // Run the function
                let result = if let Some(func) = f.take() {
                    func()
                } else {
                    Err(anyhow!("Function already called"))
                };

                // On failure, write the full error chain to the error pipe
                // before signalling done, so the parent can reconstruct it.
                if let Err(ref e) = result {
                    let msg = format!("{:#}", e);
                    let bytes = msg.as_bytes();
                    libc::write(ew, bytes.as_ptr() as *const libc::c_void, bytes.len());
                }
                libc::close(ew);

                // Signal completion
                libc::write(cw, c"done".as_ptr() as *const libc::c_void, 4);
                libc::close(cw);

                if result.is_ok() {
                    0
                } else {
                    1
                }
            }),
            &mut stack,
            CLONE_FLAGS,
            Some(Signal::SIGCHLD as i32),
        )
    }
    .context("Failed to clone with new namespaces")?;
    // clone() failure: all six AutoCloseFds drop here, closing every fd. ✓

    // Parent: drop the child-side ends now that clone has succeeded.
    // The child process has its own copies; dropping here closes the parent's.
    drop(parent_read);
    drop(child_write);
    drop(error_write);

    // Set up UID/GID mappings for the child
    // setup_user_namespace failure: parent_write, child_read, error_read
    // auto-closed by AutoCloseFd drop. ✓
    setup_user_namespace(pid)?;

    // Signal child to proceed
    unsafe {
        libc::write(parent_write.raw(), c"go".as_ptr() as *const libc::c_void, 2);
    }
    drop(parent_write);

    // Wait for child to complete (or the exec'd shell to exit)
    let mut buf = [0u8; 4];
    unsafe {
        libc::read(child_read.raw(), buf.as_mut_ptr() as *mut libc::c_void, 4);
    }
    drop(child_read);

    // Read the error message written by the child, if any.
    // error_write was either closed explicitly (on error) or auto-closed via
    // O_CLOEXEC (on successful exec), so this read always terminates.
    let child_error: Option<String> = unsafe {
        let mut error_bytes = Vec::new();
        let mut tmp = [0u8; crate::utils::ERROR_BUFFER_SIZE];
        loop {
            let n = libc::read(
                error_read.raw(),
                tmp.as_mut_ptr() as *mut libc::c_void,
                tmp.len(),
            );
            if n <= 0 {
                break;
            }
            error_bytes.extend_from_slice(&tmp[..n as usize]);
        }
        if error_bytes.is_empty() {
            None
        } else {
            Some(String::from_utf8_lossy(&error_bytes).into_owned())
        }
    };
    drop(error_read);

    // Wait for child process
    let status = nix::sys::wait::waitpid(pid, None)?;

    match status {
        nix::sys::wait::WaitStatus::Exited(_, 0) => Ok(0),
        nix::sys::wait::WaitStatus::Exited(_, code) => {
            // If the child reported an error (e.g., setup failure), return it.
            // Otherwise, just forward the exit code without an error message.
            if let Some(msg) = child_error {
                Err(anyhow!("{}", msg))
            } else {
                Ok(code)
            }
        }
        nix::sys::wait::WaitStatus::Signaled(_, sig, _) => Ok(128 + sig as i32),
        _ => Ok(0),
    }
}

/// Set up UID/GID mappings for user namespace
fn setup_user_namespace(pid: Pid) -> Result<()> {
    let uid = getuid();
    let gid = getgid();

    // Get the subordinate UID/GID ranges from /etc/subuid and /etc/subgid
    // For unprivileged users, we need to use these ranges
    let (sub_uid_start, sub_uid_count) = get_subuid_range(uid)?;
    let (sub_gid_start, sub_gid_count) = get_subgid_range(gid)?;

    // We map UID 0 (root inside the namespace) to the host user, then map
    // IDs 1..sub_uid_count-1 to the subordinate range.  A count of 0 or 1
    // leaves no subordinate IDs to map and indicates a malformed /etc/subuid.
    if sub_uid_count < 2 {
        return Err(anyhow!(
            "subuid count {} for uid {} is too small (need at least 2); \
             check /etc/subuid",
            sub_uid_count,
            uid
        ));
    }
    if sub_gid_count < 2 {
        return Err(anyhow!(
            "subgid count {} for gid {} is too small (need at least 2); \
             check /etc/subgid",
            sub_gid_count,
            gid
        ));
    }

    // Allow setgroups so apt and other tools can drop privileges
    let setgroups_path = format!("/proc/{}/setgroups", pid);
    std::fs::write(&setgroups_path, "allow\n")
        .with_context(|| format!("Failed to write {}", setgroups_path))?;

    // Use newuidmap and newgidmap for setting up mappings
    // These are setuid binaries that allow unprivileged users to map subuid/subgid ranges
    let pid_str = pid.to_string();

    // newuidmap format: newuidmap pid ns_start host_start count ...
    // Map current user to root (0), then subordinate UIDs starting from 1
    let uid_result = std::process::Command::new("newuidmap")
        .arg(&pid_str)
        .arg("0")
        .arg(uid.to_string())
        .arg("1")
        .arg("1")
        .arg(sub_uid_start.to_string())
        .arg((sub_uid_count - 1).to_string())
        .status()
        .context("Failed to execute newuidmap")?;

    if !uid_result.success() {
        return Err(anyhow!(
            "newuidmap failed - ensure subuid entry exists in /etc/subuid"
        ));
    }

    // newgidmap format: newgidmap pid ns_start host_start count ...
    // Map current group to root (0), then subordinate GIDs starting from 1
    let gid_result = std::process::Command::new("newgidmap")
        .arg(&pid_str)
        .arg("0")
        .arg(gid.to_string())
        .arg("1")
        .arg("1")
        .arg(sub_gid_start.to_string())
        .arg((sub_gid_count - 1).to_string())
        .status()
        .context("Failed to execute newgidmap")?;

    if !gid_result.success() {
        return Err(anyhow!(
            "newgidmap failed - ensure subgid entry exists in /etc/subgid"
        ));
    }

    Ok(())
}

/// Get subordinate UID range for a user from /etc/subuid
fn get_subuid_range(uid: nix::unistd::Uid) -> Result<(u32, u32)> {
    let content = std::fs::read_to_string("/etc/subuid").context("Failed to read /etc/subuid")?;

    let username =
        users::get_user_by_uid(uid.as_raw()).map(|u| u.name().to_string_lossy().to_string());

    for line in content.lines() {
        let parts: Vec<&str> = line.split(':').collect();
        if parts.len() >= 3 {
            // Check if this line matches our user (by name or UID)
            let matches = parts[0] == username.as_deref().unwrap_or("")
                || parts[0].parse::<u32>().ok() == Some(uid.as_raw());

            if matches {
                let start: u32 = parts[1].parse().context("Invalid subuid start")?;
                let count: u32 = parts[2].parse().context("Invalid subuid count")?;
                return Ok((start, count));
            }
        }
    }

    Err(anyhow!(
        "No subuid entry found for user {} (uid {}). \
        Add one to /etc/subuid, e.g.:\n  {}:100000:65536",
        username.as_deref().unwrap_or("<unknown>"),
        uid,
        username.as_deref().unwrap_or(&uid.to_string()),
    ))
}

/// Get subordinate GID range for a group from /etc/subgid
fn get_subgid_range(gid: nix::unistd::Gid) -> Result<(u32, u32)> {
    let content = std::fs::read_to_string("/etc/subgid").context("Failed to read /etc/subgid")?;

    let groupname =
        users::get_group_by_gid(gid.as_raw()).map(|g| g.name().to_string_lossy().to_string());

    for line in content.lines() {
        let parts: Vec<&str> = line.split(':').collect();
        if parts.len() >= 3 {
            // Check if this line matches our group (by name or GID)
            let matches = parts[0] == groupname.as_deref().unwrap_or("")
                || parts[0].parse::<u32>().ok() == Some(gid.as_raw());

            if matches {
                let start: u32 = parts[1].parse().context("Invalid subgid start")?;
                let count: u32 = parts[2].parse().context("Invalid subgid count")?;
                return Ok((start, count));
            }
        }
    }

    Err(anyhow!(
        "No subgid entry found for group {} (gid {}). \
        Add one to /etc/subgid, e.g.:\n  {}:100000:65536",
        groupname.as_deref().unwrap_or("<unknown>"),
        gid,
        groupname.as_deref().unwrap_or(&gid.to_string()),
    ))
}

/// Set hostname in UTS namespace
pub fn set_hostname(distro: &str) -> Result<()> {
    use nix::unistd::sethostname;
    use std::collections::hash_map::DefaultHasher;
    use std::hash::{Hash, Hasher};

    // Seed from both the current time and the PID so each invocation gets a
    // distinct suffix even when called in rapid succession.  Time alone is
    // not sufficient: truncating nanoseconds to u8 in a tight loop produces
    // the same byte every iteration.
    let mut hasher = DefaultHasher::new();
    std::time::SystemTime::now().hash(&mut hasher);
    std::process::id().hash(&mut hasher);
    let mut state = hasher.finish();

    let chars = b"abcdefghijklmnopqrstuvwxyz0123456789";
    // Use HOSTNAME_SUFFIX_BITS for entropy (6 hex chars = 24 bits)
    let suffix_len = (crate::utils::HOSTNAME_SUFFIX_BITS as f64).log2() as usize / 4;
    let random_suffix: String = (0..suffix_len)
        .map(|_| {
            // Knuth multiplicative LCG — each step advances the full 64-bit state.
            state = state
                .wrapping_mul(6364136223846793005)
                .wrapping_add(1442695040888963407);
            chars[(state >> 33) as usize % chars.len()] as char
        })
        .collect();

    let hostname = format!("ecr-{}-{}", distro, random_suffix);

    sethostname(&hostname).with_context(|| format!("Failed to set hostname to {}", hostname))?;

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::hash::{Hash, Hasher};

    #[test]
    fn test_check_user_namespace_returns_ok() {
        // This test verifies the function runs without panicking
        // On most modern Linux systems with user namespaces enabled, this should pass
        let result = check_user_namespace();
        // We can't assert success because it depends on system configuration
        // But we can verify it doesn't panic and returns a Result
        assert!(result.is_ok() || result.is_err());
    }

    #[test]
    fn test_hostname_format() {
        // Test that hostname generation produces valid format
        use std::collections::HashSet;

        let mut hostnames = HashSet::new();
        for _ in 0..100 {
            let mut hasher = std::collections::hash_map::DefaultHasher::new();
            std::time::SystemTime::now().hash(&mut hasher);
            std::process::id().hash(&mut hasher);
            let mut state = hasher.finish();

            let chars = b"abcdefghijklmnopqrstuvwxyz0123456789";
            let suffix_len = (crate::utils::HOSTNAME_SUFFIX_BITS as f64).log2() as usize / 4;
            let random_suffix: String = (0..suffix_len)
                .map(|_| {
                    state = state
                        .wrapping_mul(6364136223846793005)
                        .wrapping_add(1442695040888963407);
                    chars[(state >> 33) as usize % chars.len()] as char
                })
                .collect();

            let hostname = format!("ecr-test-{}", random_suffix);

            // Verify hostname format
            assert!(hostname.starts_with("ecr-test-"));
            assert!(hostname.len() > 9); // "ecr-test-" + at least 1 char
            assert!(hostname
                .chars()
                .all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '-'));

            hostnames.insert(hostname);
        }

        // With 100 iterations and good entropy, we should get many unique hostnames
        assert!(
            hostnames.len() > 50,
            "Expected many unique hostnames, got {}",
            hostnames.len()
        );
    }

    #[test]
    fn test_set_hostname_uniqueness() {
        // Verify that rapid consecutive calls produce different hostnames
        use std::collections::HashSet;

        let mut hostnames = Vec::new();
        for _ in 0..10 {
            // Simulate the hostname generation logic
            let mut hasher = std::collections::hash_map::DefaultHasher::new();
            std::time::SystemTime::now().hash(&mut hasher);
            std::process::id().hash(&mut hasher);
            let mut state = hasher.finish();

            let chars = b"abcdefghijklmnopqrstuvwxyz0123456789";
            let suffix_len = (crate::utils::HOSTNAME_SUFFIX_BITS as f64).log2() as usize / 4;
            let random_suffix: String = (0..suffix_len)
                .map(|_| {
                    state = state
                        .wrapping_mul(6364136223846793005)
                        .wrapping_add(1442695040888963407);
                    chars[(state >> 33) as usize % chars.len()] as char
                })
                .collect();

            hostnames.push(format!("ecr-test-{}", random_suffix));
        }

        let unique: HashSet<_> = hostnames.iter().collect();
        // Most hostnames should be unique (high entropy)
        assert!(unique.len() >= 8, "Expected mostly unique hostnames");
    }
}