feat: make rts less verbose, support 2 different rts for nfsboot

hil/src/commands/nfsboot.rs

@@ -1,4 +1,4 @@
-use camino::Utf8PathBuf;
+use camino::{Utf8Path, Utf8PathBuf};
 use clap::Parser;
 use color_eyre::{
     eyre::{bail, WrapErr},
@@ -13,7 +13,7 @@ use crate::nfsboot::{error_detection_for_host_state, request_sudo, MountSpec};
 /// Boot the orb using NFS
 #[derive(Debug, Parser)]
 pub struct Nfsboot {
-    /// The s3 URI of the RTS to use.
+    /// The s3 URI of the RTS to use for NFS boot.
     #[arg(
         long,
         conflicts_with = "rts_path",
@@ -23,14 +23,24 @@ pub struct Nfsboot {
     /// The directory to save the s3 artifact we download.
     #[arg(long)]
     download_dir: Option<Utf8PathBuf>,
-    /// Path to a downloaded RTS (zipped .tar or an already-extracted directory).
+    /// Path to a downloaded RTS (zipped .tar or an already-extracted directory)
+    /// used for NFS boot.
     #[arg(
         long,
         conflicts_with = "s3_url",
         conflicts_with = "download_dir",
         required_unless_present = "s3_url"
     )]
     rts_path: Option<Utf8PathBuf>,
+    /// S3 URI of a separate RTS whose `rts/` directory is used for `--mount`
+    /// content instead of the boot RTS. Use this when NFS-booting with a dev
+    /// build but flashing a stage or prod build.
+    #[arg(long, conflicts_with = "mount_rts_path")]
+    mount_s3_url: Option<S3Uri>,
+    /// Local path to a separate RTS tarball whose `rts/` directory is used for
+    /// `--mount` content instead of the boot RTS.
+    #[arg(long, conflicts_with = "mount_s3_url")]
+    mount_rts_path: Option<Utf8PathBuf>,
     /// If this flag is given, overwites any existing files when downloading the rts.
     #[arg(long)]
     overwrite_existing: bool,
@@ -51,12 +61,18 @@ impl Nfsboot {
         error_detection_for_host_state()
             .await
             .wrap_err("failed to assert host's state")?;
-        let rts_path = self.maybe_download_rts().await?;
-        debug!("resolved RTS path: {rts_path}");
+        let rts_path = self.resolve_rts_path().await?;
+        debug!("resolved boot RTS path: {rts_path}");
+
+        let mount_rts_path = self.resolve_mount_rts_path().await?;
+        if let Some(ref p) = mount_rts_path {
+            debug!("resolved mount RTS path: {p}");
+        }
 
         let rng = rand::rngs::StdRng::from_rng(rand::thread_rng()).unwrap();
         let _mount_guard = crate::nfsboot::nfsboot(
             rts_path,
+            mount_rts_path,
             self.mounts,
             self.persistent_img_path.as_deref().map(|p| p.as_std_path()),
             rng,
@@ -69,24 +85,24 @@ impl Nfsboot {
         unreachable!()
     }
 
-    async fn maybe_download_rts(&self) -> Result<Utf8PathBuf> {
-        let existing_file_behavior = if self.overwrite_existing {
+    fn existing_file_behavior(&self) -> ExistingFileBehavior {
+        if self.overwrite_existing {
             ExistingFileBehavior::Overwrite
         } else {
             ExistingFileBehavior::Abort
-        };
-        // Determine RTS tarball path: download from S3 or use provided path
+        }
+    }
+
+    async fn resolve_rts_path(&self) -> Result<Utf8PathBuf> {
+        let existing_file_behavior = self.existing_file_behavior();
         let rts_path = if let Some(ref s3_url) = self.s3_url {
             assert!(
                 self.rts_path.is_none(),
                 "sanity: mutual exclusion guaranteed by clap"
             );
-            let download_dir =
-                self.download_dir.clone().unwrap_or_else(crate::current_dir);
-            let download_path = download_dir.join(
-                crate::download_s3::parse_filename(s3_url)
-                    .wrap_err("failed to parse filename")?,
-            );
+            let download_path = self
+                .download_path_for_s3_url(s3_url)
+                .wrap_err("failed to resolve boot rts download path")?;
 
             crate::download_s3::download_url(
                 s3_url,
@@ -114,4 +130,196 @@ impl Nfsboot {
 
         Ok(rts_path)
     }
+
+    async fn resolve_mount_rts_path(&self) -> Result<Option<Utf8PathBuf>> {
+        let path = if let Some(ref s3_url) = self.mount_s3_url {
+            assert!(
+                self.mount_rts_path.is_none(),
+                "sanity: mutual exclusion guaranteed by clap"
+            );
+            let download_path = self
+                .mount_download_path(s3_url)
+                .wrap_err("failed to resolve mount rts download path")?;
+
+            crate::download_s3::download_url(
+                s3_url,
+                &download_path,
+                self.existing_file_behavior(),
+            )
+            .await
+            .wrap_err("error while downloading mount rts from s3")?;
+
+            Some(download_path)
+        } else {
+            self.mount_rts_path.clone()
+        };
+
+        Ok(path)
+    }
+
+    fn download_dir(&self) -> Utf8PathBuf {
+        self.download_dir.clone().unwrap_or_else(crate::current_dir)
+    }
+
+    fn download_path_for_s3_url(&self, s3_url: &S3Uri) -> Result<Utf8PathBuf> {
+        let file_name = crate::download_s3::parse_filename(s3_url)
+            .wrap_err("failed to parse s3 filename")?;
+
+        Ok(self.download_dir().join(file_name))
+    }
+
+    fn mount_download_path(&self, mount_s3_url: &S3Uri) -> Result<Utf8PathBuf> {
+        let mount_download_path = self
+            .download_path_for_s3_url(mount_s3_url)
+            .wrap_err("failed to parse mount s3 filename")?;
+
+        let collides_with_boot = if let Some(ref boot_s3_url) = self.s3_url {
+            let boot_download_path = self
+                .download_path_for_s3_url(boot_s3_url)
+                .wrap_err("failed to parse boot s3 filename")?;
+            mount_download_path == boot_download_path
+        } else if let Some(ref rts_path) = self.rts_path {
+            same_file(rts_path, &mount_download_path)
+        } else {
+            false
+        };
+
+        if !collides_with_boot {
+            return Ok(mount_download_path);
+        }
+
+        let mount_file_name = mount_download_path
+            .file_name()
+            .ok_or_else(|| color_eyre::eyre::eyre!("mount rts path has no filename"))?;
+
+        Ok(self.download_dir().join(format!("mount-{mount_file_name}")))
+    }
+}
+
+/// Compare two paths for equality by filename only, to catch cases where one
+/// is relative (e.g. `./rts.tar.zstd`) and the other absolute
+/// (e.g. `/home/user/rts.tar.zstd`).
+fn same_file(a: &Utf8Path, b: &Utf8Path) -> bool {
+    match (a.canonicalize_utf8(), b.canonicalize_utf8()) {
+        (Ok(a), Ok(b)) => a == b,
+        _ => a.file_name() == b.file_name(),
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::Nfsboot;
+    use camino::Utf8PathBuf;
+    use orb_s3_helpers::S3Uri;
+
+    const DEV_S3: &str = "s3://worldcoin-orb-resources/worldcoin/orb-os/rts/2025-08-14-heads-main-0-g0a8d01b-diamond/rts-diamond-dev.tar.zstd";
+    const STAGE_S3: &str = "s3://worldcoin-orb-resources/worldcoin/orb-os/rts/2025-08-14-heads-main-0-g0a8d01b-diamond/rts-diamond-stage.tar.zstd";
+
+    const DEV_FILENAME: &str =
+        "2025-08-14-heads-main-0-g0a8d01b-diamond-rts-diamond-dev.tar.zstd";
+    const STAGE_FILENAME: &str =
+        "2025-08-14-heads-main-0-g0a8d01b-diamond-rts-diamond-stage.tar.zstd";
+
+    #[test]
+    fn mount_download_path_disambiguated_when_both_s3_urls_collide() {
+        let s3_url: S3Uri = DEV_S3.parse().unwrap();
+        let command = Nfsboot {
+            s3_url: Some(s3_url.clone()),
+            download_dir: Some(Utf8PathBuf::from("/tmp/dl")),
+            rts_path: None,
+            mount_s3_url: Some(s3_url.clone()),
+            mount_rts_path: None,
+            overwrite_existing: false,
+            mounts: vec![],
+            persistent_img_path: None,
+        };
+
+        let path = command.mount_download_path(&s3_url).unwrap();
+        assert_eq!(
+            path,
+            Utf8PathBuf::from(format!("/tmp/dl/mount-{DEV_FILENAME}"))
+        );
+    }
+
+    #[test]
+    fn mount_download_path_unchanged_when_s3_urls_differ() {
+        let boot: S3Uri = DEV_S3.parse().unwrap();
+        let mount: S3Uri = STAGE_S3.parse().unwrap();
+        let command = Nfsboot {
+            s3_url: Some(boot),
+            download_dir: Some(Utf8PathBuf::from("/tmp/dl")),
+            rts_path: None,
+            mount_s3_url: Some(mount.clone()),
+            mount_rts_path: None,
+            overwrite_existing: false,
+            mounts: vec![],
+            persistent_img_path: None,
+        };
+
+        let path = command.mount_download_path(&mount).unwrap();
+        assert_eq!(path, Utf8PathBuf::from(format!("/tmp/dl/{STAGE_FILENAME}")));
+    }
+
+    #[test]
+    fn mount_download_path_disambiguated_when_collides_with_local_rts() {
+        let mount: S3Uri = DEV_S3.parse().unwrap();
+        let local_rts_path = Utf8PathBuf::from(format!("/tmp/dl/{DEV_FILENAME}"));
+        let command = Nfsboot {
+            s3_url: None,
+            download_dir: Some(Utf8PathBuf::from("/tmp/dl")),
+            rts_path: Some(local_rts_path),
+            mount_s3_url: Some(mount.clone()),
+            mount_rts_path: None,
+            overwrite_existing: false,
+            mounts: vec![],
+            persistent_img_path: None,
+        };
+
+        let path = command.mount_download_path(&mount).unwrap();
+        assert_eq!(
+            path,
+            Utf8PathBuf::from(format!("/tmp/dl/mount-{DEV_FILENAME}"))
+        );
+    }
+
+    #[test]
+    fn mount_download_path_disambiguated_with_relative_local_rts() {
+        let mount: S3Uri = DEV_S3.parse().unwrap();
+        let local_rts_path = Utf8PathBuf::from(format!("./{DEV_FILENAME}"));
+        let command = Nfsboot {
+            s3_url: None,
+            download_dir: Some(Utf8PathBuf::from("/tmp/dl")),
+            rts_path: Some(local_rts_path),
+            mount_s3_url: Some(mount.clone()),
+            mount_rts_path: None,
+            overwrite_existing: false,
+            mounts: vec![],
+            persistent_img_path: None,
+        };
+
+        let path = command.mount_download_path(&mount).unwrap();
+        assert_eq!(
+            path,
+            Utf8PathBuf::from(format!("/tmp/dl/mount-{DEV_FILENAME}"))
+        );
+    }
+
+    #[test]
+    fn mount_download_path_unchanged_when_local_rts_differs() {
+        let mount: S3Uri = STAGE_S3.parse().unwrap();
+        let local_rts_path = Utf8PathBuf::from(format!("/tmp/dl/{DEV_FILENAME}"));
+        let command = Nfsboot {
+            s3_url: None,
+            download_dir: Some(Utf8PathBuf::from("/tmp/dl")),
+            rts_path: Some(local_rts_path),
+            mount_s3_url: Some(mount.clone()),
+            mount_rts_path: None,
+            overwrite_existing: false,
+            mounts: vec![],
+            persistent_img_path: None,
+        };
+
+        let path = command.mount_download_path(&mount).unwrap();
+        assert_eq!(path, Utf8PathBuf::from(format!("/tmp/dl/{STAGE_FILENAME}")));
+    }
 }

hil/src/nfsboot.rs

@@ -41,8 +41,14 @@ impl From<Mounter> for MountGuard {
 /// nfsboot.sh from the rts.
 ///
 /// The filesystems will remain mounted until `cancel` is cancelled.
+/// * `path_to_rts` - RTS tarball used for NFS boot (rootfs, nfsbootcmd).
+/// * `mount_rts_path` - Optional separate RTS tarball whose `rts/` directory
+///   is used for `--mount` content. When a stage or prod build is provided
+///   here, the dev build in `path_to_rts` handles NFS boot while this one
+///   supplies the content that gets flashed onto the orb.
 pub async fn nfsboot(
     path_to_rts: Utf8PathBuf,
+    mount_rts_path: Option<Utf8PathBuf>,
     mut mounts: Vec<MountSpec>,
     persistent_img_path: Option<&Path>,
     rng: impl rand::Rng + Send + 'static,
@@ -57,9 +63,26 @@ pub async fn nfsboot(
         "we expected a directory called `rts` after extracting"
     );
 
-    for m in mounts.iter_mut().filter(|m| m.host_path == "/rtsdir") {
-        m.host_path = rts_dir.clone().try_into().unwrap();
-    }
+    let mount_rts_tmp_dir = if let Some(mount_rts) = mount_rts_path {
+        let tmp = tokio::task::spawn_blocking(move || extract(&mount_rts))
+            .await
+            .wrap_err("task panicked")??;
+        debug!("mount rts temp dir: {tmp:?}");
+        let mount_rts_dir = tmp.path().join("rts");
+        assert!(
+            tokio::fs::try_exists(&mount_rts_dir).await.unwrap_or(false),
+            "we expected a directory called `rts` in mount RTS tarball"
+        );
+        for m in mounts.iter_mut().filter(|m| m.host_path == "/rtsdir") {
+            m.host_path = mount_rts_dir.clone().try_into().unwrap();
+        }
+        Some(tmp)
+    } else {
+        for m in mounts.iter_mut().filter(|m| m.host_path == "/rtsdir") {
+            m.host_path = rts_dir.clone().try_into().unwrap();
+        }
+        None
+    };
 
     if let Some(persistent_img_path) = persistent_img_path {
         crate::rts::populate_persistent(tmp_dir.path(), persistent_img_path, rng)
@@ -74,7 +97,7 @@ pub async fn nfsboot(
     let tmp_dir_path = tmp_dir.path().to_path_buf();
     let rts_dir = tmp_dir.path().join("rts");
     let mounter = tokio::task::spawn_blocking(move || {
-        let mut mounter = Mounter::new(tmp_dir);
+        let mut mounter = Mounter::new(tmp_dir, mount_rts_tmp_dir);
         mounter
             .do_mounting(&rts_dir, &scratch_dir, &mounts)
             .map(|()| mounter)
@@ -108,14 +131,16 @@ pub async fn nfsboot(
 struct Mounter {
     mounts: Vec<PathBuf>,
     tmp: Option<TempDir>,
+    mount_rts_tmp: Option<TempDir>,
 }
 
 #[bon::bon]
 impl Mounter {
-    fn new(temp_dir: TempDir) -> Self {
+    fn new(temp_dir: TempDir, mount_rts_tmp: Option<TempDir>) -> Self {
         Self {
             mounts: Vec::new(),
             tmp: Some(temp_dir),
+            mount_rts_tmp,
         }
     }
 
@@ -157,7 +182,7 @@ impl Mounter {
             .iter()
             .map(|m| inner_mount_dir.join(&m.orb_mount_name))
         {
-            run_fun!(sudo mkdir $d)
+            run_fun!(sudo mkdir -p $d)
                 .wrap_err_with(|| format!("failed to create {d:?}"))?;
         }
 
@@ -243,15 +268,15 @@ impl Drop for Mounter {
             }
         }
 
-        // The regular destructor of TempDir doesn't work, because the directory contains
-        // root-owned files. We need to delete manually with sudo.
-        let tmp = self.tmp.take().expect("always Some until drop");
-        let tmp_path = tmp.path();
-        debug!("deleting tempdir {tmp_path:?}");
-        let result = run_fun!(sudo rm -rf $tmp_path)
-            .wrap_err("failed to remove tempdir with sudo");
-        if let Err(err) = result {
-            warn!("{err:?}");
+        for tmp in [self.mount_rts_tmp.take(), self.tmp.take()] {
+            let Some(tmp) = tmp else { continue };
+            let tmp_path = tmp.path();
+            debug!("deleting tempdir {tmp_path:?}");
+            let result = run_fun!(sudo rm -rf $tmp_path)
+                .wrap_err("failed to remove tempdir with sudo");
+            if let Err(err) = result {
+                warn!("{err:?}");
+            }
         }
     }
 }