sd_switch/

lib.rs

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

use i18n_lib::UnitAction;
use i18n_lib::MSG;
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.
#[derive(Debug)]
struct SwitchPlan {
    stop_units: HashSet<Rc<str>>,
    start_units: HashSet<Rc<str>>,
    reload_units: HashSet<Rc<str>>,
    restart_units: HashSet<Rc<str>>,
    keep_old_units: HashSet<Rc<str>>,
    unchanged_units: HashSet<Rc<str>>,
}

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

fn build_switch_plan(
    old_dir: Option<&Path>,
    new_dir: &Path,
    service_manager: &impl systemd::ServiceManager,
) -> Result<SwitchPlan> {
    let mut stop_units = HashSet::new();
    let mut start_units = HashSet::new();
    let mut reload_units = HashSet::new();
    let mut restart_units = HashSet::new();
    let mut keep_old_units = HashSet::new();
    let mut unchanged_units = HashSet::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());

        if let Some(new_unit_path) = new_unit_path_opt {
            let new_unit_file = UnitFile::load(&new_unit_path).with_context(|| {
                format!("Failed load of new unit file {}", new_unit_path.display())
            })?;

            if let Some(old_unit_path) = old_unit_path_opt {
                let old_unit_file = UnitFile::load(&old_unit_path).with_context(|| {
                    format!("Failed load of old unit file {}", old_unit_path.display())
                })?;

                if old_unit_file.restart_eq(&new_unit_file) {
                    unchanged_units.insert(active_unit_name);
                } else if old_unit_file.reload_eq(&new_unit_file) {
                    reload_units.insert(active_unit_name);
                } else if new_unit_file.unit_type() == unit_file::UnitType::Target {
                    if new_unit_file.switch_method() == unit_file::UnitSwitchMethod::StopOnly {
                        keep_old_units.insert(active_unit_name);
                    } else {
                        start_units.insert(active_unit_name);
                    }
                } else {
                    match new_unit_file.switch_method() {
                        unit_file::UnitSwitchMethod::Reload => {
                            reload_units.insert(active_unit_name);
                        }
                        unit_file::UnitSwitchMethod::Restart => {
                            restart_units.insert(active_unit_name);
                        }
                        unit_file::UnitSwitchMethod::StopStart => {
                            if service_manager
                                .unit_manager(&active_unit)?
                                .refuse_manual_stop()?
                            {
                                keep_old_units.insert(active_unit_name);
                            } else {
                                stop_units.insert(active_unit_name.clone());
                                start_units.insert(active_unit_name);
                            }
                        }
                        unit_file::UnitSwitchMethod::StopOnly => {
                            if service_manager
                                .unit_manager(&active_unit)?
                                .refuse_manual_stop()?
                            {
                                keep_old_units.insert(active_unit_name);
                            } else {
                                stop_units.insert(active_unit_name);
                            }
                        }
                        unit_file::UnitSwitchMethod::KeepOld => {
                            keep_old_units.insert(active_unit_name);
                        }
                    }
                }
            } else if service_manager
                .unit_manager(&active_unit)?
                .refuse_manual_stop()?
                || new_unit_file.switch_method() == unit_file::UnitSwitchMethod::StopOnly
            {
                keep_old_units.insert(active_unit_name);
            } else {
                stop_units.insert(active_unit_name.clone());
                start_units.insert(active_unit_name);
            }
        } else if old_unit_path_opt.is_some() {
            if service_manager
                .unit_manager(&active_unit)?
                .refuse_manual_stop()?
            {
                keep_old_units.insert(active_unit_name);
            } else {
                stop_units.insert(active_unit_name);
            }
        }
    }

    // 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(|| {
            format!(
                "Failed load of wanted unit file {}",
                wanted_unit.unit_path.display()
            )
        })?;

        if new_unit_file.switch_method() == unit_file::UnitSwitchMethod::StopOnly {
            continue;
        }

        start_units.insert(wanted_unit.unit_name);
    }

    Ok(SwitchPlan {
        stop_units,
        start_units,
        reload_units,
        restart_units,
        keep_old_units,
        unchanged_units,
    })
}

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,
{
    if !plan.stop_units.is_empty() {
        println!(
            "{}",
            MSG.stopping_units(&pretty_unit_names(&plan.stop_units))
        );
        if !dry_run {
            let mut job_set = service_manager.new_job_set()?;

            for uf in &plan.stop_units {
                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()?;

    if !plan.reload_units.is_empty() {
        println!(
            "{}",
            MSG.reloading_units(&pretty_unit_names(&plan.reload_units))
        );
        if !dry_run {
            for uf in &plan.reload_units {
                job_set
                    .reload_unit(uf)
                    .with_context(|| MSG.err_unit_action_failed(uf, UnitAction::Reload))?;
            }
        }
    }

    if !plan.restart_units.is_empty() {
        println!(
            "{}",
            MSG.restarting_units(&pretty_unit_names(&plan.restart_units))
        );
        if !dry_run {
            for uf in &plan.restart_units {
                job_set
                    .restart_unit(uf)
                    .with_context(|| MSG.err_unit_action_failed(uf, UnitAction::Restart))?;
            }
        }
    }

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

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

    if !plan.start_units.is_empty() {
        println!(
            "{}",
            MSG.starting_units(&pretty_unit_names(&plan.start_units))
        );
        if !dry_run {
            for uf in &plan.start_units {
                job_set
                    .start_unit(uf)
                    .with_context(|| MSG.err_unit_action_failed(uf, UnitAction::Start))?;
            }
        }
    }

    job_set.wait_for_all(job_handler, timeout)?;

    Ok(())
}

/// Performs a systemd unit "switch".
pub fn switch(
    service_manager: &impl systemd::ServiceManager,
    //connection: &zbus::blocking::Connection,
    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, service_manager)
        .context("Failed to build switch 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);
    }
}