pkh/src/deb/ephemeral.rs

use crate::context::{self, Context, ContextConfig};
use crate::ui::deb::{DebUi, Phase};
use directories::ProjectDirs;
use std::error::Error;
use std::fs;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tar::Archive;
use xz2::read::XzDecoder;

/// An ephemeral unshare context guard that creates and manages a temporary chroot environment
/// for building packages with unshare permissions.
pub struct EphemeralContextGuard {
    previous_context: String,
    chroot_path: PathBuf,
    build_succeeded: bool,
    base_ctx: Arc<Context>,
}

impl EphemeralContextGuard {
    /// Create a new ephemeral unshare context with an explicit base context
    ///
    /// # Arguments
    /// * `series` - The distribution series (e.g., "noble", "sid")
    /// * `arch` - Optional target architecture. If provided and different from host,
    ///   downloads a chroot for that architecture (uses qemu_binfmt transparently)
    /// * `base_ctx` - The base context to use for creating the chroot
    pub async fn new_with_context(
        series: &str,
        arch: Option<&str>,
        base_ctx: Arc<Context>,
        ui: Option<Arc<DebUi>>,
    ) -> Result<Self, Box<dyn Error>> {
        let current_context_name = context::manager().current_name();

        // Create a temporary directory for the chroot
        let chroot_path_str = base_ctx.create_temp_dir()?;
        let chroot_path = PathBuf::from(chroot_path_str);

        log::debug!(
            "Creating new chroot for {} (arch: {:?}) at {}...",
            series,
            arch,
            chroot_path.display()
        );

        // Download and extract the chroot tarball
        Self::download_and_extract_chroot(series, arch, &chroot_path, base_ctx.clone(), &ui)
            .await?;

        // Switch to an ephemeral context to build the package in the chroot
        context::manager().set_current_ephemeral(Context::new(ContextConfig::Unshare {
            path: chroot_path.to_string_lossy().to_string(),
            parent: Some(current_context_name.clone()),
        }));

        Ok(Self {
            previous_context: current_context_name,
            chroot_path,
            build_succeeded: false,
            base_ctx,
        })
    }

    async fn download_and_extract_chroot(
        series: &str,
        arch: Option<&str>,
        chroot_path: &PathBuf,
        ctx: Arc<context::Context>,
        ui: &Option<Arc<DebUi>>,
    ) -> Result<(), Box<dyn Error>> {
        // Clone ctx for use in create_device_nodes after download_chroot_tarball consumes it
        let ctx_for_devices = ctx.clone();
        // Get project directories for caching
        let proj_dirs = ProjectDirs::from("com", "pkh", "pkh")
            .ok_or("Could not determine project directories")?;
        let cache_dir = proj_dirs.cache_dir();
        fs::create_dir_all(cache_dir)?;

        // Create tarball filename based on series and architecture
        let tarball_filename = if let Some(a) = arch {
            format!("{}-{}-buildd.tar.xz", series, a)
        } else {
            format!("{}-buildd.tar.xz", series)
        };
        let tarball_path = cache_dir.join(&tarball_filename);

        // Check for existing lockfile, and wait for a timeout if it exists
        // After timeout, warn the user
        let lockfile_path = tarball_path.with_extension("lock");

        // Check if lockfile exists and wait for it to be removed
        let mut wait_time = 0;
        let timeout = 300; // 5 minutes timeout
        let poll_interval = 5; // Check every 5 seconds

        while ctx.exists(&lockfile_path)? {
            if wait_time >= timeout {
                log::warn!(
                    "Lockfile {} exists and has been present for more than {} seconds. \
                     Another process may be downloading the chroot tarball. Continuing anyway...",
                    lockfile_path.display(),
                    timeout
                );
                break;
            }

            log::info!(
                "Lockfile {} exists, waiting for download to complete... ({}s/{})",
                lockfile_path.display(),
                wait_time,
                timeout
            );

            std::thread::sleep(std::time::Duration::from_secs(poll_interval));
            wait_time += poll_interval;
        }

        // Download tarball if it doesn't exist
        if !tarball_path.exists() {
            log::debug!(
                "Downloading chroot tarball for {} (arch: {:?})...",
                series,
                arch
            );
            if let Some(u) = ui {
                u.phase(Phase::PreparingChroot);
            }
            Self::download_chroot_tarball(series, arch, &tarball_path, ctx, ui).await?;
        } else {
            log::debug!(
                "Using cached chroot tarball for {} (arch: {:?})",
                series,
                arch
            );
        }

        // Extract tarball to chroot directory
        log::debug!("Extracting chroot tarball to {}...", chroot_path.display());
        if let Some(u) = ui {
            u.phase(Phase::ExtractingChroot);
        }
        Self::extract_tarball(&tarball_path, chroot_path, ui.as_deref())?;

        // Create device nodes in the chroot
        log::debug!("Creating device nodes in chroot...");
        if let Some(u) = ui {
            u.phase(Phase::FinalizingChroot);
        }
        Self::create_device_nodes(chroot_path, ctx_for_devices.clone())?;

        // Bind mount /proc from host into chroot (before entering unshare namespace)
        // This allows /proc to work in containers where mounting inside unshare fails
        log::debug!("Bind-mounting /proc into chroot...");
        Self::bind_mount_proc(chroot_path, ctx_for_devices)?;

        Ok(())
    }

    async fn download_chroot_tarball(
        series: &str,
        arch: Option<&str>,
        tarball_path: &Path,
        ctx: Arc<context::Context>,
        ui: &Option<Arc<DebUi>>,
    ) -> Result<(), Box<dyn Error>> {
        // Create a lock file to make sure that noone tries to use the file while it's not fully downloaded
        let lockfile_path = tarball_path.with_extension("lock");
        ctx.command("touch")
            .arg(lockfile_path.to_string_lossy().to_string())
            .status()?;

        // Download the keyring(s)
        let keyring_dir =
            crate::apt::keyring::download_cache_keyrings(Some(ctx.clone()), series).await?;

        // Use mmdebstrap to download the tarball to the cache directory
        let mut cmd = ctx.command("mmdebstrap");
        cmd.arg("--variant=buildd")
            .arg("--mode=unshare")
            .arg("--include=mount,curl,ca-certificates")
            .arg("--format=tar")
            .arg(format!("--keyring={}", keyring_dir.display()))
            // Setup hook to copy keyrings into the chroot so apt inside can use them
            .arg("--setup-hook=mkdir -p \"$1/etc/apt/trusted.gpg.d\"")
            .arg(format!(
                "--setup-hook=cp {}/*.gpg \"$1/etc/apt/trusted.gpg.d/\"",
                keyring_dir.display()
            ));

        // Add architecture if specified
        if let Some(a) = arch {
            cmd.arg(format!("--arch={}", a));
        }

        cmd.arg(series)
            .arg(tarball_path.to_string_lossy().to_string());

        if let Some(u) = ui {
            cmd.capture(u.sink());
        }

        let status = cmd.status()?;

        if !status.success() {
            // Remove file on error
            let _ = ctx
                .command("rm")
                .arg("-f")
                .arg(tarball_path.to_string_lossy().to_string())
                .status();
            let _ = ctx
                .command("rm")
                .arg("-f")
                .arg(lockfile_path.to_string_lossy().to_string())
                .status();
            return Err(format!(
                "Failed to download chroot tarball for series {} (arch: {:?})",
                series, arch
            )
            .into());
        }

        // Remove lockfile: tarball is fully downloaded
        let _ = ctx
            .command("rm")
            .arg("-f")
            .arg(lockfile_path.to_string_lossy().to_string())
            .status();

        Ok(())
    }

    fn extract_tarball(
        tarball_path: &PathBuf,
        chroot_path: &PathBuf,
        ui: Option<&DebUi>,
    ) -> Result<(), Box<dyn Error>> {
        // Create the chroot directory
        fs::create_dir_all(chroot_path)?;

        // Open the tarball file
        let tarball_file = std::fs::File::open(tarball_path)?;
        let xz_decoder = XzDecoder::new(tarball_file);
        let mut archive = Archive::new(xz_decoder);

        // Extract entries one by one so progress can be reported (a full
        // second decompression pass just to count entries upfront would be
        // too expensive for multi-hundred-MB chroot tarballs)
        let mut count = 0usize;
        for entry in archive.entries()? {
            let mut entry = entry?;
            entry.unpack_in(chroot_path)?;
            count += 1;
            if count.is_multiple_of(100)
                && let Some(u) = ui
            {
                u.progress_message(&format!("Extracting chroot… ({count} files)"));
            }
        }
        if let Some(u) = ui {
            u.progress_message(&format!("Extracting chroot… ({count} files)"));
        }

        Ok(())
    }

    fn create_device_nodes(
        chroot_path: &Path,
        ctx: Arc<context::Context>,
    ) -> Result<(), Box<dyn Error>> {
        let dev_null_path = chroot_path.join("dev/null");
        let dev_zero_path = chroot_path.join("dev/zero");

        // Ensure /dev directory exists
        fs::create_dir_all(chroot_path.join("dev"))?;

        // Remove existing device nodes if they exist
        let _ = ctx
            .command("rm")
            .arg("-f")
            .arg(dev_null_path.to_string_lossy().to_string())
            .status();

        let _ = ctx
            .command("rm")
            .arg("-f")
            .arg(dev_zero_path.to_string_lossy().to_string())
            .status();

        // Check if we're running as root
        let is_root = crate::utils::root::is_root()?;

        // Create new device nodes using mknod (with sudo if not root)
        let mut cmd_null = ctx.command(if is_root { "mknod" } else { "sudo" });
        if !is_root {
            cmd_null.arg("mknod");
        }
        let status_null = cmd_null
            .arg("-m")
            .arg("666")
            .arg(dev_null_path.to_string_lossy().to_string())
            .arg("c")
            .arg("1")
            .arg("3")
            .status()?;

        let mut cmd_zero = ctx.command(if is_root { "mknod" } else { "sudo" });
        if !is_root {
            cmd_zero.arg("mknod");
        }
        let status_zero = cmd_zero
            .arg("-m")
            .arg("666")
            .arg(dev_zero_path.to_string_lossy().to_string())
            .arg("c")
            .arg("1")
            .arg("5")
            .status()?;

        if !status_null.success() || !status_zero.success() {
            return Err("Failed to create device nodes".into());
        }

        Ok(())
    }

    /// Bind mount /proc from host into the chroot
    /// This is done before entering the unshare namespace, so it works in containers
    fn bind_mount_proc(
        chroot_path: &Path,
        ctx: Arc<context::Context>,
    ) -> Result<(), Box<dyn Error>> {
        let proc_path = chroot_path.join("proc");

        // Ensure /proc directory exists in chroot
        fs::create_dir_all(&proc_path)?;

        // Check if we're running as root
        let is_root = crate::utils::root::is_root()?;

        // Bind mount host's /proc into chroot (with sudo if not root)
        let mut cmd = ctx.command(if is_root { "mount" } else { "sudo" });
        if !is_root {
            cmd.arg("mount");
        }
        let status = cmd
            .arg("--bind")
            .arg("/proc")
            .arg(proc_path.to_string_lossy().to_string())
            .status()?;

        if !status.success() {
            log::warn!(
                "Could not bind-mount /proc into chroot at {}. Some packages may not install correctly.",
                proc_path.display()
            );
        } else {
            log::debug!("Bind-mounted /proc into chroot at {}", proc_path.display());
        }

        Ok(())
    }

    /// Mark the build as successful, which will trigger chroot cleanup on drop
    pub fn mark_build_successful(&mut self) {
        self.build_succeeded = true;
    }
}

impl Drop for EphemeralContextGuard {
    fn drop(&mut self) {
        log::debug!("Cleaning up ephemeral context ({:?})...", self.chroot_path);

        // Clean up any overlay mounts before resetting the context.
        // This must happen while the ephemeral context is still current so its
        // driver is accessible. The actual unmount commands run via the parent
        // (base) context, so they work regardless.
        let ephemeral_ctx = context::current();
        if let Err(e) = ephemeral_ctx.cleanup() {
            log::warn!("Failed to clean up overlay mounts: {}", e);
        }

        // Reset to normal context
        if let Err(e) = context::manager().set_current(&self.previous_context) {
            log::error!("Failed to restore context {}: {}", self.previous_context, e);
        }

        // Remove chroot directory only if build succeeded
        if self.build_succeeded {
            log::debug!(
                "Build succeeded, removing chroot directory: {}",
                self.chroot_path.display()
            );

            // Check if we're running as root to avoid unnecessary sudo
            let is_root = crate::utils::root::is_root().unwrap_or(false);

            // Unmount /proc from chroot before removing (ignore errors)
            let proc_path = self.chroot_path.join("proc");
            let _ = if is_root {
                self.base_ctx.command("umount").arg(&proc_path).status()
            } else {
                self.base_ctx
                    .command("sudo")
                    .arg("umount")
                    .arg(&proc_path)
                    .status()
            };

            let result = if is_root {
                self.base_ctx
                    .command("rm")
                    .arg("-rf")
                    .arg(&self.chroot_path)
                    .status()
            } else {
                self.base_ctx
                    .command("sudo")
                    .arg("rm")
                    .arg("-rf")
                    .arg(&self.chroot_path)
                    .status()
            };

            match result {
                Ok(status) => {
                    if !status.success() {
                        log::error!(
                            "Failed to remove chroot directory {}",
                            self.chroot_path.display()
                        );
                    } else {
                        log::debug!(
                            "Successfully removed chroot directory: {}",
                            self.chroot_path.display()
                        );
                    }
                }
                Err(e) => {
                    log::error!(
                        "Failed to execute cleanup command for {}: {}",
                        self.chroot_path.display(),
                        e
                    );
                }
            }
        } else {
            log::debug!(
                "Build did not succeed or was not marked as successful, keeping chroot directory: {}",
                self.chroot_path.display()
            );
        }
    }
}