bscreensaver/dbus-service/src/main.rs

use async_std::{fs::File, prelude::*, sync::{Arc, Mutex}, task};
use bscreensaver_util::init_logging;
use futures::{future::FutureExt, pin_mut, select};
use log::{debug, error, info, trace, warn};
use std::{io, process::exit, time::{Duration, Instant}};
use zbus::{dbus_interface, fdo::{self, DBusProxy, RequestNameFlags}, names::{BusName, UniqueName, WellKnownName}, ConnectionBuilder, MessageHeader};

use bscreensaver_command::{bscreensaver_command, BCommand};

const OUR_DBUS_NAME: &str = "org.freedesktop.ScreenSaver";
const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(45);

struct Inhibitor {
    cookie: u32,
    app_name: String,
    peer: Option<UniqueName<'static>>,
}

struct State {
    inhibitors: Vec<Inhibitor>,
}
    
struct ScreenSaver {
    state: Arc<Mutex<State>>,
}

#[dbus_interface(name = "org.freedesktop.ScreenSaver")]
impl ScreenSaver {
    async fn inhibit(
        &mut self,
        #[zbus(header)]
        hdr: MessageHeader<'_>,
        app_name: &str,
        reason: &str
    ) -> fdo::Result<u32> {
        debug!("Handling inhibit for app {}: {}", app_name, reason);
        if app_name.trim().is_empty() {
            return Err(fdo::Error::InvalidArgs("Application name is blank".to_string()));
        } else if reason.trim().is_empty() {
            return Err(fdo::Error::InvalidArgs("Reason is blank".to_string()));
        }

        // Firefox tries to inhibit when only audio is playing, so ignore that
        if reason.contains("audio") && !reason.contains("video") {
            info!("Ignoring audio-only inhibit from app {}", app_name);
            return Ok(0);
        }

        let peer = hdr.sender()?;
        let cookie = rand_u32().await
            .map_err(|err| fdo::Error::IOError(err.to_string()))?;
        self.state.lock().await.inhibitors.push(Inhibitor {
            cookie,
            app_name: app_name.to_string(),
            peer: peer.map(|s| s.to_owned()),
        });

        Ok(cookie)
    }

    async fn un_inhibit(&mut self, cookie: u32) -> fdo::Result<()> {
        let mut state = self.state.lock().await;

        let before = state.inhibitors.len();
        state.inhibitors.retain(|inhibitor| {
            if inhibitor.cookie == cookie {
                info!("Uninhibit received from {} for cookie {}", inhibitor.app_name, cookie);
                false
            } else {
                true
            }
        });
        if before == state.inhibitors.len() {
            info!("No inhibitor found with cookie {}", cookie);
        }

        Ok(())
    }
}

#[async_std::main]
async fn main() {
    init_logging("BSCREENSAVER_DBUS_SERVICE_LOG");

    let state = Arc::new(Mutex::new(State {
        inhibitors: Vec::new(),
    }));

    let xcb_handle = task::spawn(xcb_task()).fuse();
    let dbus_handle = task::spawn(dbus_task(Arc::clone(&state))).fuse();
    let heartbeat_handle = task::spawn(heartbeat_task(Arc::clone(&state))).fuse();

    pin_mut!(xcb_handle, dbus_handle, heartbeat_handle);

    let res = loop {
        select! {
            _ = xcb_handle => {
                info!("Lost connection to X server; quitting");
                break Ok(());
            },
            res = dbus_handle => {
                match res {
                    Err(err) => error!("Lost connection to the system bus: {}", err),
                    Ok(_) => error!("DBus task exited normally; this should not happen!"),
                }
                break Err(());
            },
            res = heartbeat_handle => {
                match res {
                    Err(err) => error!("Heartbeat task terminated with error: {}", err),
                    Ok(_) => error!("Heartbeat task exited normally; this should not happen!"),
                }
                break Err(());
            }
        };
    };

    if let Err(_) = res {
        exit(1);
    }
}

async fn xcb_task() -> anyhow::Result<()> {
    let (xcb_conn, _) = task::block_on(async { xcb::Connection::connect(None) })?;
    let mut xcb_conn = async_xcb::AsyncConnection::new(xcb_conn)?;

    // We need to drain the XCB connection periodically.  Even though we have not
    // asked for any events, we'll still get stuff like MappingNotify if the keyboard
    // settings change.
    loop {
        let mut buf = [0u8; 512];
        xcb_conn.read(&mut buf).await?;
    }
}

async fn dbus_task(state: Arc<Mutex<State>>) -> anyhow::Result<()> {
    let org_fdo_screensaver = ScreenSaver { state: Arc::clone(&state) };
    let screensaver = ScreenSaver { state: Arc::clone(&state) };

    let dbus_conn = ConnectionBuilder::session()?
        .serve_at("/org/freedesktop/ScreenSaver", org_fdo_screensaver)?
        .serve_at("/ScreenSaver", screensaver)?
        .build()
        .await?;

    let our_unique_name = dbus_conn.unique_name().unwrap();

    let dbus_proxy = DBusProxy::new(&dbus_conn).await?;
    dbus_proxy.request_name(
        WellKnownName::from_static_str(OUR_DBUS_NAME)?,
        RequestNameFlags::AllowReplacement | RequestNameFlags::ReplaceExisting | RequestNameFlags::DoNotQueue
    ).await?;
    let mut name_owner_changed_stream = dbus_proxy.receive_name_owner_changed().await?;

    loop {
        if let Some(name_owner_changed) = name_owner_changed_stream.next().await {
            let args = name_owner_changed.args()?;
            match args.name() {
                BusName::WellKnown(name) if name == OUR_DBUS_NAME => {
                    if args.new_owner().is_none() || args.new_owner().as_ref().filter(|no| no != &our_unique_name).is_some() {
                        info!("Lost bus name {}; quitting", OUR_DBUS_NAME);
                        exit(0);
                    }
                },
                BusName::Unique(name) => {
                    if args.new_owner().is_none() {
                        state.lock().await.inhibitors.retain(|inhibitor| {
                            if inhibitor.peer.as_ref().filter(|n| n == &name).is_some() {
                                info!("Canceling inhibit from {}, as the client has disappeared", inhibitor.app_name);
                                false
                            } else {
                                true
                            }
                        });
                    }
                },
                _ => (),
            }
        }
    }
}

async fn heartbeat_task(state_mtx: Arc<Mutex<State>>) -> anyhow::Result<()> {
    let mut last_heartbeat: Option<Instant> = None;

    loop {
        let state = state_mtx.lock().await;
        let next_heartbeat =
            if state.inhibitors.is_empty() {
                HEARTBEAT_INTERVAL
            } else {
                if let Some(lh) = last_heartbeat {
                    let since_last = Instant::now().duration_since(lh);
                    if since_last < HEARTBEAT_INTERVAL {
                        HEARTBEAT_INTERVAL - since_last
                    } else {
                        Duration::ZERO
                    }
                } else {
                    Duration::ZERO
                }
            };
        drop(state);

        task::sleep(next_heartbeat).await;
        debug!("Heartbeat timeout expired");

        let state = state_mtx.lock().await;
        if !state.inhibitors.is_empty() && (last_heartbeat.is_none() || last_heartbeat.as_ref().filter(|lh| lh.elapsed() < HEARTBEAT_INTERVAL).is_none()) {
            trace!("About to deactivate; active inhibitors:");
            for inhibitor in &state.inhibitors {
                trace!("    {}: {}", inhibitor.cookie, inhibitor.app_name);
            }
            drop(state);
            task::block_on(async {
                if let Err(err) = bscreensaver_command(BCommand::Deactivate) {
                    warn!("Failed to deactivate screen lock: {}", err);
                } else {
                    debug!("Successfully issued deactivate heartbeat");
                    last_heartbeat = Some(Instant::now());
                }
            });
        }
    }
}

async fn rand_u32() -> io::Result<u32> {
    let mut f = File::open("/dev/urandom").await?;
    let mut buf = [0u8; 4];
    f.read_exact(&mut buf).await?;
    Ok(((buf[0] as u32) << 24) | ((buf[1] as u32) << 16) | ((buf[2] as u32) << 8) | (buf[3] as u32))
}