fix: make multiple shells use the same overlay (#4)

This avoids conflicts between the different overlays all using the same upper & work directories. Also, now changes in any shell are visible consistently everywhere. Co-authored-by: raphaelcoeffic <[email protected]>
raphaelcoeffic · Nov 7, 2024 · 4a682d6 · 4a682d6
1 parent 2ced057
commit 4a682d6
Show file tree

Hide file tree

Showing 8 changed files with 331 additions and 95 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -39,6 +39,7 @@ serde_json = "1.0"
 tar = { workspace = true }
 tempfile = { workspace = true }
 which = "6.0.3"
+fd-lock = "4.0.2"
 
 [[bin]]
 name = "dive"

diff --git a/src/main.rs b/src/main.rs
@@ -1,28 +1,30 @@
 use std::{
     ffi::OsString,
-    fmt::Debug,
-    fs::{create_dir_all, read_link},
+    fs::read_link,
     io,
     os::unix::process::CommandExt,
     path::{Path, PathBuf},
-    process::{exit, Command},
+    process::{self, exit, Command},
 };
 
-use anyhow::{Context, Result};
+use anyhow::{bail, Context, Result};
 use clap::Parser;
-use namespaces::enter_namespaces_as_root;
-use overlay::OverlayMount;
-use pid_lookup::pid_lookup;
 use procfs::process::Process;
 use rustix::{
-    path::Arg,
-    process::geteuid,
-    thread::{unshare, UnshareFlags},
+    process::{geteuid, waitpid, WaitOptions},
+    runtime::{fork, Fork},
 };
 
 mod namespaces;
 mod overlay;
+mod pid_file;
 mod pid_lookup;
+mod shared_mount;
+
+use namespaces::*;
+use overlay::*;
+use pid_lookup::*;
+use shared_mount::*;
 
 #[cfg(feature = "embedded_image")]
 mod embedded_image;
@@ -72,18 +74,6 @@ fn get_overlay_dir() -> PathBuf {
     dirs::state_dir().unwrap().join(APP_NAME).join(OVL_DIR)
 }
 
-fn make_overlay_dirs() -> Result<(PathBuf, PathBuf, PathBuf)> {
-    let overlay_dir = get_overlay_dir();
-    let upper_dir = overlay_dir.join("upper");
-    let work_dir = overlay_dir.join("work");
-
-    create_dir_all(&upper_dir)
-        .and(create_dir_all(&work_dir))
-        .context("could not create state directory")?;
-
-    Ok((overlay_dir, upper_dir, work_dir))
-}
-
 fn init_logging() {
     env_logger::Builder::new()
         .filter_level(log::LevelFilter::Info)
@@ -111,74 +101,41 @@ fn reexec_with_sudo(
         .exec())
 }
 
-fn main() -> Result<()> {
-    let args = Args::parse();
-    init_logging();
-
-    let lead_pid = if let Ok(pid) = args.container_id.parse::<i32>() {
-        pid
-    } else if let Some(pid) = pid_lookup(&args.container_id) {
-        pid
-    } else {
-        log::error!("could not find container PID");
-        exit(exitcode::NOINPUT);
-    };
-    log::debug!("container PID: {}", lead_pid);
-
-    let img_dir = get_img_dir(&args);
-    if !img_dir.exists() || !img_dir.join("store").exists() {
-        #[cfg(feature = "embedded_image")]
-        embedded_image::install_base_image(&img_dir)
-            .context("could not unpack base image")?;
-
-        #[cfg(not(feature = "embedded_image"))]
-        BaseImageBuilder::new(&img_dir)
-            .build_base()
-            .context("could not build base image")?;
-    }
-
-    let (overlay_dir, upper_dir, work_dir) = make_overlay_dirs()?;
-
-    if !geteuid().is_root() {
-        log::debug!("re-executing with sudo...");
-        reexec_with_sudo(lead_pid, &img_dir, &overlay_dir)?
-    }
-
-    let overlay = match OverlayMount::new(img_dir, upper_dir, work_dir) {
-        Err(err) => {
-            log::error!("could not create base image mount: {:?}", err);
-            exit(exitcode::OSERR);
-        }
-        Ok(mnt) => {
-            log::debug!("detached mount created");
-            mnt
-        }
-    };
-
-    let proc_env = match Process::new(lead_pid).and_then(|p| p.environ()) {
+fn prepare_shell_environment(
+    shared_mount: &SharedMount,
+    lead_pid: i32,
+) -> Result<()> {
+    let detached_mount = match shared_mount.make_detached_mount() {
         Err(err) => {
-            log::error!("could not fetch the process environment: {err}");
-            exit(exitcode::OSERR);
+            bail!("could not make detached mount: {err}");
         }
-        Ok(env) => env,
+        Ok(m) => m,
     };
+    if let Err(err) = enter_namespaces_as_root(lead_pid) {
+        bail!("cannot enter container namespaces: {err}");
+    }
+    if let Err(err) = detached_mount.mount_in_new_namespace("/nix") {
+        bail!("cannot mount /nix: {err}");
+    }
+    Ok(())
+}
 
-    enter_namespaces_as_root(lead_pid)?;
-
-    log::debug!("mounting overlay...");
-    unshare(UnshareFlags::NEWNS)
-        .context("could not create new mount namespace")?;
-    overlay
-        .mount("/nix")
-        .context("could not mount base image")?;
-
+fn exec_shell() -> Result<()> {
+    //
     // TODO: path HOME w/ user as defined by /etc/passwd
-
+    //
     // TODO: find shell in this order:
     // - zsh
     // - bash
     // - sh at last
 
+    let proc_env = match Process::new(1).and_then(|p| p.environ()) {
+        Err(err) => {
+            bail!("could not fetch the process environment: {err}");
+        }
+        Ok(env) => env,
+    };
+
     let mut cmd = Command::new("zsh");
     cmd.env_clear();
     cmd.envs(&proc_env);
@@ -216,7 +173,67 @@ fn main() -> Result<()> {
         ("INFOPATH", format!("{nix_base}/share/info")),
     ]);
 
-    cmd.status()?;
+    let err = cmd.exec();
+    bail!("cannot exec: {}", err)
+}
 
+fn wait_for_child(child_pid: rustix::thread::Pid) -> Result<()> {
+    // TODO: propagate return code properly
+    log::debug!("parent pid = {}", process::id());
+    let _ = waitpid(Some(child_pid), WaitOptions::empty())
+        .context("waitpid failed")?;
     Ok(())
 }
+
+fn main() -> Result<()> {
+    let args = Args::parse();
+    init_logging();
+
+    let lead_pid = if let Some(pid) = pid_lookup(&args.container_id) {
+        pid
+    } else {
+        log::error!("could not find container PID");
+        exit(exitcode::NOINPUT);
+    };
+    log::debug!("container PID: {}", lead_pid);
+
+    let img_dir = get_img_dir(&args);
+    if !img_dir.exists() || !img_dir.join("store").exists() {
+        #[cfg(feature = "embedded_image")]
+        embedded_image::install_base_image(&img_dir)
+            .context("could not unpack base image")?;
+
+        #[cfg(not(feature = "embedded_image"))]
+        BaseImageBuilder::new(&img_dir)
+            .build_base()
+            .context("could not build base image")?;
+    }
+
+    let overlay_dir = get_overlay_dir();
+    let overlay_builder = OverlayBuilder::new(&overlay_dir, &img_dir)?;
+
+    if !geteuid().is_root() {
+        log::debug!("re-executing with sudo...");
+        reexec_with_sudo(lead_pid, &img_dir, &overlay_dir)?
+    }
+
+    let shared_mount = SharedMount::new(&overlay_dir, overlay_builder)
+        .context("could not init shared mount")?;
+
+    match unsafe { fork()? } {
+        Fork::Child(_) => {
+            if let Err(err) = prepare_shell_environment(&shared_mount, lead_pid)
+            {
+                log::error!("{err}");
+                exit(1);
+            }
+            // in normal cases, there is no return from exec_shell()
+            if let Err(err) = exec_shell() {
+                log::error!("cannot execute shell: {err}");
+                exit(1);
+            }
+            exit(0);
+        }
+        Fork::Parent(child_pid) => wait_for_child(child_pid),
+    }
+}
diff --git a/src/namespaces.rs b/src/namespaces.rs
@@ -3,7 +3,7 @@ use std::{
     process::exit,
 };
 
-use anyhow::Result;
+use anyhow::{bail, Result};
 use procfs::process::{Namespace, Process};
 use rustix::{
     process::{pidfd_open, Pid, PidfdFlags},
@@ -49,21 +49,18 @@ pub fn enter_namespaces_as_root(lead_pid: i32) -> Result<()> {
     let lead = Process::new(lead_pid)?;
     let lead_ns = match namespace_set(&lead) {
         Err(err) => {
-            log::error!("cannot inspect lead process namespaces: {err}");
-            exit(exitcode::OSERR);
+            bail!("cannot inspect lead process namespaces: {err}");
         }
         Ok(ns) => ns,
     };
-    log::debug!("lead_ns = {:?}", lead_ns);
 
     let me = Process::myself()?;
     let me_id = me.pid;
-    log::debug!("my pid = {}", me_id);
+    log::debug!("own pid is {}", me_id);
 
     let my_ns = match namespace_set(&me) {
         Err(err) => {
-            log::error!("cannot inspect own namespaces: {err}");
-            exit(exitcode::OSERR);
+            bail!("cannot inspect own namespaces: {err}");
         }
         Ok(ns) => ns,
     };
@@ -74,7 +71,6 @@ pub fn enter_namespaces_as_root(lead_pid: i32) -> Result<()> {
         ns_set = ns_set.union(namespace_type_by_name(&ns.0));
     }
 
-    log::debug!("ns_set: {:?}", ns_set);
     if ns_set.is_empty() {
         log::debug!("no point in entering anything: we're already there!");
         return Ok(());

diff --git a/src/overlay.rs b/src/overlay.rs
@@ -2,7 +2,7 @@ use std::fs::create_dir_all;
 use std::path::Path;
 use std::{os::fd::OwnedFd, path::PathBuf};
 
-use anyhow::Result;
+use anyhow::{Context, Result};
 use rustix::fd::AsFd;
 use rustix::fs;
 use rustix::mount::{
@@ -20,7 +20,7 @@ impl OverlayMount {
     ) -> Result<Self> {
         let lower_dir = lower_dir.into().canonicalize()?;
         let fsfd = fsopen("overlay", FsOpenFlags::FSOPEN_CLOEXEC)?;
-        fsconfig_set_string(fsfd.as_fd(), "source", "nsdb")?;
+        fsconfig_set_string(fsfd.as_fd(), "source", "user-data")?;
         fsconfig_set_string(fsfd.as_fd(), "lowerdir", lower_dir)?;
         fsconfig_set_string(fsfd.as_fd(), "upperdir", upper_dir.as_ref())?;
         fsconfig_set_string(fsfd.as_fd(), "workdir", work_dir.as_ref())?;
@@ -50,3 +50,41 @@ impl OverlayMount {
         Ok(())
     }
 }
+
+pub struct OverlayBuilder {
+    pub lower_dir: PathBuf,
+    pub merged_dir: PathBuf,
+    pub upper_dir: PathBuf,
+    pub work_dir: PathBuf,
+    pub pids_dir: PathBuf,
+}
+
+impl OverlayBuilder {
+    pub fn new<P, Q>(base_dir: P, lower_dir: Q) -> Result<Self>
+    where
+        P: AsRef<Path>,
+        Q: AsRef<Path>,
+    {
+        let base_dir = base_dir.as_ref();
+        let lower_dir = lower_dir.as_ref().to_owned();
+
+        let merged_dir = base_dir.join("merged");
+        let upper_dir = base_dir.join("upper");
+        let work_dir = base_dir.join("work");
+        let pids_dir = base_dir.join("pids");
+
+        create_dir_all(&merged_dir)
+            .and(create_dir_all(&upper_dir))
+            .and(create_dir_all(&work_dir))
+            .and(create_dir_all(&pids_dir))
+            .context("could not create overlay directories")?;
+
+        Ok(OverlayBuilder {
+            lower_dir,
+            merged_dir,
+            upper_dir,
+            work_dir,
+            pids_dir,
+        })
+    }
+}