One-KVM/src/otg/msd.rs

use std::fs;
use std::path::{Path, PathBuf};
use tracing::{debug, info, warn};

use super::configfs::{create_dir, create_symlink, remove_dir, remove_file, write_file};
use super::function::GadgetFunction;
use crate::error::{AppError, Result};

#[derive(Debug, Clone)]
pub struct MsdLunConfig {
    pub file: PathBuf,
    pub cdrom: bool,
    pub ro: bool,
    pub removable: bool,
    pub nofua: bool,
}

impl Default for MsdLunConfig {
    fn default() -> Self {
        Self {
            file: PathBuf::new(),
            cdrom: false,
            ro: false,
            removable: true,
            nofua: true,
        }
    }
}

impl MsdLunConfig {
    pub fn cdrom(file: PathBuf) -> Self {
        Self {
            file,
            cdrom: true,
            ro: true,
            removable: true,
            nofua: true,
        }
    }

    pub fn disk(file: PathBuf, read_only: bool) -> Self {
        Self {
            file,
            cdrom: false,
            ro: read_only,
            removable: true,
            // nofua=false: enforce Force Unit Access so the USB host (e.g. Windows)
            // gets proper write-completion acknowledgements when writing to the
            // virtual .img file. nofua=true can cause write-verify failures
            // that manifest as Windows error 0x80070570 on writable drives.
            nofua: false,
        }
    }
}

#[derive(Debug, Clone)]
pub struct MsdFunction {
    name: String,
    lun_capacity: u8,
}

impl MsdFunction {
    pub fn new(instance: u8, lun_capacity: u8) -> Result<Self> {
        if lun_capacity != 1 && lun_capacity != 8 {
            return Err(AppError::BadRequest(format!(
                "MSD LUN capacity must be 1 or 8, got {lun_capacity}"
            )));
        }

        Ok(Self {
            name: format!("mass_storage.usb{}", instance),
            lun_capacity,
        })
    }

    fn function_path(&self, gadget_path: &Path) -> PathBuf {
        gadget_path.join("functions").join(self.name())
    }

    fn lun_path(&self, gadget_path: &Path, lun: u8) -> PathBuf {
        self.function_path(gadget_path).join(format!("lun.{}", lun))
    }

    fn existing_lun_paths(&self, gadget_path: &Path) -> Result<Vec<(u16, PathBuf)>> {
        let func_path = self.function_path(gadget_path);
        if !func_path.exists() {
            return Ok(Vec::new());
        }

        let entries = fs::read_dir(&func_path).map_err(|e| {
            AppError::Internal(format!(
                "Failed to read MSD function directory {}: {}",
                func_path.display(),
                e
            ))
        })?;
        let mut luns = entries
            .filter_map(|entry| {
                let entry = entry.ok()?;
                let name = entry.file_name();
                let name = name.to_str()?;
                let lun = name.strip_prefix("lun.")?.parse::<u16>().ok()?;
                Some((lun, entry.path()))
            })
            .collect::<Vec<_>>();
        luns.sort_by_key(|(lun, _)| *lun);
        Ok(luns)
    }

    pub async fn configure_lun_async(
        &self,
        gadget_path: &Path,
        lun: u8,
        config: &MsdLunConfig,
    ) -> Result<()> {
        let gadget_path = gadget_path.to_path_buf();
        let config = config.clone();
        let this = self.clone();

        tokio::task::spawn_blocking(move || this.configure_lun(&gadget_path, lun, &config))
            .await
            .map_err(|e| AppError::Internal(format!("Task join error: {}", e)))?
    }

    fn clear_lun_unbound(&self, gadget_path: &Path, lun: u8) -> Result<()> {
        let lun_path = self.lun_path(gadget_path, lun);
        if !lun_path.exists() {
            create_dir(&lun_path)?;
        }
        write_file(&lun_path.join("file"), "")?;
        let _ = write_file(&lun_path.join("cdrom"), "0");
        let _ = write_file(&lun_path.join("ro"), "0");
        let _ = write_file(&lun_path.join("removable"), "1");
        let _ = write_file(&lun_path.join("nofua"), "1");
        Ok(())
    }

    pub fn configure_lun(&self, gadget_path: &Path, lun: u8, config: &MsdLunConfig) -> Result<()> {
        if lun >= self.lun_capacity {
            return Err(AppError::BadRequest(format!(
                "LUN {lun} is outside MSD capacity {}",
                self.lun_capacity
            )));
        }
        let lun_path = self.lun_path(gadget_path, lun);

        if !lun_path.exists() {
            return Err(AppError::Internal(format!(
                "Configured MSD LUN {lun} does not exist"
            )));
        }

        let read_attr = |attr: &str| -> String {
            fs::read_to_string(lun_path.join(attr))
                .unwrap_or_default()
                .trim()
                .to_string()
        };

        let current_cdrom = read_attr("cdrom");
        let current_ro = read_attr("ro");
        let current_removable = read_attr("removable");
        let current_nofua = read_attr("nofua");

        let new_cdrom = if config.cdrom { "1" } else { "0" };
        let new_ro = if config.ro { "1" } else { "0" };
        let new_removable = if config.removable { "1" } else { "0" };
        let new_nofua = if config.nofua { "1" } else { "0" };

        let forced_eject_path = lun_path.join("forced_eject");
        if forced_eject_path.exists() {
            debug!("Using forced_eject to clear LUN {}", lun);
            let _ = write_file(&forced_eject_path, "1");
        } else {
            let _ = write_file(&lun_path.join("file"), "");
        }

        std::thread::sleep(std::time::Duration::from_millis(50));

        let cdrom_changed = current_cdrom != new_cdrom;
        if cdrom_changed {
            debug!(
                "Updating LUN {} cdrom: {} -> {}",
                lun, current_cdrom, new_cdrom
            );
            write_file(&lun_path.join("cdrom"), new_cdrom)?;
        }
        if current_ro != new_ro {
            debug!("Updating LUN {} ro: {} -> {}", lun, current_ro, new_ro);
            write_file(&lun_path.join("ro"), new_ro)?;
        }
        if current_removable != new_removable {
            debug!(
                "Updating LUN {} removable: {} -> {}",
                lun, current_removable, new_removable
            );
            write_file(&lun_path.join("removable"), new_removable)?;
        }
        if current_nofua != new_nofua {
            debug!(
                "Updating LUN {} nofua: {} -> {}",
                lun, current_nofua, new_nofua
            );
            write_file(&lun_path.join("nofua"), new_nofua)?;
        }

        if cdrom_changed {
            debug!("CDROM mode changed, brief yield for USB host");
            std::thread::sleep(std::time::Duration::from_millis(50));
        }

        if config.file.exists() {
            let file_path = config.file.to_string_lossy();
            let mut last_error = None;

            for attempt in 0..5 {
                match write_file(&lun_path.join("file"), file_path.as_ref()) {
                    Ok(_) => {
                        info!(
                            "LUN {} configured with file: {} (cdrom={}, ro={})",
                            lun,
                            config.file.display(),
                            config.cdrom,
                            config.ro
                        );
                        return Ok(());
                    }
                    Err(e) => {
                        let is_busy = e.to_string().contains("Device or resource busy")
                            || e.to_string().contains("os error 16");

                        if is_busy && attempt < 4 {
                            warn!(
                                "LUN {} file write busy, retrying (attempt {}/5)",
                                lun,
                                attempt + 1
                            );
                            std::thread::sleep(std::time::Duration::from_millis(50 << attempt));
                            last_error = Some(e);
                            continue;
                        }

                        return Err(e);
                    }
                }
            }

            if let Some(e) = last_error {
                return Err(e);
            }
        } else if !config.file.as_os_str().is_empty() {
            warn!("LUN {} file does not exist: {}", lun, config.file.display());
        }

        Ok(())
    }

    pub async fn disconnect_lun_async(&self, gadget_path: &Path, lun: u8) -> Result<()> {
        let gadget_path = gadget_path.to_path_buf();
        let this = self.clone();

        tokio::task::spawn_blocking(move || this.disconnect_lun(&gadget_path, lun))
            .await
            .map_err(|e| AppError::Internal(format!("Task join error: {}", e)))?
    }

    pub fn disconnect_lun(&self, gadget_path: &Path, lun: u8) -> Result<()> {
        if lun >= self.lun_capacity {
            return Err(AppError::BadRequest(format!(
                "LUN {lun} is outside MSD capacity {}",
                self.lun_capacity
            )));
        }
        let lun_path = self.lun_path(gadget_path, lun);

        self.disconnect_lun_path(&lun_path, lun as u16)
    }

    fn disconnect_lun_path(&self, lun_path: &Path, lun: u16) -> Result<()> {
        if lun_path.exists() {
            let forced_eject_path = lun_path.join("forced_eject");
            if forced_eject_path.exists() {
                debug!(
                    "Using forced_eject to disconnect LUN {} at {:?}",
                    lun, forced_eject_path
                );
                match write_file(&forced_eject_path, "1") {
                    Ok(_) => debug!("forced_eject write succeeded"),
                    Err(e) => {
                        warn!(
                            "forced_eject write failed: {}, falling back to clearing file",
                            e
                        );
                        let file_path = lun_path.join("file");
                        if file_path.exists() {
                            write_file(&file_path, "")?;
                        }
                    }
                }
            } else {
                let file_path = lun_path.join("file");
                if file_path.exists() {
                    write_file(&file_path, "")?;
                }
            }
            info!("LUN {} disconnected", lun);
        }

        Ok(())
    }

    pub fn get_lun_file(&self, gadget_path: &Path, lun: u8) -> Option<PathBuf> {
        let lun_path = self.lun_path(gadget_path, lun);
        let file_path = lun_path.join("file");

        if let Ok(content) = fs::read_to_string(&file_path) {
            let content = content.trim();
            if !content.is_empty() {
                return Some(PathBuf::from(content));
            }
        }

        None
    }

    pub fn is_lun_connected(&self, gadget_path: &Path, lun: u8) -> bool {
        self.get_lun_file(gadget_path, lun).is_some()
    }
}

impl GadgetFunction for MsdFunction {
    fn name(&self) -> &str {
        &self.name
    }

    fn create(&self, gadget_path: &Path) -> Result<()> {
        let func_path = self.function_path(gadget_path);
        create_dir(&func_path)?;

        let stall_path = func_path.join("stall");
        if stall_path.exists() {
            let _ = write_file(&stall_path, "0");
        }

        for lun in 0..self.lun_capacity {
            self.clear_lun_unbound(gadget_path, lun)?;
        }

        debug!("Created MSD function: {}", self.name());
        Ok(())
    }

    fn link(&self, config_path: &Path, gadget_path: &Path) -> Result<()> {
        let func_path = self.function_path(gadget_path);
        let link_path = config_path.join(self.name());

        if !link_path.exists() {
            create_symlink(&func_path, &link_path)?;
            debug!("Linked MSD function {} to config", self.name());
        }

        Ok(())
    }

    fn unlink(&self, config_path: &Path) -> Result<()> {
        let link_path = config_path.join(self.name());
        remove_file(&link_path)?;
        debug!("Unlinked MSD function {}", self.name());
        Ok(())
    }

    fn cleanup(&self, gadget_path: &Path) -> Result<()> {
        let func_path = self.function_path(gadget_path);
        let mut errors = Vec::new();

        let lun_paths = match self.existing_lun_paths(gadget_path) {
            Ok(luns) => luns,
            Err(e) => {
                errors.push(format!("could not enumerate MSD LUN directories: {e}"));
                Vec::new()
            }
        };
        for (lun, lun_path) in lun_paths {
            if let Err(e) = self.disconnect_lun_path(&lun_path, lun) {
                errors.push(format!("could not disconnect LUN {lun}: {e}"));
            }
            // lun.0 is the mass-storage function's configfs default group. It
            // cannot be removed directly and is released with the function.
            if lun == 0 {
                continue;
            }
            if let Err(e) = remove_dir(&lun_path) {
                errors.push(format!("could not remove LUN {lun} directory: {e}"));
            }
        }

        if let Err(e) = remove_dir(&func_path) {
            errors.push(format!("could not remove MSD function directory: {e}"));
        }

        if !errors.is_empty() {
            return Err(AppError::Config(format!(
                "MSD cleanup incomplete: {}",
                errors.join("; ")
            )));
        }

        debug!("Cleaned up MSD function {}", self.name());
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::TempDir;

    #[test]
    fn test_lun_config_cdrom() {
        let config = MsdLunConfig::cdrom(PathBuf::from("/tmp/test.iso"));
        assert!(config.cdrom);
        assert!(config.ro);
        assert!(config.removable);
    }

    #[test]
    fn test_lun_config_disk() {
        let config = MsdLunConfig::disk(PathBuf::from("/tmp/test.img"), false);
        assert!(!config.cdrom);
        assert!(!config.ro);
        assert!(config.removable);
    }

    #[test]
    fn test_msd_function_name() {
        let msd = MsdFunction::new(0, 1).unwrap();
        assert_eq!(msd.name(), "mass_storage.usb0");
        assert_eq!(msd.lun_capacity, 1);

        let multi = MsdFunction::new(0, 8).unwrap();
        assert_eq!(multi.lun_capacity, 8);
    }

    #[test]
    fn test_msd_function_rejects_invalid_capacity() {
        assert!(MsdFunction::new(0, 0).is_err());
        assert!(MsdFunction::new(0, 2).is_err());
        assert!(MsdFunction::new(0, 9).is_err());
    }

    #[test]
    fn create_uses_configured_lun_capacity() {
        for capacity in [1, 8] {
            let temp_dir = TempDir::new().unwrap();
            std::fs::create_dir_all(temp_dir.path().join("functions")).unwrap();
            let msd = MsdFunction::new(0, capacity).unwrap();

            msd.create(temp_dir.path()).unwrap();

            for lun in 0..capacity {
                assert!(msd.lun_path(temp_dir.path(), lun).exists());
            }
            assert!(!msd.lun_path(temp_dir.path(), capacity).exists());
        }
    }

    #[test]
    fn configure_lun_does_not_rebind_udc() {
        let temp_dir = TempDir::new().unwrap();
        let lun_path = temp_dir.path().join("functions/mass_storage.usb0/lun.0");
        std::fs::create_dir_all(&lun_path).unwrap();
        for attr in ["file", "cdrom", "ro", "removable", "nofua"] {
            std::fs::write(lun_path.join(attr), b"0\n").unwrap();
        }
        std::fs::write(temp_dir.path().join("UDC"), b"test.udc\n").unwrap();
        let image_path = temp_dir.path().join("test.img");
        std::fs::write(&image_path, b"image").unwrap();
        let msd = MsdFunction::new(0, 1).unwrap();

        msd.configure_lun(temp_dir.path(), 0, &MsdLunConfig::disk(image_path, false))
            .unwrap();

        assert_eq!(
            std::fs::read_to_string(temp_dir.path().join("UDC")).unwrap(),
            "test.udc\n"
        );
    }

    #[test]
    fn cleanup_removes_all_dynamic_luns_including_stale_capacity() {
        let temp_dir = TempDir::new().unwrap();
        let func_path = temp_dir.path().join("functions/mass_storage.usb0");
        for lun in 1..8 {
            std::fs::create_dir_all(func_path.join(format!("lun.{lun}"))).unwrap();
        }
        let msd = MsdFunction::new(0, 1).unwrap();

        msd.cleanup(temp_dir.path()).unwrap();

        assert!(!func_path.exists());
    }

    #[test]
    fn cleanup_reports_when_non_configfs_cannot_release_default_lun() {
        let temp_dir = TempDir::new().unwrap();
        let func_path = temp_dir.path().join("functions/mass_storage.usb0");
        for lun in 0..2 {
            std::fs::create_dir_all(func_path.join(format!("lun.{lun}"))).unwrap();
        }
        let msd = MsdFunction::new(0, 1).unwrap();

        let error = msd.cleanup(temp_dir.path()).unwrap_err();

        assert!(error.to_string().contains("MSD cleanup incomplete"));
        assert!(func_path.join("lun.0").exists());
        assert!(!func_path.join("lun.1").exists());
    }
}