sd-switch 0.6.4

use crate::systemd::JobSet;
use crate::systemd::UnitManager;
mod error;
mod i18n_lib;
pub mod systemd;
mod unit_file;

use i18n_lib::MSG;
use std::collections::BTreeMap;
use std::{
    collections::HashSet,
    path::{Path, PathBuf},
    rc::Rc,
    time::Duration,
};
use systemd::UnitStatus;
use unit_file::UnitFile;

use anyhow::{Context, Result};

fn pretty_unit_names<I>(unit_names: I) -> String
where
    I: IntoIterator,
    I::Item: AsRef<str>,
{
    let mut str_vec = unit_names
        .into_iter()
        .map(|s| String::from(s.as_ref()))
        .collect::<Vec<_>>();
    str_vec.sort();
    str_vec.join(", ")
}

fn is_unit_available(unit_path: &Path) -> bool {
    unit_path.exists()
        && !unit_path
            .canonicalize()
            .map(|p| p == Path::new("/dev/null"))
            .unwrap_or(true)
}

/// Given an active unit name this returns the actual unit file name. In the
/// case of a parameterized unit, e.g., `foo@bar.service` this function returns
/// `foo@.service`. For nonparameterized unit names it returns none.
fn parameterized_base_name(unit_name: &str) -> Option<String> {
    let res = unit_name.splitn(2, '@').collect::<Vec<_>>();
    match res[..] {
        [base_name, arg_and_suffix] => {
            let res = arg_and_suffix.rsplitn(2, '.').collect::<Vec<_>>();
            match res[..] {
                [suffix, arg] if !arg.is_empty() => Some(format!("{base_name}@.{suffix}")),
                _ => None,
            }
        }
        _ => None,
    }
}

/// Returns the file path of the given unit name within the given directory.
///
/// If no matching file is found then `None` is returned.
///
/// If the given unit name is a parameterized named then an exactly matching
/// file is returned, if it exists, otherwise the template file path is
/// returned.
fn find_unit_file_path(unit_directory: &Path, unit_name: &str) -> Option<PathBuf> {
    Some(unit_directory.join(unit_name))
        .filter(|e| is_unit_available(e))
        .or_else(|| {
            parameterized_base_name(unit_name)
                .map(|n| unit_directory.join(n))
                .filter(|e| is_unit_available(e))
        })
}

/// A plan of unit actions needed to accomplish the switch.
struct SwitchPlan {
    unit_plan: BTreeMap<Rc<str>, UnitPlan>,
}

impl SwitchPlan {
    fn build_unit_plan(
        &'_ mut self,
        name: Rc<str>,
        populate_decisions: bool,
    ) -> UnitPlanBuilder<'_> {
        UnitPlanBuilder {
            switch_plan: self,
            populate_decisions,
            name,
            decisions: Vec::new(),
        }
    }

    fn stop_units(&self) -> BTreeMap<&str, &UnitPlan> {
        self.units_with_action(|a| *a == UnitAction::Stop || *a == UnitAction::StopStart)
    }

    fn start_units(&self) -> BTreeMap<&str, &UnitPlan> {
        self.units_with_action(|a| *a == UnitAction::Start || *a == UnitAction::StopStart)
    }

    fn reload_units(&self) -> BTreeMap<&str, &UnitPlan> {
        self.units_with_action(|a| *a == UnitAction::Reload)
    }

    fn restart_units(&self) -> BTreeMap<&str, &UnitPlan> {
        self.units_with_action(|a| *a == UnitAction::Restart)
    }

    fn keep_old_units(&self) -> BTreeMap<&str, &UnitPlan> {
        self.units_with_action(|a| *a == UnitAction::KeepOld)
    }

    fn unchanged_units(&self) -> BTreeMap<&str, &UnitPlan> {
        self.units_with_action(|a| *a == UnitAction::NoAction)
    }

    fn units_with_action(
        &self,
        predicate: impl Fn(&UnitAction) -> bool,
    ) -> BTreeMap<&str, &UnitPlan> {
        self.unit_plan
            .iter()
            .filter(|(_, v)| predicate(&v.action))
            .map(|(k, v)| (k.as_ref(), v))
            .collect()
    }
}

/// The plan for a given unit. The plan consists of the action to perform plus a
/// chain of decisions why that action was chosen. Since populating the the
/// decision chain is pretty expensive, it is only actually populated when
/// specifically needed such as when verbose output is desired.
struct UnitPlan {
    action: UnitAction,
    /// The decision trace showing the decision taken to arrive at the given
    /// action.
    decisions: Vec<UnitDecision>,
}

struct UnitPlanBuilder<'a> {
    switch_plan: &'a mut SwitchPlan,
    populate_decisions: bool,
    name: Rc<str>,
    decisions: Vec<UnitDecision>,
}

impl<'a> UnitPlanBuilder<'a> {
    fn push_decision(mut self, decision: UnitDecision) -> Self {
        if self.populate_decisions {
            self.decisions.push(decision);
        }
        self
    }

    fn action(self, action: UnitAction) {
        self.switch_plan.unit_plan.insert(
            self.name,
            UnitPlan {
                action,
                decisions: self.decisions,
            },
        );
    }
}

/// The available actions for a given unit.
#[derive(Debug, PartialEq, Eq)]
enum UnitAction {
    NoAction,
    StopStart,
    Start,
    Stop,
    Restart,
    Reload,
    KeepOld,
}

/// A description of why a given action was chosen for a unit.
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
enum UnitDecision {
    /// Whether the new unit exists.
    NewExists,
    /// Whether the old unit exists.
    OldExists,
    /// There is no old unit but there is an active unit with the same name does.
    OldNotExists,
    /// New and old units are equal using restart semantics.
    RestartEq,
    /// New and old units are equal using reload semantics.
    ReloadEq,
    /// The new unit is having a special target unit.
    UnitType(unit_file::UnitType),
    /// The new unit has a special switch method.
    SwitchMethod(unit_file::UnitSwitchMethod),
    /// The active unit refuse manual stop.
    ActiveRefuseManualStop,
    /// Wanted by active target.
    WantedByActiveTarget,
}

struct UnitWithTarget {
    unit_path: PathBuf,
    unit_name: Rc<str>,
    target_name: Rc<str>,
}

fn build_switch_plan(
    old_dir: Option<&Path>,
    new_dir: &Path,
    populate_decisions: bool,
    service_manager: &impl systemd::ServiceManager,
) -> Result<SwitchPlan> {
    let mut switch_plan = SwitchPlan {
        unit_plan: BTreeMap::new(),
    };

    let mut active_unit_names = HashSet::new();

    let active_units = service_manager
        .list_units_by_states(&["active", "activating"])
        .with_context(|| MSG.err_listing_active_units())?;

    // Handle units that are currently active, typically this implies restarting
    // the unit in some way.
    for active_unit in active_units {
        let new_unit_path_opt = find_unit_file_path(new_dir, active_unit.name());
        let old_unit_path_opt = old_dir
            .as_ref()
            .and_then(|d| find_unit_file_path(d, active_unit.name()));

        let active_unit_name: Rc<str> = active_unit.name().into();
        active_unit_names.insert(active_unit_name.clone());

        let mut upb = switch_plan.build_unit_plan(active_unit_name, populate_decisions);

        if let Some(new_unit_path) = new_unit_path_opt {
            let new_unit_file = UnitFile::load(&new_unit_path)
                .with_context(|| MSG.err_read_unit_file(&new_unit_path))?;

            upb = upb.push_decision(UnitDecision::NewExists);

            if let Some(old_unit_path) = old_unit_path_opt {
                let old_unit_file = UnitFile::load(&old_unit_path)
                    .with_context(|| MSG.err_read_unit_file(&old_unit_path))?;

                upb = upb.push_decision(UnitDecision::OldExists);

                if old_unit_file.restart_eq(&new_unit_file) {
                    upb.push_decision(UnitDecision::RestartEq)
                        .action(UnitAction::NoAction);
                } else if old_unit_file.reload_eq(&new_unit_file) {
                    upb.push_decision(UnitDecision::ReloadEq)
                        .action(UnitAction::Reload);
                } else if new_unit_file.unit_type() == unit_file::UnitType::Target {
                    upb = upb.push_decision(UnitDecision::UnitType(new_unit_file.unit_type()));

                    if new_unit_file.switch_method() == unit_file::UnitSwitchMethod::StopOnly {
                        upb.push_decision(UnitDecision::SwitchMethod(
                            new_unit_file.switch_method(),
                        ))
                        .action(UnitAction::KeepOld);
                    } else {
                        upb.action(UnitAction::Start);
                    }
                } else {
                    upb = upb
                        .push_decision(UnitDecision::SwitchMethod(new_unit_file.switch_method()));

                    match new_unit_file.switch_method() {
                        unit_file::UnitSwitchMethod::Reload => {
                            upb.action(UnitAction::Reload);
                        }
                        unit_file::UnitSwitchMethod::Restart => {
                            upb.action(UnitAction::Restart);
                        }
                        unit_file::UnitSwitchMethod::StopStart => {
                            if service_manager
                                .unit_manager(&active_unit)?
                                .refuse_manual_stop()?
                            {
                                upb.push_decision(UnitDecision::ActiveRefuseManualStop)
                                    .action(UnitAction::NoAction);
                            } else {
                                upb.action(UnitAction::StopStart);
                            }
                        }
                        unit_file::UnitSwitchMethod::StopOnly => {
                            if service_manager
                                .unit_manager(&active_unit)?
                                .refuse_manual_stop()?
                            {
                                upb.push_decision(UnitDecision::ActiveRefuseManualStop)
                                    .action(UnitAction::NoAction);
                            } else {
                                upb.action(UnitAction::Stop);
                            }
                        }
                        unit_file::UnitSwitchMethod::KeepOld => {
                            upb.action(UnitAction::KeepOld);
                        }
                    }
                }
            } else {
                upb = upb.push_decision(UnitDecision::OldNotExists);

                if service_manager
                    .unit_manager(&active_unit)?
                    .refuse_manual_stop()?
                {
                    upb.push_decision(UnitDecision::ActiveRefuseManualStop)
                        .action(UnitAction::KeepOld);
                } else if new_unit_file.switch_method() == unit_file::UnitSwitchMethod::StopOnly {
                    upb.push_decision(UnitDecision::SwitchMethod(new_unit_file.switch_method()))
                        .action(UnitAction::KeepOld);
                } else {
                    upb.action(UnitAction::StopStart);
                }
            }
        } else if old_unit_path_opt.is_some() {
            upb = upb.push_decision(UnitDecision::OldExists);

            if service_manager
                .unit_manager(&active_unit)?
                .refuse_manual_stop()?
            {
                upb.push_decision(UnitDecision::ActiveRefuseManualStop)
                    .action(UnitAction::KeepOld);
            } else {
                upb.action(UnitAction::Stop);
            }
        }
    }

    // Handle units that are not currently active but are wanted by an active
    // target. Typically this is simply a matter of starting the new unit.
    for wanted_unit in find_wanted_units(new_dir)? {
        // Skip if the unit is not wanted by an active target.
        if !active_unit_names.contains(&wanted_unit.target_name) {
            continue;
        }

        // Skip if the unit is actually active.
        if active_unit_names.contains(&wanted_unit.unit_name) {
            continue;
        }

        let new_unit_file = UnitFile::load(&wanted_unit.unit_path)
            .with_context(|| MSG.err_read_unit_file(&wanted_unit.unit_path))?;

        let mut upb = switch_plan.build_unit_plan(wanted_unit.unit_name, populate_decisions);

        upb = upb
            .push_decision(UnitDecision::NewExists)
            .push_decision(UnitDecision::WantedByActiveTarget);

        if new_unit_file.switch_method() == unit_file::UnitSwitchMethod::StopOnly {
            upb.push_decision(UnitDecision::SwitchMethod(new_unit_file.switch_method()))
                .action(UnitAction::NoAction);
        } else {
            upb.action(UnitAction::Start);
        }
    }

    Ok(switch_plan)
}

fn find_wanted_units(new_dir: &Path) -> Result<Vec<UnitWithTarget>> {
    let mut result = Vec::new();

    for dir_entry in std::fs::read_dir(new_dir)? {
        let dir_entry = dir_entry.with_context(|| MSG.err_read_dir_entry(new_dir))?;

        // Get the file name as a string. Ignore the string if we cannot parse it.
        let entry_file_name = dir_entry
            .file_name()
            .into_string()
            .expect("unit with valid Unicode file name");

        if dir_entry.metadata()?.is_dir() && entry_file_name.ends_with(".target.wants") {
            let dir_name = entry_file_name;
            let target_name: Rc<str> = dir_name
                .strip_suffix(".wants")
                .expect("directory name should end in .wants")
                .into();
            for wants_entry in std::fs::read_dir(dir_entry.path())? {
                let wants_entry = wants_entry
                    .with_context(|| MSG.err_read_dir_entry(dir_entry.path().as_path()))?;

                let unit_name = wants_entry
                    .file_name()
                    .into_string()
                    .expect("unit with valid Unicode file name")
                    .into();
                result.push(UnitWithTarget {
                    unit_path: wants_entry.path(),
                    unit_name,
                    target_name: target_name.clone(),
                });
            }
        }
    }

    Ok(result)
}

fn exec_pre_reload<F>(
    plan: &SwitchPlan,
    service_manager: &impl systemd::ServiceManager,
    job_handler: F,
    dry_run: bool,
    timeout: Duration,
) -> Result<()>
where
    F: Fn(&str, &str) + Send + 'static,
{
    let stop_units = plan.stop_units();

    if stop_units.is_empty() {
        return Ok(());
    }

    println!(
        "{}",
        MSG.stopping_units(&pretty_unit_names(stop_units.keys()))
    );

    if !dry_run {
        let mut job_set = service_manager.new_job_set()?;

        for uf in stop_units.keys() {
            job_set
                .stop_unit(uf)
                .with_context(|| MSG.err_unit_action_failed(uf, UnitAction::Stop))?;
        }

        job_set.wait_for_all(job_handler, timeout)?;
    }

    Ok(())
}

fn exec_reload(
    service_manager: &impl systemd::ServiceManager,
    dry_run: bool,
    verbose: bool,
) -> Result<()> {
    if !dry_run {
        if verbose {
            println!("{}", MSG.resetting_failed_units());
        }
        service_manager
            .reset_failed()
            .with_context(|| MSG.err_resetting_failed_units())?;

        if verbose {
            println!("{}", MSG.reloading_systemd());
        }
        service_manager
            .daemon_reload()
            .with_context(|| MSG.err_reloading_systemd())?;
    }

    Ok(())
}

fn exec_post_reload<F>(
    plan: &SwitchPlan,
    service_manager: &impl systemd::ServiceManager,
    job_handler: F,
    dry_run: bool,
    verbose: bool,
    timeout: Duration,
) -> Result<()>
where
    F: Fn(&str, &str) + Send + 'static,
{
    let mut job_set = service_manager.new_job_set()?;

    {
        let units = plan.reload_units();
        if !units.is_empty() {
            println!("{}", MSG.reloading_units(&pretty_unit_names(units.keys())));

            if !dry_run {
                for uf in units.keys() {
                    job_set
                        .reload_unit(uf)
                        .with_context(|| MSG.err_unit_action_failed(uf, UnitAction::Reload))?;
                }
            }
        }
    }

    {
        let units = plan.restart_units();
        if !units.is_empty() {
            println!("{}", MSG.restarting_units(&pretty_unit_names(units.keys())));

            if !dry_run {
                for uf in units.keys() {
                    job_set
                        .restart_unit(uf)
                        .with_context(|| MSG.err_unit_action_failed(uf, UnitAction::Restart))?;
                }
            }
        }
    }

    {
        let units = plan.keep_old_units();
        if !units.is_empty() {
            println!(
                "{}",
                MSG.keeping_old_units(&pretty_unit_names(units.keys()))
            );
        }
    }

    if verbose {
        let units = plan.unchanged_units();
        if !units.is_empty() {
            println!("{}", MSG.unchanged_units(&pretty_unit_names(units.keys())));
        }
    }

    {
        let units = plan.start_units();
        if !units.is_empty() {
            println!("{}", MSG.starting_units(&pretty_unit_names(units.keys())));

            if !dry_run {
                for uf in units.keys() {
                    job_set
                        .start_unit(uf)
                        .with_context(|| MSG.err_unit_action_failed(uf, UnitAction::Start))?;
                }
            }
        }
    }

    job_set.wait_for_all(job_handler, timeout)?;

    Ok(())
}

fn print_plan(plan: &SwitchPlan) {
    if plan.unit_plan.is_empty() {
        println!("The calculated switch plan is empty.");
        return;
    }

    println!("Unit switch plan:");

    for (unit_file, unit_plan) in &plan.unit_plan {
        let action = match unit_plan.action {
            UnitAction::NoAction => "No action",
            UnitAction::StopStart => "Stop/Start",
            UnitAction::Start => "Start",
            UnitAction::Stop => "Stop",
            UnitAction::Restart => "Restart",
            UnitAction::Reload => "Reload",
            UnitAction::KeepOld => "Keep active",
        };
        let decisions = pretty_unit_names(unit_plan.decisions.iter().map(|v| format!("{:?}", v)));
        println!("  {action} {unit_file} ({decisions})");
    }
}

/// Performs a systemd unit "switch".
pub fn switch(
    service_manager: &impl systemd::ServiceManager,
    old_dir: Option<&Path>,
    new_dir: &Path,
    dry_run: bool,
    verbose: bool,
    timeout: Duration,
) -> Result<()> {
    let system_status = service_manager.system_status()?;

    // If the systemd manager status is degraded then inform the user about the
    // failed units. We will still attempt to perform the switch, though.
    if matches!(system_status, systemd::SystemStatus::Degraded) {
        let units_by_states = service_manager.list_units_by_states(&["failed"])?;
        let failed: Vec<&str> = units_by_states.iter().map(|status| status.name()).collect();
        let failed = failed.join(", ");
        eprintln!(
            "The service manager is degraded.\n\
             Failed services: {failed}\n\
             Attempting to continue anyway..."
        );
    }

    let do_switch = {
        use systemd::SystemStatus::{
            Degraded, Initializing, Maintenance, Running, Starting, Stopping,
        };
        match system_status {
            Initializing | Starting | Running | Degraded => true,
            Maintenance | Stopping => false,
        }
    };

    if !do_switch {
        if verbose {
            println!("Skipping switch since systemd has {system_status} status");
        }
        return Ok(());
    }

    let plan = build_switch_plan(old_dir, new_dir, verbose, service_manager)
        .context("Failed to build switch plan")?;

    if verbose {
        print_plan(&plan);
    }

    let job_handler = move |name: &str, state: &str| {
        if verbose || state != "done" {
            println!("{name} {state}");
        }
    };

    exec_pre_reload(&plan, service_manager, job_handler, dry_run, timeout)
        .context("Failed to perform pre-reload tasks")?;

    exec_reload(service_manager, dry_run, verbose)?;

    exec_post_reload(
        &plan,
        service_manager,
        job_handler,
        dry_run,
        verbose,
        timeout,
    )
    .context("Failed to perform post-reload tasks")?;

    Ok(())
}

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

    #[test]
    fn can_get_base_name_for_parameterized_unit() {
        assert_eq!(
            parameterized_base_name("foo@bar.service"),
            Some(String::from("foo@.service"))
        );
        assert_eq!(
            parameterized_base_name("foo@bar.baz.service"),
            Some(String::from("foo@.service"))
        );
    }

    #[test]
    fn no_base_name_for_nonparameterized_units() {
        assert_eq!(parameterized_base_name("foo@.service"), None);
        assert_eq!(parameterized_base_name("foo.service"), None);
        assert_eq!(parameterized_base_name("foo@barservice"), None);
    }
}