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::{interface, fdo::{self, DBusProxy, RequestNameFlags}, names::{BusName, UniqueName, WellKnownName}, ConnectionBuilder, MessageHeader};

use bscreensaver_command::{bscreensaver_command, BCommand};
use bscreensaver_util::opt_contains;

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>>,
}

#[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(())
    }
}

#[derive(Clone, Copy, PartialEq)]
enum DpmsState {
    Unknown,
    WasDisabled,
    WeDisabled,
    WeEnabled,
}

struct DpmsHandling {
    conn: Option<xcb::Connection>,
    state: DpmsState,
}

impl DpmsHandling {
    fn init() -> Self {
        xcb::Connection::connect_with_extensions(None, &[xcb::Extension::Dpms],&[])
            .map(|(conn, _)| Some(conn))
            .unwrap_or_else(|err| match err {
                xcb::ConnError::ClosedExtNotSupported => None,
                _ => {
                    warn!("Failed to connect to X display; we will not handle DPMS on inhibit: {}", err);
                    None
                },
            })
        .and_then(|conn| {
            let cookie = conn.send_request(&xcb::dpms::Capable {});
            conn.wait_for_reply(cookie)
                .map(|capable| {
                    if capable.capable() {
                        Some(Self {
                            conn: Some(conn),
                            state: DpmsState::Unknown,
                        })
                    } else {
                        None
                    }
                })
            .unwrap_or_else(|err| {
                warn!("Failed to check if X server is DPMS-capable: {}", err);
                None
            })
        })
        .unwrap_or_else(|| Self {
            conn: None,
            state: DpmsState::Unknown,
        })
    }

    fn maybe_disable_dpms(&mut self) {
        if let Some(conn) = &self.conn {
            let cookie = conn.send_request(&xcb::dpms::Info {});
            match conn.wait_for_reply(cookie) {
                Err(err) => warn!("Failed to query DPMS state: {}", err),
                Ok(info) if info.state() => match conn.send_and_check_request(&xcb::dpms::Disable {}) {
                    Err(err) => warn!("Failed to disable DPMS: {}", err),
                    Ok(_) => {
                        debug!("Successfully disabled DPMS");
                        self.state = DpmsState::WeDisabled;
                    },
                }
                Ok(_) => {
                    debug!("DPMS was already disabled");
                    if self.state != DpmsState::WeDisabled {
                        self.state = DpmsState::WasDisabled;
                    }
                },
            }
        }
    }

    fn maybe_enable_dpms(&mut self) {
        if let Some(conn) = &self.conn {
            if self.state == DpmsState::WeDisabled {
                match conn.send_and_check_request(&xcb::dpms::Enable {}) {
                    Err(err) => warn!("Failed to enable DPMS: {}", err),
                    Ok(_) => {
                        debug!("Successfully enabled DPMS");
                        self.state = DpmsState::WeEnabled;
                    },
                }
            } else {
                debug!("We didn't disable DPMS, so we're not going to re-enable it");
            }
        }
    }
}

#[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() || opt_contains(&args.new_owner(), &our_unique_name) {
                        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;
    let mut dpms_handling = task::block_on(async { DpmsHandling::init() });

    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() {
            task::block_on(async { dpms_handling.maybe_disable_dpms() });

            if 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, Some(Duration::from_secs(4))) {
                        warn!("Failed to deactivate screen lock: {}", err);
                    } else {
                        debug!("Successfully issued deactivate heartbeat");
                        last_heartbeat = Some(Instant::now());
                    }
                });
            }
        } else {
            task::block_on(async { dpms_handling.maybe_enable_dpms() });
        }
    }
}

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))
}