use std::cell::Cell;
use std::collections::HashMap;
use std::future::poll_fn;
use std::pin::Pin;
use std::rc::Rc;
use std::task::{Poll, Waker};
use std::time::Duration;
use crabshell_reactive::{Signal, untrack};
use zbus::export::futures_core::Stream;
use zbus::message::Type as Kind;
use zbus::{Connection, Message, MessageStream, OwnedVariant};
use crate::Shared;
pub(crate) const BUS: &str = "org.freedesktop.DBus";
pub(crate) const BUS_PATH: &str = "/org/freedesktop/DBus";
pub(crate) const PROPERTIES: &str = "org.freedesktop.DBus.Properties";
pub(crate) const OBJECT_MANAGER: &str = "org.freedesktop.DBus.ObjectManager";
pub(crate) const CALL_TIMEOUT: Duration = Duration::from_secs(25);
pub(crate) type Props = HashMap<String, OwnedVariant>;
pub(crate) async fn bus_call(
connection: &Connection,
method: &str,
rule: &str,
) -> zbus::Result<()> {
let call = connection.call_method(Some(BUS), BUS_PATH, Some(BUS), method, &rule);
call.await.map(drop)
}
pub(crate) fn owner_rule(bus: &str) -> String {
format!("type='signal',sender='{BUS}',interface='{BUS}',member='NameOwnerChanged',arg0='{bus}'")
}
pub(crate) async fn owner(connection: &Connection, bus: &str) -> Option<String> {
let call = connection.call_method(Some(BUS), BUS_PATH, Some(BUS), "GetNameOwner", &bus);
call.await.ok()?.body().deserialize::<String>().ok()
}
pub(crate) fn owner_changed(message: &Message) -> Option<(String, String)> {
let header = message.header();
if header.sender().map(|sender| sender.as_str()) != Some(BUS)
|| header.member().map(|member| member.as_str()) != Some("NameOwnerChanged")
{
return None;
}
let (name, _, owner): (String, String, String) = message.body().deserialize().ok()?;
Some((name, owner))
}
pub(crate) async fn get_all(
shared: &Shared,
connection: &Connection,
destination: &str,
path: &str,
interface: &str,
) -> zbus::Result<Props> {
let call = connection.call_method(
Some(destination),
path,
Some(PROPERTIES),
"GetAll",
&interface,
);
let reply = shared.timeout(CALL_TIMEOUT, call).await?;
reply.body().deserialize::<Props>()
}
pub(crate) async fn set(
shared: &Shared,
connection: &Connection,
destination: &str,
(path, interface): (&str, &str),
name: &str,
value: zbus::Variant<'_>,
) -> zbus::Result<()> {
let body = (interface, name, value);
let call = connection.call_method(Some(destination), path, Some(PROPERTIES), "Set", &body);
shared.timeout(CALL_TIMEOUT, call).await.map(drop)
}
pub(crate) fn properties_changed(
message: &Message,
) -> Option<(String, String, Props, Vec<String>)> {
let header = message.header();
if header.interface().map(|interface| interface.as_str()) != Some(PROPERTIES)
|| header.member().map(|member| member.as_str()) != Some("PropertiesChanged")
{
return None;
}
let path = header.path()?.as_str().to_owned();
let (interface, changed, dropped): (String, Props, Vec<String>) =
message.body().deserialize().ok()?;
Some((path, interface, changed, dropped))
}
pub(crate) fn apply(props: &mut Props, changed: Props, dropped: Vec<String>) {
for name in dropped {
props.remove(&name);
}
props.extend(changed);
}
pub(crate) fn get<'a, T>(props: &'a Props, name: &str) -> Option<T>
where
T: TryFrom<&'a OwnedVariant>,
{
T::try_from(props.get(name)?).ok()
}
pub(crate) fn text(props: &Props, name: &str) -> String {
get::<&str>(props, name).unwrap_or_default().to_owned()
}
pub(crate) fn set_changed<T: PartialEq + 'static>(signal: Signal<T>, value: T) {
if untrack(|| signal.try_with(|old| *old != value)) == Some(true) {
signal.try_set(value);
}
}
pub(crate) trait Daemon: 'static {
const NAME: &'static str;
fn alive(&self) -> bool;
fn clear(&self) -> u64;
async fn read(self: Rc<Self>, connection: Connection, generation: u64);
fn signal(self: &Rc<Self>, connection: &Connection, message: &Message);
fn rules() -> Vec<String> {
vec![format!("type='signal',sender='{}'", Self::NAME)]
}
fn reads_again(&self, _message: &Message) -> bool {
false
}
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum Reading {
Idle,
Busy,
Stale,
}
fn read<D: Daemon>(
shared: &Rc<Shared>,
daemon: &Rc<D>,
connection: &Connection,
(reading, generation): (&Rc<Cell<Reading>>, &Rc<Cell<u64>>),
) {
if reading.get() != Reading::Idle {
return reading.set(Reading::Stale);
}
reading.set(Reading::Busy);
let (daemon, connection) = (Rc::clone(daemon), connection.clone());
let (reading, generation) = (Rc::clone(reading), Rc::clone(generation));
shared.spawn(async move {
loop {
Rc::clone(&daemon)
.read(connection.clone(), generation.get())
.await;
if reading.get() != Reading::Stale {
return reading.set(Reading::Idle);
}
reading.set(Reading::Busy);
}
});
}
struct WakeOnDispose(Rc<Cell<Option<Waker>>>);
impl Drop for WakeOnDispose {
fn drop(&mut self) {
if let Some(waker) = self.0.take() {
waker.wake();
}
}
}
pub(crate) async fn follow<D: Daemon>(shared: Rc<Shared>, daemon: Rc<D>, address: zbus::Address) {
let name = D::NAME;
if shared.root.is_disposed() {
return;
}
let waker = Rc::new(Cell::new(None));
shared.root.signal(WakeOnDispose(Rc::clone(&waker)));
let connection = match zbus::connection::Builder::address(address).build().await {
Ok(connection) => connection,
Err(error) => return tracing::info!(%error, name, "can't reach the system bus"),
};
let mut messages = MessageStream::from(&connection);
for rule in D::rules().into_iter().chain([owner_rule(name)]) {
if let Err(error) = bus_call(&connection, "AddMatch", &rule).await {
return tracing::warn!(%error, name, "couldn't follow it on the bus");
}
}
let reading = Rc::new(Cell::new(Reading::Idle));
let generation = Rc::new(Cell::new(daemon.clear()));
let state = (&reading, &generation);
let mut current = owner(&connection, name).await;
match ¤t {
Some(_) => read(&shared, &daemon, &connection, state),
None => tracing::info!(name, "not running; waiting for it"),
}
loop {
let next = poll_fn(|cx| {
if !daemon.alive() {
return Poll::Ready(None);
}
waker.set(Some(cx.waker().clone()));
Pin::new(&mut messages).poll_next(cx)
})
.await;
if !daemon.alive() {
daemon.clear();
return;
}
let message = match next {
Some(Ok(message)) => message,
Some(Err(error)) => {
tracing::warn!(%error, name, "a bad message on the bus");
continue;
}
None => break,
};
let header = message.header();
if header.message_type() != Kind::Signal {
continue;
}
if let Some((changed, new)) = owner_changed(&message) {
if changed == name {
generation.set(daemon.clear());
current = (!new.is_empty()).then_some(new);
if current.is_some() {
tracing::info!(name, "started; reading it");
read(&shared, &daemon, &connection, state);
} else {
tracing::info!(name, "left the bus");
}
}
continue;
}
let from_daemon = current.is_some()
&& header.sender().map(|sender| sender.as_str()) == current.as_deref()
&& header.destination().is_none();
if !from_daemon {
continue;
}
if reading.get() != Reading::Idle || daemon.reads_again(&message) {
read(&shared, &daemon, &connection, state);
} else {
daemon.signal(&connection, &message);
}
}
tracing::warn!(name, "the bus closed");
daemon.clear();
}