sd-switch 0.6.3

use std::fs::{read_dir, File};
use std::io::Read;
use std::path::{Path, PathBuf};

use crate::error::Error;
pub use crate::systemd::ini::UnitSwitchMethod;
use crate::systemd::ini::{self, SystemdIni};

/// A systemd unit file path and content.
pub struct UnitFile {
    /// Path to the unit file.
    path: PathBuf,
    /// Unit configuration with overrides applied.
    config: SystemdIni,
}

#[derive(Debug, PartialEq)]
pub enum UnitType {
    Automount,
    Device,
    Mount,
    Path,
    Scope,
    Service,
    Slice,
    Socket,
    Swap,
    Target,
    Timer,
}

impl UnitFile {
    pub fn load(path: &Path) -> Result<Self, Error> {
        // Validate that the path has an OK file extension.
        unit_path_to_unit_type(path)?;

        let mut ini_content = String::new();
        File::open(path)?.read_to_string(&mut ini_content)?;

        let mut config = ini::parse(&ini_content).map_err(|err| Error::from_ini(path, err))?;

        if let Some(override_files) = collect_overrides(path)? {
            for file in override_files {
                let mut ini_content = String::new();
                File::open(&file)?.read_to_string(&mut ini_content)?;
                let overrides =
                    ini::parse(&ini_content).map_err(|err| Error::from_ini(&file, err))?;
                config.extend(overrides);
            }
        }

        Ok(Self {
            path: path.to_owned(),
            config,
        })
    }

    pub fn path(&self) -> &Path {
        &self.path
    }

    pub fn unit_type(&self) -> UnitType {
        // The error case should not occur since the same call succeeded in the
        // `load` constructor.
        unit_path_to_unit_type(self.path()).expect("Unexpected error")
    }

    fn refuse_manual_start(&self) -> bool {
        self.get_bool_option("Unit", ini::KEY_REFUSEMANUALSTART)
            .unwrap_or(false)
    }

    /// Determine how the unit should be switched.
    ///
    /// If the systemd unit option `RefuseManualStart` is set, then
    /// we only attempt to stop the old unit without starting the new one.
    ///
    /// The switch method can be set in the unit file, or the corresponding
    /// `overrides.conf` file, using
    ///
    /// ```ini
    /// [Unit]
    /// X-SwitchMethod = A
    /// ```
    ///
    /// where `A` is one of `reload`, `restart`, `stop-start`, and `keep-old`.
    /// See `UnitSwitchMethod` for a description of the different switch
    /// methods.
    ///
    /// If `X-SwitchMethod=` is not present then the legacy `Unit` options
    /// `X-ReloadIfChanged=`, `X-RestartIfChanged=`, and `X-StopIfChanged` will
    /// be used.
    ///
    /// If none of the above options are present then the switch method
    /// `UnitSwitchMethod::StopStart` is returned.
    pub fn switch_method(&self) -> UnitSwitchMethod {
        if let Some(method) = self.switch_method_from_refuse_manual_start() {
            method
        } else if let Some(method) = self.switch_method_from_xswitchmethod() {
            method
        } else if let Some(method) = self.switch_method_from_legacy() {
            method
        } else {
            UnitSwitchMethod::StopStart
        }
    }

    fn switch_method_from_xswitchmethod(&self) -> Option<UnitSwitchMethod> {
        self.config.get_unit_switch_method()
    }

    /// Get switch method based on unit's `RefuseManualStart` value. We don't consider `RefuseManualStop` here since it does not affect a new unit.
    fn switch_method_from_refuse_manual_start(&self) -> Option<UnitSwitchMethod> {
        if self.refuse_manual_start() {
            Some(UnitSwitchMethod::StopOnly)
        } else {
            None
        }
    }

    fn switch_method_from_legacy(&self) -> Option<UnitSwitchMethod> {
        // Check legacy options.
        let reload_if_changed = self.get_bool_option("Unit", ini::KEY_X_RELOADIFCHANGED);
        let restart_if_changed = self.get_bool_option("Unit", ini::KEY_X_RESTARTIFCHANGED);
        let stop_if_changed = self.get_bool_option("Unit", ini::KEY_X_STOPIFCHANGED);

        match (reload_if_changed, restart_if_changed, stop_if_changed) {
            (Some(true), _, _) => Some(UnitSwitchMethod::Reload),
            (_, Some(false), _) => Some(UnitSwitchMethod::KeepOld),
            (_, _, Some(false)) => Some(UnitSwitchMethod::Restart),
            (_, _, _) => None,
        }
    }

    /// Whether this unit is equal to the other unit using restart semantics.
    ///
    /// This simply checks whether the units are the same with changes to the
    /// following options being ignored:
    ///
    /// - `Unit.Description`
    /// - `Unit.Documentation`
    pub fn restart_eq(&self, other: &UnitFile) -> bool {
        self.config.eq_excluding(
            &other.config,
            &[("Unit", "Description"), ("Unit", "Documentation")],
        )
    }

    /// Whether this unit is equal to the other unit using reload semantics.
    ///
    /// This checks whether the units are equal as per `restart_eq` but also
    /// tolerate differences in the `Unit.X-Reload-Triggers` option.
    pub fn reload_eq(&self, other: &UnitFile) -> bool {
        self.config.eq_excluding(
            &other.config,
            &[
                ("Unit", "Description"),
                ("Unit", "Documentation"),
                ("Unit", "X-Reload-Triggers"),
            ],
        )
    }

    /// Gets a named Boolean option with priority given to the override file.
    fn get_bool_option(&self, section: &str, key: &str) -> Option<bool> {
        self.config.get_bool(section, key)
    }
}

fn unit_path_to_unit_type(unit_path: &Path) -> Result<UnitType, Error> {
    unit_path
        .extension()
        .and_then(|v| v.to_str())
        .and_then(|v| match v {
            "service" => Some(UnitType::Service),
            "socket" => Some(UnitType::Socket),
            "device" => Some(UnitType::Device),
            "mount" => Some(UnitType::Mount),
            "automount" => Some(UnitType::Automount),
            "swap" => Some(UnitType::Swap),
            "target" => Some(UnitType::Target),
            "path" => Some(UnitType::Path),
            "timer" => Some(UnitType::Timer),
            "slice" => Some(UnitType::Slice),
            "scope" => Some(UnitType::Scope),
            _ => None,
        })
        .ok_or_else(|| {
            Error::SdSwitch(format!(
                "Could not determine unit type from extension of {}",
                unit_path.display()
            ))
        })
}

fn collect_overrides(unit_path: &Path) -> Result<Option<Vec<PathBuf>>, Error> {
    // `expect` is fine here as `unit_path.file_name()` is already called in `load` by
    // `unit_path_to_unit_type`
    let mut drop_in_dir_name = unit_path.file_name().expect("XXX").to_os_string();
    drop_in_dir_name.push(".d");

    let drop_in_path = unit_path.with_file_name(drop_in_dir_name);
    if drop_in_path.is_dir() {
        let mut conf_files: Vec<_> = read_dir(drop_in_path)?
            .filter_map(|entry| {
                entry.ok().and_then(|entry| {
                    let path = entry.path();

                    if path.is_file() && path.extension().is_some_and(|ext| ext == "conf") {
                        return Some(path);
                    }
                    None
                })
            })
            .collect();

        conf_files.sort();
        return Ok(Some(conf_files));
    }
    Ok(None)
}

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

    fn test_data_path(file_name: &str) -> PathBuf {
        PathBuf::from("testdata/unit_file").join(file_name)
    }

    fn switch_method_file_path(method: &str) -> PathBuf {
        test_data_path(&format!("switch-method-{method}.service"))
    }

    #[test]
    fn switch_method_is_stop_only_on_refusemanualstart() {
        let uf = UnitFile::load(&test_data_path("refuse-manual-start.service")).unwrap();
        assert_eq!(uf.switch_method(), UnitSwitchMethod::StopOnly);
    }

    // Having refuse manual stop in the unit should not prevent regular restart since we don't know if the refuse actually applies to the currently running service.
    #[test]
    fn switch_method_is_stop_start_on_refusemanualstop() {
        let uf = UnitFile::load(&test_data_path("refuse-manual-stop.service")).unwrap();
        assert_eq!(uf.switch_method(), UnitSwitchMethod::StopStart);
    }

    #[test]
    fn can_read_switch_method_reload() {
        let uf = UnitFile::load(&switch_method_file_path("reload")).unwrap();
        assert_eq!(uf.switch_method(), UnitSwitchMethod::Reload);
    }

    #[test]
    fn can_read_switch_method_restart() {
        let uf = UnitFile::load(&switch_method_file_path("restart")).unwrap();
        assert_eq!(uf.switch_method(), UnitSwitchMethod::Restart);
    }

    #[test]
    fn can_read_switch_method_stop_start() {
        let uf = UnitFile::load(&switch_method_file_path("stop-start")).unwrap();
        assert_eq!(uf.switch_method(), UnitSwitchMethod::StopStart);
    }

    #[test]
    fn can_read_switch_method_keep_old() {
        let uf = UnitFile::load(&switch_method_file_path("keep-old")).unwrap();
        assert_eq!(uf.switch_method(), UnitSwitchMethod::KeepOld);
    }

    #[test]
    fn invalid_unit_fields_yields_error() {
        let uf = UnitFile::load(&test_data_path("invalid.service"));
        match uf {
            Ok(_) => panic!("Unexpected success"),
            Err(e) => assert_eq!(
                format!("{e}"),
                "testdata/unit_file/invalid.service: expected '[' but got 'T' at line 0, column 0"
            ),
        }
    }

    #[test]
    fn can_handle_legacy_reload() {
        let uf = UnitFile::load(&test_data_path("legacy-reload.service")).unwrap();
        assert_eq!(uf.switch_method(), UnitSwitchMethod::Reload);
    }

    #[test]
    fn can_handle_legacy_restart() {
        let uf = UnitFile::load(&test_data_path("legacy-restart.service")).unwrap();
        assert_eq!(uf.switch_method(), UnitSwitchMethod::Restart);
    }

    #[test]
    fn can_handle_legacy_stop_start() {
        let uf = UnitFile::load(&test_data_path("legacy-stop-start.service")).unwrap();
        assert_eq!(uf.switch_method(), UnitSwitchMethod::StopStart);
    }

    #[test]
    fn can_handle_legacy_keep_old() {
        let uf = UnitFile::load(&test_data_path("legacy-keep-old.service")).unwrap();
        assert_eq!(uf.switch_method(), UnitSwitchMethod::KeepOld);
    }

    #[test]
    fn switch_method_defaults_to_stop_start() {
        let uf = UnitFile::load(&test_data_path("empty.service")).unwrap();
        assert_eq!(uf.switch_method(), UnitSwitchMethod::StopStart);
    }

    #[test]
    fn can_add_option_using_overrides() {
        let uf = UnitFile::load(&test_data_path("unit_with_override.service")).unwrap();
        assert!(uf.config.get_unit_switch_method().is_some());
    }

    #[test]
    fn can_override_option_using_overrides() {
        let uf = UnitFile::load(&test_data_path("unit_with_override.service")).unwrap();
        let result = uf.get_bool_option("Unit", ini::KEY_REFUSEMANUALSTOP);
        assert_eq!(result, Some(true));
    }

    #[test]
    fn can_override_option_in_correct_order() {
        let uf = UnitFile::load(&test_data_path("unit_with_override.service")).unwrap();
        assert_eq!(uf.switch_method(), UnitSwitchMethod::KeepOld);
    }

    #[test]
    fn can_handle_reload_triggers_in_override_file() {
        let uf_a = UnitFile::load(&test_data_path("reload_overrides_a.service")).unwrap();
        let uf_b = UnitFile::load(&test_data_path("reload_overrides_b.service")).unwrap();

        let result = uf_a.switch_method();
        assert_eq!(result, UnitSwitchMethod::KeepOld);

        let result = uf_b.switch_method();
        assert_eq!(result, UnitSwitchMethod::KeepOld);

        // The units are not restart equal since they differ in the
        // reload-triggers field.
        assert!(!uf_a.restart_eq(&uf_b));

        // But they _are_ reload equal since they only differ in the description
        // and reload-triggers fields.
        assert!(uf_a.reload_eq(&uf_b));
    }

    #[test]
    fn can_load_file_with_escaped_characters() {
        let uf = UnitFile::load(&test_data_path("escaped-values.service")).unwrap();
        assert_eq!(uf.switch_method(), UnitSwitchMethod::Restart);
    }
}