Swapping get_mount_devices from while loop to filter_map (#103)

* Swapping get_mount_devices from while loop to filter_map

The while loop for get_mount_devices was stuck in a loop where it was  re-fetching the list of mounted devices every time if anything returned true and then would repeatedly find the same device. 

This should eliminate the issue and allow it to only add the device once.

* Adding unit test for getting mount media

* Write unit tests and mount_media with fstab instead of block_utils

Co-authored-by: Anh Vo <[email protected]>

Author: peytonr18

libazureinit/Cargo.toml

@@ -19,6 +19,7 @@ nix = {version = "0.29.0", features = ["fs", "user"]}
 block-utils = "0.11.1"
 tracing = "0.1.40"
 strum = { version = "0.26.3", features = ["derive"] }
+fstab = "0.4.0"
 
 [dev-dependencies]
 tempfile = "3"

libazureinit/src/media.rs

@@ -6,6 +6,7 @@ use std::fs::create_dir_all;
 use std::fs::File;
 use std::io::Read;
 use std::os::unix::fs::PermissionsExt;
+use std::path::Path;
 use std::path::PathBuf;
 use std::process::Command;
 
@@ -16,6 +17,7 @@ use tracing;
 use tracing::instrument;
 
 use crate::error::Error;
+use fstab::FsTab;
 
 #[derive(Debug, Default, Deserialize, PartialEq, Clone)]
 pub struct Environment {
@@ -75,20 +77,27 @@ pub const PATH_MOUNT_DEVICE: &str = "/dev/sr0";
 pub const PATH_MOUNT_POINT: &str = "/run/azure-init/media/";
 
 const CDROM_VALID_FS: &[&str] = &["iso9660", "udf"];
+const MTAB_PATH: &str = "/etc/mtab";
 
 // Get a mounted device with any filesystem for CDROM
 #[instrument]
-pub fn get_mount_device() -> Result<Vec<String>, Error> {
-    let mut list_devices: Vec<String> = Vec::new();
+pub fn get_mount_device(path: Option<&Path>) -> Result<Vec<String>, Error> {
+    let fstab = FsTab::new(path.unwrap_or_else(|| Path::new(MTAB_PATH)));
+    let entries = fstab.get_entries()?;
 
-    while let Some(device) = block_utils::get_mounted_devices()?
+    // Retrieve the names of all devices that have cdrom-type filesystem (e.g., udf)
+    let cdrom_devices = entries
         .into_iter()
-        .find(|dev| CDROM_VALID_FS.contains(&dev.fs_type.to_str()))
-    {
-        list_devices.push(device.name);
-    }
-
-    Ok(list_devices)
+        .filter_map(|entry| {
+            if CDROM_VALID_FS.contains(&entry.vfs_type.as_str()) {
+                Some(entry.fs_spec)
+            } else {
+                None
+            }
+        })
+        .collect();
+
+    Ok(cdrom_devices)
 }
 
 // Some zero-sized structs that just provide states for our state machine
@@ -222,6 +231,9 @@ pub fn mount_parse_ovf_env(dev: String) -> Result<Environment, Error> {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::error::Error;
+    use std::io::Write;
+    use tempfile::NamedTempFile;
 
     #[test]
     fn test_get_ovf_env_none_missing() {
@@ -406,4 +418,80 @@ mod tests {
             _ => panic!("Non-empty passwords aren't allowed"),
         };
     }
+
+    #[test]
+    fn test_get_mount_device_with_cdrom_entries() {
+        let mut temp_file =
+            NamedTempFile::new().expect("Failed to create temporary file");
+        let mount_table = r#"
+            /dev/sr0 /mnt/cdrom iso9660 ro,user,noauto 0 0
+            /dev/sr1 /mnt/cdrom2 udf ro,user,noauto 0 0
+        "#;
+        temp_file
+            .write_all(mount_table.as_bytes())
+            .expect("Failed to write to temporary file");
+        let temp_path = temp_file.into_temp_path();
+        let result = get_mount_device(Some(temp_path.as_ref()));
+
+        let list_devices = result.unwrap();
+        assert_eq!(
+            list_devices,
+            vec!["/dev/sr0".to_string(), "/dev/sr1".to_string()]
+        );
+    }
+
+    #[test]
+    fn test_get_mount_device_without_cdrom_entries() {
+        let mut temp_file =
+            NamedTempFile::new().expect("Failed to create temporary file");
+        let mount_table = r#"
+            /dev/sda1 / ext4 defaults 0 0
+            /dev/sda2 /home ext4 defaults 0 0
+        "#;
+        temp_file
+            .write_all(mount_table.as_bytes())
+            .expect("Failed to write to temporary file");
+        let temp_path = temp_file.into_temp_path();
+        let result = get_mount_device(Some(temp_path.as_ref()));
+
+        let list_devices = result.unwrap();
+        assert!(list_devices.is_empty());
+    }
+
+    #[test]
+    fn test_get_mount_device_with_mixed_entries() {
+        let mut temp_file =
+            NamedTempFile::new().expect("Failed to create temporary file");
+        let mount_table = r#"
+            /dev/sr0 /mnt/cdrom iso9660 ro,user,noauto 0 0
+            /dev/sda1 / ext4 defaults 0 0
+            /dev/sr1 /mnt/cdrom2 udf ro,user,noauto 0 0
+        "#;
+        temp_file
+            .write_all(mount_table.as_bytes())
+            .expect("Failed to write to temporary file");
+        let temp_path = temp_file.into_temp_path();
+        let result = get_mount_device(Some(temp_path.as_ref()));
+
+        let list_devices = result.unwrap();
+        assert_eq!(
+            list_devices,
+            vec!["/dev/sr0".to_string(), "/dev/sr1".to_string()]
+        );
+    }
+
+    #[test]
+    fn test_get_mount_device_with_empty_table() {
+        let mut temp_file =
+            NamedTempFile::new().expect("Failed to create temporary file");
+        let mount_table = "";
+        temp_file
+            .write_all(mount_table.as_bytes())
+            .expect("Failed to write to temporary file");
+        let temp_path = temp_file.into_temp_path();
+        let result = get_mount_device(Some(temp_path.as_ref()));
+
+        let list_devices = result.unwrap();
+        assert!(list_devices.is_empty());
+    }
 }

src/main.rs

@@ -4,13 +4,12 @@
 use std::process::ExitCode;
 
 use anyhow::Context;
-
 use libazureinit::imds::InstanceMetadata;
 use libazureinit::User;
 use libazureinit::{
     error::Error as LibError,
     goalstate, imds, media,
-    media::Environment,
+    media::{get_mount_device, Environment},
     reqwest::{header, Client},
     HostnameProvisioner, PasswordProvisioner, Provision, UserProvisioner,
 };
@@ -23,7 +22,7 @@ const VERSION: &str = env!("CARGO_PKG_VERSION");
 
 #[instrument]
 fn get_environment() -> Result<Environment, anyhow::Error> {
-    let ovf_devices = media::get_mount_device()?;
+    let ovf_devices = get_mount_device(None)?;
     let mut environment: Option<Environment> = None;
 
     // loop until it finds a correct device.