bscreensaver/locker/src/main.rs

use clap::Parser;
use log::{debug, error, info, trace, warn};
use nix::{
    poll::{poll, PollFd, PollFlags, PollTimeout}, sys::{
        signal::{sigprocmask, SigSet, SigmaskHow, Signal},
        signalfd::{SfdFlags, SignalFd},
    }, unistd::{execv, fork, setsid, ForkResult}
};
use std::{
    env, ffi::CString, fs::read_link, io, os::{fd::AsFd, unix::io::AsRawFd}, path::PathBuf, process::exit, rc::Rc, sync::Mutex, time::{Duration, Instant}
};
use xcb::{randr, x, xfixes, xinput};

use bscreensaver::{screensaver::{BlankerState, CommandHandlers, Screensaver}, subservice::Subservices, xkb_ext};
use bscreensaver_util::{*, settings::Configuration};

const BLANKED_ARG: &str = "blanked";
const LOCKED_ARG: &str = "locked";

macro_rules! maybe_add_fd {
    ($pfds:expr, $fd:expr) => {
        if let Some(fd) = $fd {
            if fd >= 0 {
                $pfds.push(PollFd::new(borrow_raw_fd(fd), PollFlags::POLLIN));
                Some(fd)
            } else {
                None
            }
        } else {
            None
        }
    };
}

#[derive(Parser, Debug)]
#[command(name = "bscreensaver", version, author, about = "Blanks and locks the screen after a period of time", long_about = None)]
struct Args {
    /// Starts up in already blanked
    #[arg(long = BLANKED_ARG)]
    blanked: bool,

    /// Stats up already blanked and locked
    #[arg(long = LOCKED_ARG)]
    locked: bool,
}

fn main() -> anyhow::Result<()> {
    init_logging("BSCREENSAVER_LOG");

    let config = Configuration::load()?;
    let args = Args::parse();

    let mut signal_fd = init_signals()?;

    let (conn, screen_num) = xcb::Connection::connect_with_extensions(
        None,
        &[xcb::Extension::RandR, xcb::Extension::XFixes, xcb::Extension::Input, xcb::Extension::Xkb],
        &[]
    )?;
    let setup = conn.get_setup();
    let screen = setup.roots().nth(screen_num as usize).unwrap();

    init_xfixes(&conn)?;
    init_xinput(&conn)?;
    init_randr(&conn)?;
    xkb_ext::x11_setup(&conn, xkb::x11::MIN_MAJOR_XKB_VERSION, xkb::x11::MIN_MINOR_XKB_VERSION)
        .map_err(|_| anyhow::anyhow!("Failed to initialize XKB extension"))?;

    let helper_dir = PathBuf::from(env!("HELPER_DIR"));

    let subservices = Rc::new(Mutex::new(Subservices::start_all(&helper_dir)?));
    let command_handlers = {
        let subservices1 = Rc::clone(&subservices);
        let subservices2 = Rc::clone(&subservices);
        CommandHandlers {
            restart_handler: &move |screensaver, trigger| {
                restart_daemon(screensaver, &mut subservices1.lock().unwrap(), trigger)
            },
            exit_handler: &move |screensaver, trigger| {
                exit_daemon(screensaver, &mut subservices2.lock().unwrap(), trigger)
            },
        }
    };
    let mut screensaver = Screensaver::new(&config, &helper_dir, &command_handlers, &conn, screen)?;

    if args.locked {
        match screensaver.lock_screen() {
            Err(err) => error!("POSSIBLY FAILED TO LOCK SCREEN ON STARTUP: {}", err),
            Ok(_) => debug!("Got --{} arg; screen locked on startup", LOCKED_ARG),
        }
    } else if args.blanked {
        match screensaver.blank_screen() {
            Err(err) => warn!("Possibly failed to blank screen on startup: {}", err),
            Ok(_) => debug!("Got --{} arg; screen locked on startup", BLANKED_ARG),
        }
    }

    let _ = conn.send_and_check_request(&x::SetScreenSaver {
        timeout: 0,
        interval: 0,
        prefer_blanking: x::Blanking::NotPreferred,
        allow_exposures: x::Exposures::NotAllowed,
    });


    loop {
        if let Err(err) = screensaver.handle_xcb_events() {
            if conn.has_error().is_err() {
                error!("Lost connection to X server; attempting to restart");
                (command_handlers.restart_handler)(&mut screensaver, None)?;
            }
            warn!("Error handling event: {}", err);
        }

        let conn_fd = conn.as_raw_fd();
        if conn_fd < 0 {
            error!("Lost connection to X server; attempting to restart");
            (command_handlers.restart_handler)(&mut screensaver, None)?;
        }
        let conn_fd = borrow_raw_fd(conn_fd);
            

        let mut pfds = Vec::new();
        pfds.push(PollFd::new(borrow_raw_fd(signal_fd.as_raw_fd()), PollFlags::POLLIN));
        pfds.push(PollFd::new(conn_fd.as_fd(), PollFlags::POLLIN));
        let dbus_service_fd = maybe_add_fd!(&mut pfds, subservices.lock().unwrap().dbus_service().map(|ds| ds.pidfd().as_raw_fd()));
        let systemd_service_fd = maybe_add_fd!(&mut pfds, subservices.lock().unwrap().systemd_service().map(|ds| ds.pidfd().as_raw_fd()));
        let dialog_fd = maybe_add_fd!(&mut pfds, screensaver.unlock_dialog_pidfd().map(|udpfd| udpfd.as_raw_fd()));

        let since_last_activity = Instant::now().duration_since(screensaver.last_user_activity());
        let poll_timeout = match screensaver.blanker_state() {
            BlankerState::Idle if since_last_activity > config.lock_timeout - config.blank_before_locking => Some(Duration::ZERO),
            BlankerState::Idle => Some(config.lock_timeout - config.blank_before_locking - since_last_activity),
            BlankerState::Blanked if since_last_activity > config.lock_timeout => Some(Duration::ZERO),
            BlankerState::Blanked => Some(config.lock_timeout - since_last_activity),
            BlankerState::Locked => None,
        };
        let poll_timeout = poll_timeout
            .map(|pt| PollTimeout::try_from(pt).unwrap_or(PollTimeout::MAX))
            .unwrap_or(PollTimeout::NONE);

        trace!("about to poll (timeout={})", poll_timeout.as_millis().map(|pt| pt.to_string()).unwrap_or("(?)".into()));
        let nready = poll(pfds.as_mut_slice(), poll_timeout)?;  // FIXME: maybe shouldn't quit here on errors if screen is locked
        trace!("polled; {} FD ready", nready);
        if nready > 0 {
            for pfd in pfds {
                if pfd.revents().filter(|pf| pf.contains(PollFlags::POLLIN)).is_some() {
                    let result = match pfd.as_fd().as_raw_fd() {
                        fd if fd == signal_fd.as_raw_fd() => handle_signals(&mut screensaver, &mut subservices.lock().unwrap(), &mut signal_fd),
                        fd if fd == conn_fd.as_raw_fd() => screensaver.handle_xcb_events(),
                        fd if opt_contains(&dbus_service_fd, &fd) => subservices.lock().unwrap().handle_quit(),
                        fd if opt_contains(&systemd_service_fd, &fd) => subservices.lock().unwrap().handle_quit(),
                        fd if opt_contains(&dialog_fd, &fd) => screensaver.handle_unlock_dialog_quit(),
                        _ => Ok(()),
                    };

                    if let Err(err) = result {
                        if conn.has_error().is_err() {
                            error!("Lost connection to X server; atempting to restart");
                            (command_handlers.restart_handler)(&mut screensaver, None)?;
                        }
                        warn!("Error handling event: {}", err);
                    }
                }
            }
        }

        let since_last_activity = Instant::now().duration_since(screensaver.last_user_activity());

        if screensaver.blanker_state() < BlankerState::Blanked && since_last_activity > config.lock_timeout - config.blank_before_locking {
            if let Err(err) = screensaver.blank_screen() {
                error!("POSSIBLY FAILED TO BLANK SCREEN: {}", err);
            }
        }

        if screensaver.blanker_state() < BlankerState::Locked && since_last_activity > config.lock_timeout {
            if let Err(err) = screensaver.lock_screen() {
                error!("POSSIBLY FAILED TO LOCK SCREEN: {}", err);
            }
        }
    }
}

fn init_signals() -> anyhow::Result<SignalFd> {
    let sigs = {
        let mut s = SigSet::empty();
        s.add(Signal::SIGHUP);
        s.add(Signal::SIGINT);
        s.add(Signal::SIGQUIT);
        s.add(Signal::SIGTERM);
        s
    };
    sigprocmask(SigmaskHow::SIG_BLOCK, Some(&sigs), None)?;
    let flags = SfdFlags::SFD_NONBLOCK | SfdFlags::SFD_CLOEXEC;
    let fd = SignalFd::with_flags(&sigs, flags)?;
    Ok(fd)
}

fn init_xfixes(conn: &xcb::Connection) -> xcb::Result<()> {
    let cookie = conn.send_request(&xfixes::QueryVersion {
        client_major_version: 4,
        client_minor_version: 0,
    });
    let reply = conn.wait_for_reply(cookie)?;
    if reply.major_version() < 4 {
        warn!("XFIXES version 4 or better is not supported by the X server (max supported: {}.{})", reply.major_version(), reply.minor_version());
    }
    Ok(())
}

fn init_xinput(conn: &xcb::Connection) -> xcb::Result<()> {
    let cookie = conn.send_request(&xcb::xinput::XiQueryVersion {
        major_version: 2,
        minor_version: 2,
    });
    let reply = conn.wait_for_reply(cookie)?;
    if reply.major_version() < 2 {
        error!("Version 2 or greater of the Xinput extension is required (got {}.{})", reply.major_version(), reply.minor_version());
        exit(1);
    }

    let mut cookies = Vec::new();
    for screen in conn.get_setup().roots() {
        cookies.push(conn.send_request_checked(&xinput::XiSelectEvents {
            window: screen.root(),
            masks: &[
                xinput::EventMaskBuf::new(
                    xinput::Device::AllMaster,
                    &[
                        xinput::XiEventMask::RAW_KEY_PRESS |
                        xinput::XiEventMask::RAW_KEY_RELEASE |
                        xinput::XiEventMask::RAW_BUTTON_PRESS |
                        xinput::XiEventMask::RAW_BUTTON_RELEASE |
                        xinput::XiEventMask::RAW_TOUCH_BEGIN |
                        xinput::XiEventMask::RAW_TOUCH_UPDATE |
                        xinput::XiEventMask::RAW_TOUCH_END |
                        xinput::XiEventMask::RAW_MOTION
                    ]
                )
            ],
        }));
    }
    for cookie in cookies {
        conn.check_request(cookie)?;
    }

    Ok(())
}

fn init_randr(conn: &xcb::Connection) -> anyhow::Result<()> {
    let cookie = conn.send_request(&randr::QueryVersion {
        major_version: 1,
        minor_version: 2,
    });
    let reply = conn.wait_for_reply(cookie)?;
    if reply.major_version() < 1 || (reply.major_version() == 1 && reply.minor_version() < 2) {
        Err(anyhow::anyhow!("XRandR extension version 1.2 or greater is required, but the X server only supports {}.{}", reply.major_version(), reply.minor_version()))?;
    }

    for screen in conn.get_setup().roots() {
        conn.send_and_check_request(&randr::SelectInput {
            window: screen.root(),
            enable: randr::NotifyMask::SCREEN_CHANGE | randr::NotifyMask::CRTC_CHANGE | randr::NotifyMask::OUTPUT_CHANGE,
        })?;
    }
    Ok(())
}

fn handle_signals(screensaver: &mut Screensaver, subservices: &mut Subservices, signal_fd: &mut SignalFd) -> anyhow::Result<()> {
    match signal_fd.read_signal()? {
        None => (),
        Some(info) if info.ssi_signo == Signal::SIGHUP as u32 => restart_daemon(screensaver, subservices, None)?,
        Some(info) if info.ssi_signo == Signal::SIGINT as u32 => exit_daemon(screensaver, subservices, None)?,
        Some(info) if info.ssi_signo == Signal::SIGQUIT as u32 => exit_daemon(screensaver, subservices, None)?,
        Some(info) if info.ssi_signo == Signal::SIGTERM as u32 => exit_daemon(screensaver, subservices, None)?,
        Some(info) => trace!("Unexpected signal {}", info.ssi_signo),
    }
    Ok(())
}

fn restart_daemon(screensaver: &mut Screensaver, subservices: &mut Subservices, trigger: Option<&x::ClientMessageEvent>) -> anyhow::Result<()> {
    info!("Restarting");

    let exe = read_link("/proc/self/exe")?;
    let exe = CString::new(exe.to_str().ok_or(io::Error::new(io::ErrorKind::InvalidData, "Path cannot be converted to str".to_string()))?)?;

    let argv = match screensaver.blanker_state() {
        BlankerState::Idle => vec![],
        BlankerState::Blanked => vec![CString::new("--".to_string() + BLANKED_ARG)?],
        BlankerState::Locked => vec![CString::new("--".to_string() + LOCKED_ARG)?],
    };

    kill_child_processes(screensaver, subservices);

    match unsafe { fork() } {
        Err(err) => {
            error!("Failed to fork: {}", err);
            if let Some(ev) = trigger {
                bscreensaver::send_command_reply(screensaver.xcb_connection(), ev, false);
            }
            Err(err)?;
        },
        Ok(ForkResult::Parent { .. }) => {
            if let Some(ev) = trigger {
                bscreensaver::send_command_reply(screensaver.xcb_connection(), ev, true);
            }
            exit(0);
        },
        Ok(ForkResult::Child) => {
            if let Err(err) = setsid() {
                warn!("Failed to start new session: {}", err);
            }
            execv(exe.as_c_str(), &argv)?;
        },
    }

    Ok(())
}

fn exit_daemon(screensaver: &mut Screensaver, subservices: &mut Subservices, trigger: Option<&x::ClientMessageEvent>) -> anyhow::Result<()> {
    info!("Quitting");
    kill_child_processes(screensaver, subservices);
    if let Some(ev) = trigger {
        bscreensaver::send_command_reply(screensaver.xcb_connection(), ev, true);
    }
    exit(0);
}

fn kill_child_processes(screensaver: &mut Screensaver, subservices: &mut Subservices) {
    screensaver.stop_unlock_dialog();
    subservices.stop_all();
}