steamos-manager/src/systemd.rs
2025-05-30 16:04:16 -07:00

457 lines
15 KiB
Rust

/*
* Copyright © 2023 Collabora Ltd.
* Copyright © 2024 Valve Software
*
* SPDX-License-Identifier: MIT
*/
use anyhow::{anyhow, Result};
use std::path::PathBuf;
use std::str::FromStr;
use strum::{Display, EnumString};
use zbus::proxy::CacheProperties;
use zbus::zvariant::OwnedObjectPath;
use zbus::Connection;
#[zbus::proxy(
interface = "org.freedesktop.systemd1.Unit",
default_service = "org.freedesktop.systemd1"
)]
trait SystemdUnit {
#[zbus(property)]
fn active_state(&self) -> Result<String>;
#[zbus(property)]
fn unit_file_state(&self) -> Result<String>;
async fn restart(&self, mode: &str) -> Result<OwnedObjectPath>;
async fn start(&self, mode: &str) -> Result<OwnedObjectPath>;
async fn stop(&self, mode: &str) -> Result<OwnedObjectPath>;
}
#[zbus::proxy(
interface = "org.freedesktop.systemd1.Manager",
default_service = "org.freedesktop.systemd1",
default_path = "/org/freedesktop/systemd1"
)]
trait SystemdManager {
#[allow(clippy::type_complexity)]
async fn enable_unit_files(
&self,
files: &[&str],
runtime: bool,
force: bool,
) -> Result<(bool, Vec<(String, String, String)>)>;
async fn disable_unit_files(
&self,
files: &[&str],
runtime: bool,
) -> Result<Vec<(String, String, String)>>;
async fn mask_unit_files(
&self,
files: &[&str],
runtime: bool,
force: bool,
) -> Result<Vec<(String, String, String)>>;
async fn unmask_unit_files(
&self,
files: &[&str],
runtime: bool,
) -> Result<Vec<(String, String, String)>>;
async fn reload(&self) -> Result<()>;
}
#[derive(Display, EnumString, PartialEq, Debug, Copy, Clone)]
#[strum(serialize_all = "lowercase")]
pub enum EnableState {
Disabled,
Enabled,
Masked,
Static,
}
pub struct SystemdUnit<'dbus> {
connection: Connection,
proxy: SystemdUnitProxy<'dbus>,
name: String,
}
pub async fn daemon_reload(connection: &Connection) -> Result<()> {
let proxy = SystemdManagerProxy::new(connection).await?;
proxy.reload().await?;
Ok(())
}
impl<'dbus> SystemdUnit<'dbus> {
pub async fn new(connection: Connection, name: &str) -> Result<SystemdUnit<'dbus>> {
let path = PathBuf::from("/org/freedesktop/systemd1/unit").join(escape(name));
let path = String::from(path.to_str().ok_or(anyhow!("Unit name {name} invalid"))?);
Ok(SystemdUnit {
proxy: SystemdUnitProxy::builder(&connection)
.cache_properties(CacheProperties::No)
.path(path)?
.build()
.await?,
connection,
name: String::from(name),
})
}
pub async fn restart(&self) -> Result<()> {
self.proxy.restart("fail").await?;
Ok(())
}
pub async fn start(&self) -> Result<()> {
self.proxy.start("fail").await?;
Ok(())
}
pub async fn stop(&self) -> Result<()> {
self.proxy.stop("fail").await?;
Ok(())
}
#[allow(unused)]
pub async fn enable(&self) -> Result<bool> {
let manager = SystemdManagerProxy::new(&self.connection).await?;
let (_, res) = manager
.enable_unit_files(&[self.name.as_str()], false, false)
.await?;
Ok(!res.is_empty())
}
#[allow(unused)]
pub async fn disable(&self) -> Result<bool> {
let manager = SystemdManagerProxy::new(&self.connection).await?;
let res = manager
.disable_unit_files(&[self.name.as_str()], false)
.await?;
Ok(!res.is_empty())
}
pub async fn mask(&self) -> Result<bool> {
let manager = SystemdManagerProxy::new(&self.connection).await?;
let res = manager
.mask_unit_files(&[self.name.as_str()], false, false)
.await?;
Ok(!res.is_empty())
}
pub async fn unmask(&self) -> Result<bool> {
let manager = SystemdManagerProxy::new(&self.connection).await?;
let res = manager
.unmask_unit_files(&[self.name.as_str()], false)
.await?;
Ok(!res.is_empty())
}
pub async fn active(&self) -> Result<bool> {
Ok(self.proxy.active_state().await? == "active")
}
pub async fn enabled(&self) -> Result<EnableState> {
Ok(EnableState::from_str(
self.proxy.unit_file_state().await?.as_str(),
)?)
}
}
pub fn escape(name: &str) -> String {
let mut parts = String::new();
for c in name.chars() {
if c.is_ascii_alphanumeric() {
parts.push(c);
} else {
let escaped = format!("_{:02x}", u32::from(c));
parts.push_str(escaped.as_str());
}
}
parts
}
#[cfg(test)]
pub mod test {
use super::*;
use crate::error::to_zbus_fdo_error;
use crate::{enum_roundtrip, testing};
use std::collections::HashMap;
use std::time::Duration;
use tokio::time::sleep;
use zbus::{fdo, ObjectServer};
use zbus::zvariant::ObjectPath;
#[test]
fn enable_state_roundtrip() {
enum_roundtrip!(EnableState {
"disabled": str = Disabled,
"enabled": str = Enabled,
"masked": str = Masked,
"static": str = Static,
});
assert!(EnableState::from_str("loaded").is_err());
}
#[test]
fn test_escape() {
assert_eq!(escape("systemd"), "systemd");
assert_eq!(escape("system d"), "system_20d");
}
#[derive(Default)]
pub struct MockUnit {
pub active: String,
pub unit_file: String,
job: u32,
}
#[derive(Default)]
pub struct MockManager {
states: HashMap<String, EnableState>,
}
#[zbus::interface(name = "org.freedesktop.systemd1.Unit")]
impl MockUnit {
#[zbus(property)]
fn active_state(&self) -> fdo::Result<String> {
Ok(self.active.clone())
}
#[zbus(property)]
fn unit_file_state(&self) -> fdo::Result<String> {
Ok(self.unit_file.clone())
}
async fn restart(&mut self, mode: &str) -> fdo::Result<OwnedObjectPath> {
if mode != "fail" {
return Err(to_zbus_fdo_error("Invalid mode"));
}
let path = ObjectPath::try_from(format!("/restart/{mode}/{}", self.job))
.map_err(to_zbus_fdo_error)?;
self.job += 1;
self.active = String::from("active");
Ok(path.into())
}
async fn start(&mut self, mode: &str) -> fdo::Result<OwnedObjectPath> {
if mode != "fail" {
return Err(to_zbus_fdo_error("Invalid mode"));
}
let path = ObjectPath::try_from(format!("/start/{mode}/{}", self.job))
.map_err(to_zbus_fdo_error)?;
self.job += 1;
self.active = String::from("active");
Ok(path.into())
}
async fn stop(&mut self, mode: &str) -> fdo::Result<OwnedObjectPath> {
if mode != "fail" {
return Err(to_zbus_fdo_error("Invalid mode"));
}
let path = ObjectPath::try_from(format!("/stop/{mode}/{}", self.job))
.map_err(to_zbus_fdo_error)?;
self.job += 1;
self.active = String::from("inactive");
Ok(path.into())
}
}
#[zbus::interface(name = "org.freedesktop.systemd1.Manager")]
impl MockManager {
#[allow(clippy::type_complexity)]
async fn enable_unit_files(
&mut self,
files: Vec<String>,
_runtime: bool,
_force: bool,
#[zbus(object_server)] object_server: &ObjectServer,
) -> fdo::Result<(bool, Vec<(String, String, String)>)> {
let mut res = Vec::new();
for file in files {
if let Some(state) = self.states.get(&file) {
if *state == EnableState::Disabled {
self.states.insert(file.to_string(), EnableState::Enabled);
res.push((String::default(), String::default(), file.to_string()));
}
} else {
self.states.insert(file.to_string(), EnableState::Enabled);
res.push((String::default(), String::default(), file.to_string()));
}
let path = PathBuf::from("/org/freedesktop/systemd1/unit").join(escape(&file));
let mock_unit = object_server
.interface::<_, MockUnit>(path.to_string_lossy())
.await;
if let Ok(mock_unit) = mock_unit {
dbg!();
mock_unit.get_mut().await.unit_file = String::from("enabled");
}
}
Ok((true, res))
}
async fn disable_unit_files(
&mut self,
files: Vec<String>,
_runtime: bool,
#[zbus(object_server)] object_server: &ObjectServer,
) -> fdo::Result<Vec<(String, String, String)>> {
let mut res = Vec::new();
for file in files {
if let Some(state) = self.states.get(&file) {
if *state == EnableState::Enabled {
self.states.insert(file.to_string(), EnableState::Disabled);
res.push((String::default(), String::default(), file.to_string()));
}
} else {
self.states.insert(file.to_string(), EnableState::Disabled);
res.push((String::default(), String::default(), file.to_string()));
}
let path = PathBuf::from("/org/freedesktop/systemd1/unit").join(escape(&file));
let mock_unit = object_server
.interface::<_, MockUnit>(path.to_string_lossy())
.await;
if let Ok(mock_unit) = mock_unit {
mock_unit.get_mut().await.unit_file = String::from("disabled");
}
}
Ok(res)
}
async fn mask_unit_files(
&mut self,
files: Vec<String>,
_runtime: bool,
_force: bool,
) -> fdo::Result<Vec<(String, String, String)>> {
let mut res = Vec::new();
for file in files {
if let Some(state) = self.states.get(&file) {
if *state != EnableState::Masked {
self.states.insert(file.to_string(), EnableState::Masked);
res.push((String::default(), String::default(), file.to_string()));
}
} else {
self.states.insert(file.to_string(), EnableState::Masked);
res.push((String::default(), String::default(), file.to_string()));
}
}
Ok(res)
}
async fn unmask_unit_files(
&mut self,
files: Vec<String>,
_runtime: bool,
) -> fdo::Result<Vec<(String, String, String)>> {
let mut res = Vec::new();
for file in files {
if let Some(state) = self.states.get(&file) {
if *state == EnableState::Masked {
self.states.remove(&file);
res.push((String::default(), String::default(), file.to_string()));
}
}
}
Ok(res)
}
async fn reload(&self) -> fdo::Result<()> {
Ok(())
}
}
#[tokio::test]
async fn test_unit() {
let mut h = testing::start();
let mut unit = MockUnit::default();
unit.active = String::from("inactive");
unit.unit_file = String::from("disabled");
let connection = h.new_dbus().await.expect("dbus");
connection
.request_name("org.freedesktop.systemd1")
.await
.expect("request_name");
let object_server = connection.object_server();
object_server
.at("/org/freedesktop/systemd1/unit/test_2eservice", unit)
.await
.expect("at");
object_server
.at("/org/freedesktop/systemd1", MockManager::default())
.await
.expect("at");
let mock_unit = object_server
.interface::<_, MockUnit>("/org/freedesktop/systemd1/unit/test_2eservice")
.await
.unwrap();
sleep(Duration::from_millis(10)).await;
let unit = SystemdUnit::new(connection.clone(), "test.service")
.await
.expect("unit");
assert_eq!(unit.active().await.unwrap(), false);
assert!(unit.start().await.is_ok());
assert_eq!(mock_unit.get().await.active, "active");
assert_eq!(unit.active().await.unwrap(), true);
assert!(unit.restart().await.is_ok());
assert_eq!(mock_unit.get().await.active, "active");
assert_eq!(unit.active().await.unwrap(), true);
assert!(unit.stop().await.is_ok());
assert_eq!(mock_unit.get().await.active, "inactive");
assert_eq!(unit.active().await.unwrap(), false);
assert_eq!(mock_unit.get().await.unit_file, "disabled");
assert_eq!(unit.enabled().await.unwrap(), EnableState::Disabled);
assert!(unit.enable().await.unwrap());
assert_eq!(mock_unit.get().await.unit_file, "enabled");
assert_eq!(unit.enabled().await.unwrap(), EnableState::Enabled);
assert!(unit.disable().await.unwrap());
assert_eq!(mock_unit.get().await.unit_file, "disabled");
assert_eq!(unit.enabled().await.unwrap(), EnableState::Disabled);
}
#[tokio::test]
async fn test_manager() {
let mut h = testing::start();
let mut unit = MockUnit::default();
unit.active = String::from("active");
unit.unit_file = String::from("enabled");
let connection = h.new_dbus().await.expect("dbus");
connection
.request_name("org.freedesktop.systemd1")
.await
.expect("request_name");
let object_server = connection.object_server();
object_server
.at("/org/freedesktop/systemd1/unit/test_2eservice", unit)
.await
.expect("at");
object_server
.at("/org/freedesktop/systemd1", MockManager::default())
.await
.expect("at");
sleep(Duration::from_millis(10)).await;
let unit = SystemdUnit::new(connection.clone(), "test.service")
.await
.expect("unit");
assert!(unit.enable().await.unwrap());
assert!(!unit.enable().await.unwrap());
assert!(unit.disable().await.unwrap());
assert!(!unit.disable().await.unwrap());
assert!(unit.mask().await.unwrap());
assert!(!unit.mask().await.unwrap());
assert!(unit.unmask().await.unwrap());
assert!(!unit.unmask().await.unwrap());
}
}