feat: mount and umount disk image (#78)

Signed-off-by: longjin <longjin@dragonos.org>

dadk-config/src/rootfs/partition.rs

@@ -21,10 +21,10 @@ pub enum PartitionType {
 }
 
 impl PartitionConfig {
-    /// Determines whether a partitioned image should be created.
+    /// Determines whether the disk image should be partitioned
     ///
     /// Returns `true` if the partition type is not `None`, otherwise returns `false`.
-    pub fn should_create_partitioned_image(&self) -> bool {
+    pub fn image_should_be_partitioned(&self) -> bool {
         self.partition_type != PartitionType::None
     }
 }

dadk/Cargo.toml

@@ -35,6 +35,7 @@ dadk-user = { path = "../dadk-user" }
 derive_builder = "0.20.0"
 env_logger = "0.11.5"
 log = "0.4.22"
+regex = "1.9.1"
 
 [dev-dependencies]
 tempfile = "3.13.0"

dadk/src/actions/rootfs/disk_img.rs

@@ -1,4 +1,4 @@
-use std::{fs::File, io::Write, path::PathBuf, process::Command};
+use std::{fs::File, io::Write, mem::ManuallyDrop, path::PathBuf, process::Command};
 
 use crate::context::DADKExecContext;
 use anyhow::{anyhow, Result};
@@ -22,7 +22,7 @@ pub(super) fn create(ctx: &DADKExecContext) -> Result<()> {
 
     // 判断是否需要分区？
 
-    let r = if ctx.rootfs().partition.should_create_partitioned_image() {
+    let r = if ctx.rootfs().partition.image_should_be_partitioned() {
         create_partitioned_image(ctx, &disk_image_path)
     } else {
         create_unpartitioned_image(ctx, &disk_image_path)
@@ -33,6 +33,130 @@ pub(super) fn create(ctx: &DADKExecContext) -> Result<()> {
     }
     r
 }
+
+pub fn mount(ctx: &DADKExecContext) -> Result<()> {
+    let disk_image_path = ctx.disk_image_path();
+    if !disk_image_path.exists() {
+        return Err(anyhow!(
+            "Disk image does not exist: {}",
+            disk_image_path.display()
+        ));
+    }
+    let disk_mount_path = ctx.disk_mount_path();
+
+    // 尝试创建挂载点
+    std::fs::create_dir_all(&disk_mount_path)
+        .map_err(|e| anyhow!("Failed to create disk mount path: {}", e))?;
+
+    let partitioned = ctx.rootfs().partition.image_should_be_partitioned();
+    log::trace!("Disk image is partitioned: {}", partitioned);
+    if partitioned {
+        mount_partitioned_image(ctx, &disk_image_path, &disk_mount_path)?
+    } else {
+        mount_unpartitioned_image(ctx, &disk_image_path, &disk_mount_path)?
+    }
+    log::info!("Disk image mounted at {}", disk_mount_path.display());
+    Ok(())
+}
+
+fn mount_partitioned_image(
+    ctx: &DADKExecContext,
+    disk_image_path: &PathBuf,
+    disk_mount_path: &PathBuf,
+) -> Result<()> {
+    let mut loop_device = ManuallyDrop::new(
+        LoopDeviceBuilder::new()
+            .img_path(disk_image_path.clone())
+            .build()
+            .map_err(|e| anyhow!("Failed to create loop device: {}", e))?,
+    );
+
+    loop_device
+        .attach()
+        .map_err(|e| anyhow!("Failed to attach loop device: {}", e))?;
+
+    let dev_path = loop_device.partition_path(1)?;
+    mount_unpartitioned_image(ctx, &dev_path, disk_mount_path)?;
+
+    Ok(())
+}
+
+fn mount_unpartitioned_image(
+    _ctx: &DADKExecContext,
+    disk_image_path: &PathBuf,
+    disk_mount_path: &PathBuf,
+) -> Result<()> {
+    let cmd = Command::new("mount")
+        .arg(disk_image_path)
+        .arg(disk_mount_path)
+        .output()
+        .map_err(|e| anyhow!("Failed to mount disk image: {}", e))?;
+    if !cmd.status.success() {
+        return Err(anyhow!(
+            "Failed to mount disk image: {}",
+            String::from_utf8_lossy(&cmd.stderr)
+        ));
+    }
+    Ok(())
+}
+
+pub fn umount(ctx: &DADKExecContext) -> Result<()> {
+    let disk_img_path = ctx.disk_image_path();
+    let disk_mount_path = ctx.disk_mount_path();
+    let mut loop_device = LoopDeviceBuilder::new().img_path(disk_img_path).build();
+
+    let should_detach_loop_device: bool;
+    if let Ok(loop_device) = loop_device.as_mut() {
+        if let Err(e) = loop_device.attach_by_exists() {
+            log::trace!("umount: Failed to attach loop device: {}", e);
+        }
+
+        should_detach_loop_device = loop_device.attached();
+    } else {
+        should_detach_loop_device = false;
+    }
+
+    if disk_mount_path.exists() {
+        let cmd = Command::new("umount")
+            .arg(disk_mount_path)
+            .output()
+            .map_err(|e| anyhow!("Failed to umount disk image: {}", e));
+        match cmd {
+            Ok(cmd) => {
+                if !cmd.status.success() {
+                    let e = anyhow!(
+                        "Failed to umount disk image: {}",
+                        String::from_utf8_lossy(&cmd.stderr)
+                    );
+                    if should_detach_loop_device {
+                        log::error!("{}", e);
+                    } else {
+                        return Err(e);
+                    }
+                }
+            }
+            Err(e) => {
+                if should_detach_loop_device {
+                    log::error!("{}", e);
+                } else {
+                    return Err(e);
+                }
+            }
+        }
+    }
+
+    if let Ok(mut loop_device) = loop_device {
+        let loop_dev_path = loop_device.dev_path().cloned();
+        if let Err(e) = loop_device.detach().map_err(|e| anyhow!("{}", e)) {
+            if ctx.rootfs().partition.image_should_be_partitioned() {
+                return Err(e);
+            }
+        }
+        log::info!("Loop device detached: {:?}", loop_dev_path);
+    }
+    Ok(())
+}
+
 /// Ensures the provided disk image path is not a device node.
 fn disk_path_safety_check(disk_image_path: &PathBuf) -> Result<()> {
     const DONT_ALLOWED_PREFIX: [&str; 5] =
@@ -57,6 +181,9 @@ fn create_partitioned_image(ctx: &DADKExecContext, disk_image_path: &PathBuf) ->
         .img_path(disk_image_path.clone())
         .build()
         .map_err(|e| anyhow!("Failed to create loop device: {}", e))?;
+    loop_device
+        .attach()
+        .map_err(|e| anyhow!("Failed to attach loop device: {}", e))?;
 
     let partition_path = loop_device.partition_path(1)?;
     let fs_type = ctx.rootfs().metadata.fs_type;

dadk/src/actions/rootfs/loopdev.rs

@@ -1,21 +1,39 @@
+use core::str;
 use std::{path::PathBuf, process::Command};
 
 use anyhow::{anyhow, Result};
+use regex::Regex;
+
+use crate::utils::abs_path;
+
+const LOOP_DEVICE_LOSETUP_A_REGEX: &str = r"^/dev/loop(\d+)";
 
 pub struct LoopDevice {
-    img_path: PathBuf,
+    img_path: Option<PathBuf>,
     loop_device_path: Option<String>,
 }
 impl LoopDevice {
+    pub fn attached(&self) -> bool {
+        self.loop_device_path.is_some()
+    }
+
+    pub fn dev_path(&self) -> Option<&String> {
+        self.loop_device_path.as_ref()
+    }
+
     pub fn attach(&mut self) -> Result<()> {
-        if self.loop_device_path.is_some() {
+        if self.attached() {
             return Ok(());
         }
+        if self.img_path.is_none() {
+            return Err(anyhow!("Image path not set"));
+        }
+
         let output = Command::new("losetup")
             .arg("-f")
             .arg("--show")
             .arg("-P")
-            .arg(&self.img_path)
+            .arg(self.img_path.as_ref().unwrap())
             .output()?;
 
         if output.status.success() {
@@ -33,6 +51,34 @@ impl LoopDevice {
             ))
         }
     }
+
+    /// 尝试连接已经存在的loop device
+    pub fn attach_by_exists(&mut self) -> Result<()> {
+        if self.attached() {
+            return Ok(());
+        }
+        if self.img_path.is_none() {
+            return Err(anyhow!("Image path not set"));
+        }
+        log::trace!(
+            "Try to attach loop device by exists: image path: {}",
+            self.img_path.as_ref().unwrap().display()
+        );
+        // losetup -a 查看是否有已经attach了的，如果有，就附着上去
+        let cmd = Command::new("losetup")
+            .arg("-a")
+            .output()
+            .map_err(|e| anyhow!("Failed to run losetup -a: {}", e))?;
+        let output = String::from_utf8(cmd.stdout)?;
+        let s = __loop_device_path_by_disk_image_path(
+            self.img_path.as_ref().unwrap().to_str().unwrap(),
+            &output,
+        )
+        .map_err(|e| anyhow!("Failed to find loop device: {}", e))?;
+        self.loop_device_path = Some(s);
+        Ok(())
+    }
+
     /// 获取指定分区的路径
     ///
     /// # 参数
@@ -47,7 +93,7 @@ impl LoopDevice {
     ///
     /// 如果循环设备未附加，则返回 `anyhow!("Loop device not attached")` 错误。
     pub fn partition_path(&self, nth: u8) -> Result<PathBuf> {
-        if self.loop_device_path.is_none() {
+        if !self.attached() {
             return Err(anyhow!("Loop device not attached"));
         }
         let s = format!("{}p{}", self.loop_device_path.as_ref().unwrap(), nth);
@@ -73,37 +119,118 @@ impl LoopDevice {
             self.loop_device_path = None;
             Ok(())
         } else {
-            Err(anyhow::anyhow!("Failed to detach loop device"))
+            Err(anyhow::anyhow!(
+                "Failed to detach loop device: {}, {}",
+                output.status,
+                str::from_utf8(output.stderr.as_slice()).unwrap_or("<Unknown>")
+            ))
         }
     }
 }
 
 impl Drop for LoopDevice {
     fn drop(&mut self) {
-        self.detach().expect("Failed to detach loop device");
+        if let Err(e) = self.detach() {
+            log::warn!("Failed to detach loop device: {}", e);
+        }
     }
 }
 
 pub struct LoopDeviceBuilder {
     img_path: Option<PathBuf>,
+    loop_device_path: Option<String>,
 }
 
 impl LoopDeviceBuilder {
     pub fn new() -> Self {
-        LoopDeviceBuilder { img_path: None }
+        LoopDeviceBuilder {
+            img_path: None,
+            loop_device_path: None,
+        }
     }
 
     pub fn img_path(mut self, img_path: PathBuf) -> Self {
-        self.img_path = Some(img_path);
+        self.img_path = Some(abs_path(&img_path));
         self
     }
 
     pub fn build(self) -> Result<LoopDevice> {
-        let mut loop_dev = LoopDevice {
-            img_path: self.img_path.unwrap(),
-            loop_device_path: None,
+        let loop_dev = LoopDevice {
+            img_path: self.img_path,
+            loop_device_path: self.loop_device_path,
         };
-        loop_dev.attach()?;
+
         Ok(loop_dev)
     }
 }
+
+fn __loop_device_path_by_disk_image_path(
+    disk_img_path: &str,
+    losetup_a_output: &str,
+) -> Result<String> {
+    let re = Regex::new(LOOP_DEVICE_LOSETUP_A_REGEX)?;
+    for line in losetup_a_output.lines() {
+        if !line.contains(disk_img_path) {
+            continue;
+        }
+        let caps = re.captures(line);
+        if caps.is_none() {
+            continue;
+        }
+        let caps = caps.unwrap();
+        let loop_device = caps.get(1).unwrap().as_str();
+        let loop_device = format!("/dev/loop{}", loop_device);
+        return Ok(loop_device);
+    }
+    Err(anyhow!("Loop device not found"))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_regex_find_loop_device() {
+        const DEVICE_NAME_SHOULD_MATCH: [&str; 3] =
+            ["/dev/loop11", "/dev/loop11p1", "/dev/loop11p1 "];
+        let device_name = "/dev/loop11";
+        let re = Regex::new(LOOP_DEVICE_LOSETUP_A_REGEX).unwrap();
+        for name in DEVICE_NAME_SHOULD_MATCH {
+            assert!(re.find(name).is_some(), "{} should match", name);
+            assert_eq!(
+                re.find(name).unwrap().as_str(),
+                device_name,
+                "{} should match {}",
+                name,
+                device_name
+            );
+        }
+    }
+
+    #[test]
+    fn test_parse_losetup_a_output() {
+        let losetup_a_output = r#"/dev/loop1: []: (/data/bin/disk-image-x86_64.img)
+/dev/loop29: []: (/var/lib/abc.img)
+/dev/loop13: []: (/var/lib/snapd/snaps/gtk-common-themes_1535.snap
+/dev/loop19: []: (/var/lib/snapd/snaps/gnome-42-2204_172.snap)"#;
+        let disk_img_path = "/data/bin/disk-image-x86_64.img";
+        let loop_device_path =
+            __loop_device_path_by_disk_image_path(disk_img_path, losetup_a_output).unwrap();
+        assert_eq!(loop_device_path, "/dev/loop1");
+    }
+
+    #[test]
+    fn test_parse_lsblk_output_not_match() {
+        let losetup_a_output = r#"/dev/loop1: []: (/data/bin/disk-image-x86_64.img)
+/dev/loop29: []: (/var/lib/abc.img)
+/dev/loop13: []: (/var/lib/snapd/snaps/gtk-common-themes_1535.snap
+/dev/loop19: []: (/var/lib/snapd/snaps/gnome-42-2204_172.snap)"#;
+        let disk_img_path = "/data/bin/disk-image-riscv64.img";
+        let loop_device_path =
+            __loop_device_path_by_disk_image_path(disk_img_path, losetup_a_output);
+        assert!(
+            loop_device_path.is_err(),
+            "should not match any loop device"
+        );
+    }
+}

dadk/src/actions/rootfs/mod.rs

@@ -9,7 +9,7 @@ pub(super) fn run(ctx: &DADKExecContext, rootfs_cmd: &RootFSCommand) -> Result<(
         RootFSCommand::Create => disk_img::create(ctx),
         RootFSCommand::Delete => todo!(),
         RootFSCommand::DeleteSysroot => todo!(),
-        RootFSCommand::Mount => todo!(),
-        RootFSCommand::Unmount => todo!(),
+        RootFSCommand::Mount => disk_img::mount(ctx),
+        RootFSCommand::Umount => disk_img::umount(ctx),
     }
 }

dadk/src/console/rootfs.rs

@@ -12,5 +12,5 @@ pub enum RootFSCommand {
     /// 挂载根文件系统（磁盘镜像）
     Mount,
     /// 卸载根文件系统（磁盘镜像）
-    Unmount,
+    Umount,
 }

dadk/src/context/mod.rs

@@ -86,8 +86,19 @@ impl DADKExecContext {
 
     /// 获取磁盘镜像的路径，路径由工作目录、架构和固定文件名组成
     pub fn disk_image_path(&self) -> PathBuf {
+        self.workdir()
+            .join(format!("bin/{}.img", self.disk_image_basename()))
+    }
+
+    /// 获取磁盘挂载路径
+    pub fn disk_mount_path(&self) -> PathBuf {
+        self.workdir()
+            .join(format!("bin/mnt/{}", self.disk_image_basename()))
+    }
+
+    fn disk_image_basename(&self) -> String {
         let arch: String = self.target_arch().into();
-        self.workdir().join(format!("bin/disk-image-{}.img", arch))
+        format!("disk-image-{}", arch)
     }
 
     /// 获取磁盘镜像大小

dadk/src/utils.rs

@@ -19,6 +19,7 @@ pub fn abs_path(path: &PathBuf) -> PathBuf {
     if path.is_absolute() {
         path.to_path_buf()
     } else {
-        std::env::current_dir().unwrap().join(path)
+        let origin = std::env::current_dir().unwrap().join(path);
+        origin.canonicalize().unwrap_or(origin)
     }
 }